scico.numpy.util¶

Utility functions for working with jax arrays and BlockArrays.

Functions

`array_info`(x)	Return a string providing information about an array.
`array_to_namedtuple`(array)	Convert an array representation of a namedtuple back to a namedtuple.
`broadcast_nested_shapes`(shape_a, shape_b)	Compute the result of broadcasting on array shapes.
`complex_dtype`(dtype)	Construct the corresponding complex dtype for a given real dtype.
`dtype_name`(dtype)	Return the name of a dtype.
`indexed_shape`(shape, idx)	Determine the shape of an array after indexing/slicing.
`is_array`(x)	Check if input is of type `jax.Array` or `numpy.ndarray`.
`is_arraylike`(x)	Check if input is of type `jax.typing.ArrayLike`.
`is_blockable`(shapes)	Determine whether a sequence of shapes could be a `BlockArray` shape.
`is_collapsible`(shapes)	Determine whether a sequence of shapes can be collapsed.
`is_complex_dtype`(dtype)	Determine whether a dtype is complex.
`is_nested`(x)	Check if input is a list/tuple containing at least one list/tuple.
`is_real_dtype`(dtype)	Determine whether a dtype is real.
`is_scalar_equiv`(s)	Determine whether an object is a scalar or is scalar-equivalent.
`jax_indexed_shape`(shape, idx)	Determine the shape of an array after indexing/slicing.
`namedtuple_to_array`(ntpl)	Convert a namedtuple to an array.
`no_nan_divide`(x, y)	Return x/y, with 0 instead of `NaN` where y is 0.
`normalize_axes`(axes[, shape, default, sort])	Normalize axes to a sequence and optionally ensure correctness.
`real_dtype`(dtype)	Construct the corresponding real dtype for a given complex dtype.
`shape_dtype_rep`(shape, dtype)	Construct a representation of array or blockarray shape and dtype.
`shape_to_size`(shape)	Compute array size corresponding to a specified shape.
`slice_length`(length, idx)	Determine the length of an array axis after indexing.
`transpose_list_of_ntpl`(ntlist)	Convert a list of namedtuples to namedtuple of lists.
`transpose_ntpl_of_list`(ntpl)	Convert a namedtuple of lists/arrays to a list of namedtuples.

scico.numpy.util.transpose_ntpl_of_list(ntpl)[source]¶

Convert a namedtuple of lists/arrays to a list of namedtuples.

Parameters:: ntpl (NamedTuple) – Named tuple object to be transposed.
Return type:: List[NamedTuple]
Returns:: List of namedtuple objects.

scico.numpy.util.transpose_list_of_ntpl(ntlist)[source]¶

Convert a list of namedtuples to namedtuple of lists.

Parameters:: ntpl – List of namedtuple objects to be transposed.
Return type:: NamedTuple
Returns:: Named tuple of lists.

scico.numpy.util.namedtuple_to_array(ntpl)[source]¶

Convert a namedtuple to an array.

Convert a collections.namedtuple object to a numpy.ndarray object that can be saved using numpy.savez.

Parameters:: ntpl (NamedTuple) – Named tuple object to be converted to ndarray.
Return type:: Array
Returns:: Array representation of input named tuple.

scico.numpy.util.array_to_namedtuple(array)[source]¶

Convert an array representation of a namedtuple back to a namedtuple.

Convert a numpy.ndarray object constructed by namedtuple_to_array back to the original collections.namedtuple representation.

Parameters:: by (Array representation of named tuple constructed) – namedtuple_to_array.
Return type:: NamedTuple
Returns:: Named tuple object with the same name and fields as the original named tuple object provided to namedtuple_to_array.

scico.numpy.util.normalize_axes(axes, shape=None, default=None, sort=False)[source]¶

Normalize axes to a sequence and optionally ensure correctness.

Normalize axes to a tuple or list and (optionally) ensure that entries refer to axes that exist in shape.

Parameters:

axes (Union[int, Tuple[int, ...], List[int], None]) – User specification of one or more axes: int, list, tuple, or None. Negative values count from the last to the first axis.
shape (Optional[Tuple[int, ...]]) – The shape of the array of which axes are being specified. If not None, axes is checked to make sure its entries refer to axes that exist in shape.
default (Optional[List[int]]) – Default value to return if axes is None. By default, tuple(range(len(shape))).
sort (bool) – If True, sort the returned axis indices.

