extract

Contents

extract#

class saiunit.math.extract(condition, arr, *, size=None, fill_value=None, **kwargs)#

Return the elements of an array that satisfy some condition.

Parameters:

  • condition (Array) – An array of boolean values that selects which elements to extract.

  • arr (Union[Array, saiunit.Quantity]) – The array from which to extract elements.

  • size (Optional[int]) – optional static size for output. Must be specified in order for extract to be compatible with JAX transformations like jit() or vmap().

  • fill_value (Union[Array, ndarray, number, bool, saiunit.Quantity, None]) – if size is specified, fill padded entries with this value (default: 0).

Returns:

res – The extracted elements. The shape of res is the same as that of condition.

Return type:

Union[Array, saiunit.Quantity]

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> a = jnp.array([1, 2, 3]) * u.meter
>>> u.math.extract(a.mantissa > 1, a)
Contents