choose

Contents

choose#

class saiunit.math.choose(a, choices, mode='raise', **kwargs)#

Construct a quantity or an array from an index array and a set of arrays to choose from.

Parameters:

  • a (Array | saiunit.Quantity) – This array must be an integer array of the same shape as choices. The elements of a are used to select elements from choices.

  • choices (Sequence[saiunit.Quantity | Array]) – Choice arrays. a and all choices must be broadcastable to the same shape.

  • mode (str) – Specifies how indices outside [0, n-1] will be treated: - ‘raise’ : raise an error (default) - ‘wrap’ : wrap around - ‘clip’ : clip to the range [0, n-1]

Returns:

res – The constructed array. The shape is identical to the shape of a, and the data type is the data type of choices.

Return type:

Array | saiunit.Quantity

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> choices = [jnp.array([1, 2, 3]), jnp.array([4, 5, 6])]
>>> u.math.choose(jnp.array([0, 1, 0]), choices)
Contents