extract

Contents

extract#

class brainunit.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 (Array | saiunit.Quantity) – The array from which to extract elements.

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

  • fill_value (Array | ndarray | bool | number | bool | int | float | complex | 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:

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