Return type:

Sequence[int]

Returns:

Tuple or list of axes (never an int, never None). The output will only be a list if the input is a list or if the input is None and defaults is a list.

scico.numpy.util.slice_length(length, idx)[source]¶

Determine the length of an array axis after indexing.

Determine the length of an array axis after slicing. An exception is raised if the indexing expression is an integer that is out of bounds for the specified axis length. A value of None is returned for valid integer indexing expressions as an indication that the corresponding axis shape is an empty tuple; this value should be converted to a unit integer if the axis size is required.

Parameters:

length (int) – Length of axis being sliced.
idx (AxisIndex) – Indexing/slice to be applied to axis.

Return type:

Optional[int]

Returns:

Length of indexed/sliced axis.

Raises:

ValueError – If idx is an integer index that is out bounds for the axis length or if the type of idx is not one of Ellipsis, int, or slice.

scico.numpy.util.indexed_shape(shape, idx)[source]¶

Determine the shape of an array after indexing/slicing.

The indexed shape is determined by replicating the observed effects of NumPy/JAX array indexing/slicing syntax. It is significantly faster than jax_indexed_shape, and has a minimal memory footprint in all circumstances.

Parameters:

shape (Tuple[int, ...]) – Shape of array.
idx (ArrayIndex) – Indexing expression (singleton or tuple of Ellipsis, int, slice, or None (np.newaxis)).

Return type:

Tuple[int, ...]

Returns:

Shape of indexed/sliced array.

Raises:

ValueError – If any element of idx is not one of Ellipsis,
int –
index is out bounds for the corresponding axis length. –

scico.numpy.util.jax_indexed_shape(shape, idx)[source]¶

Determine the shape of an array after indexing/slicing.

The indexed shape is determined by constructing and indexing an array of the appropriate shape, relying on jax.jit to avoid memory allocation. It is potentially more reliable than indexed_shape because the indexing/slicing calculations are referred to JAX, but is significantly slower, and will involved potentially significant memory allocations if JIT is disabled, e.g. for debugging purposes.

Parameters:

shape (Tuple[int, ...]) – Shape of array.
idx (ArrayIndex) – Indexing expression (singleton or tuple of Ellipsis, int, slice, or None (np.newaxis)).

Return type:

Tuple[int, ...]

Returns:

Shape of indexed/sliced array.

scico.numpy.util.no_nan_divide(x, y)[source]¶

Return x/y, with 0 instead of NaN where y is 0.

Parameters:

x (Union[BlockArray, Array]) – Numerator.
y (Union[BlockArray, Array]) – Denominator.

Return type:

Union[BlockArray, Array]

Returns:

x / y with 0 wherever y == 0.

scico.numpy.util.array_info(x)[source]¶

Return a string providing information about an array.

Parameters:: x (Union[BlockArray, Array]) – A numpy or jax array or scico BlockArray.
Return type:: str
Returns:: A string containing information on the array.
Raises:: TypeError – If the array is not of a recognized type.

scico.numpy.util.shape_to_size(shape)[source]¶

Compute array size corresponding to a specified shape.

Compute array size corresponding to a specified shape, which may be nested, i.e. corresponding to a BlockArray.

Parameters:: shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – A shape tuple.
Return type:: int
Returns:: The number of elements in an array or BlockArray with shape shape.

scico.numpy.util.is_array(x)[source]¶

Check if input is of type jax.Array or numpy.ndarray.

Check if input is an array, of type jax.Array or numpy.ndarray.

Parameters:: x (Any) – Object to be tested.
Return type:: bool
Returns:: True if x is an array, False otherwise.

scico.numpy.util.is_arraylike(x)[source]¶

Check if input is of type jax.typing.ArrayLike.

isinstance(x, jax.typing.ArrayLike) does not work in Python < 3.10, see https://jax.readthedocs.io/en/latest/jax.typing.html#jax-typing-best-practices.

Parameters:: x (Any) – Object to be tested.
Return type:: bool
Returns:: True if x is an ArrayLike, False otherwise.

scico.numpy.util.is_nested(x)[source]¶

Check if input is a list/tuple containing at least one list/tuple.

