Source code for mars.tensor.indexing.compress

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

from ..utils import validate_axis
from ..datasource import tensor as astensor
from ..core import Tensor

[docs]def compress(condition, a, axis=None, out=None): """ Return selected slices of a tensor along given axis. When working along a given axis, a slice along that axis is returned in `output` for each index where `condition` evaluates to True. When working on a 1-D array, `compress` is equivalent to `extract`. Parameters ---------- condition : 1-D tensor of bools Tensor that selects which entries to return. If len(condition) is less than the size of `a` along the given axis, then output is truncated to the length of the condition tensor. a : array_like Tensor from which to extract a part. axis : int, optional Axis along which to take slices. If None (default), work on the flattened tensor. out : Tensor, optional Output tensor. Its type is preserved and it must be of the right shape to hold the output. Returns ------- compressed_array : Tensor A copy of `a` without the slices along axis for which `condition` is false. See Also -------- take, choose, diag, diagonal, select Tensor.compress : Equivalent method in ndarray mt.extract: Equivalent method when working on 1-D arrays Examples -------- >>> import mars.tensor as mt >>> a = mt.array([[1, 2], [3, 4], [5, 6]]) >>> a.execute() array([[1, 2], [3, 4], [5, 6]]) >>> mt.compress([0, 1], a, axis=0).execute() array([[3, 4]]) >>> mt.compress([False, True, True], a, axis=0).execute() array([[3, 4], [5, 6]]) >>> mt.compress([False, True], a, axis=1).execute() array([[2], [4], [6]]) Working on the flattened tensor does not return slices along an axis but selects elements. >>> mt.compress([False, True], a).execute() array([2]) """ a = astensor(a) condition = astensor(condition, dtype=bool) if condition.ndim != 1: raise ValueError("condition must be an 1-d tensor") if axis is None: a = a.ravel() if len(condition) < a.size: a = a[: len(condition)] return a[condition] try: axis = validate_axis(a.ndim, axis) except ValueError: raise np.AxisError( f"axis {axis} is out of bounds for tensor of dimension {a.ndim}" ) try: if len(condition) < a.shape[axis]: a = a[(slice(None),) * axis + (slice(len(condition)),)] t = a[(slice(None),) * axis + (condition,)] if out is None: return t if out is not None and not isinstance(out, Tensor): raise TypeError(f"out should be Tensor object, got {type(out)} instead") if not np.can_cast(out.dtype, t.dtype, "safe"): raise TypeError( f"Cannot cast array data from dtype('{out.dtype}') to dtype('{t.dtype}') " "according to the rule 'safe'" ) # skip shape check because out shape is unknown = t.astype(out.dtype, order=out.order.value).data return out except IndexError: raise np.AxisError( f"axis {len(condition)} is out of bounds for tensor of dimension 1" )