compress

Contents

compress#

class saiunit.math.compress(condition, a, axis=None, *, size=None, fill_value=None, **kwargs)#

Return selected slices of a quantity or an array along given axis.

Parameters:
  • condition (Array) – An array of boolean values that selects which slices to return. If the shape of condition is not the same as a, it must be broadcastable to a.

  • a (Array | saiunit.Quantity) – Array from which to extract a part.

  • axis (int | None) – The axis along which to take slices. If axis is None, condition must be a 1-D array with the same length as a. If axis is an integer, condition must be broadcastable to the same shape as a along all axes except axis.

  • size (int | None) – The length of the returned axis. By default, the length of the input array along the axis is used.

  • fill_value (Array | ndarray | bool | number | bool | int | float | complex | None) – The value to use for elements in the output array that are not selected. If None, the output array has the same type as a and is filled with zeros.

Returns:

res – A new array that has the same number of dimensions as a, and the same shape as a with axis axis removed.

Return type:

Array | saiunit.Quantity

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> a = [1, 2, 3, 4] * u.meter
>>> u.math.compress(jnp.array([0, 1, 1, 0]), a)