Parameters:: x (Any) – Object to be tested.
Return type:: bool
Returns:: True if x is a list/tuple containing at least one list/tuple, False otherwise.

Example

>>> is_nested([1, 2, 3])
False
>>> is_nested([(1,2), (3,)])
True
>>> is_nested([[1, 2], 3])
True

scico.numpy.util.is_collapsible(shapes)[source]¶

Determine whether a sequence of shapes can be collapsed.

Return True if the a list of shapes represent arrays that can be stacked, i.e., they are all the same.

Parameters:: shapes (Sequence[Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]]) – A sequence of shapes.
Return type:: bool
Returns:: A boolean value indicating whether the shapes are all the same.

scico.numpy.util.is_blockable(shapes)[source]¶

Determine whether a sequence of shapes could be a BlockArray shape.

Return True if the sequence of shapes represent arrays that can be combined into a BlockArray, i.e., none are nested.

Parameters:: shapes (Sequence[Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]]) – A sequence of shapes.
Return type:: TypeGuard[Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]]
Returns:: A boolean value indicating whether any of the shapes are nested.

scico.numpy.util.shape_dtype_rep(shape, dtype)[source]¶

Construct a representation of array or blockarray shape and dtype.

Construct a representation of array or block array shape and dtype that is suitable for both jax arrays and scico blockarrays.

Parameters:

shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Array or blockarray shape.
dtype (DType) – Array or blockarray dtype.

Return type:

Union[ShapeDtypeStruct, BlockArray]

Returns:

A jax.ShapeDtypeStruct or a BlockArray containing objects of type jax.ShapeDtypeStruct.

scico.numpy.util.broadcast_nested_shapes(shape_a, shape_b)[source]¶

Compute the result of broadcasting on array shapes.

Compute the result of applying a broadcasting binary operator to (block) arrays with (possibly nested) shapes shape_a and shape_b. Extends numpy.broadcast_shapes to also support the nested tuple shapes of BlockArrays.

Parameters:

shape_a (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – First array shape.
shape_b (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Second array shape.

Return type:

Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]

Returns:

A (possibly nested) shape tuple.

Example

>>> broadcast_nested_shapes(((1, 1, 3), (2, 3, 1)), ((2, 3,), (2, 1, 4)))
((1, 2, 3), (2, 3, 4))

scico.numpy.util.is_real_dtype(dtype)[source]¶

Determine whether a dtype is real.

Parameters:: dtype (DType) – A numpy or scico.numpy dtype (e.g. float32, complex64).
Return type:: bool
Returns:: False if the dtype is complex, otherwise True.

scico.numpy.util.is_complex_dtype(dtype)[source]¶

Determine whether a dtype is complex.

Parameters:: dtype (DType) – A numpy or scico.numpy dtype (e.g. float32, complex64).
Return type:: bool
Returns:: True if the dtype is complex, otherwise False.

scico.numpy.util.real_dtype(dtype)[source]¶

Construct the corresponding real dtype for a given complex dtype.

Construct the corresponding real dtype for a given complex dtype, e.g. the real dtype corresponding to complex64 is float32.

Parameters:: dtype (DType) – A complex numpy or scico.numpy dtype (e.g. complex64, complex128).
Return type:: DType
Returns:: The real dtype corresponding to the input dtype

scico.numpy.util.complex_dtype(dtype)[source]¶

Construct the corresponding complex dtype for a given real dtype.

Construct the corresponding complex dtype for a given real dtype, e.g. the complex dtype corresponding to float32 is complex64.

Parameters:: dtype (DType) – A real numpy or scico.numpy dtype (e.g. float32, float64).
Return type:: DType
Returns:: The complex dtype corresponding to the input dtype.

scico.numpy.util.dtype_name(dtype)[source]¶

Return the name of a dtype.

Construct a string representation of a dtype name.

Parameters:: dtype (DType) – The dtype for which the name is required.
Return type:: str
Returns:: The name of the dtype.

scico.numpy.util.is_scalar_equiv(s)[source]¶

Determine whether an object is a scalar or is scalar-equivalent.

Determine whether an object is a scalar or a singleton array.

Parameters:: s (Any) – Object to be tested.
Return type:: bool
Returns:: True if the object is a scalar or a singleton array, otherwise False.