Shape Tools¶
ndindex contains several helper functions for working with and manipulating array shapes.
Functions¶
- ndindex.iter_indices(*shapes, skip_axes=(), _debug=False)[source]¶
Iterate indices for every element of an arrays of shape
shapes
.Each shape in
shapes
should be a shape tuple, which are broadcast compatible along the non-skipped axes. Each iteration step will produce a tuple of indices, one for each shape, which would correspond to the same elements if the arrays of the given shapes were first broadcast together.This is a generalization of the NumPy
np.ndindex()
function (which otherwise has no relation).np.ndindex()
only iterates indices for a single shape, whereasiter_indices()
supports generating indices for multiple broadcast compatible shapes at once. This is equivalent to first broadcasting the arrays then generating indices for the single broadcasted shape.Additionally, this function supports the ability to skip axes of the shapes using
skip_axes
. These axes will be fully sliced in each index. The remaining axes will be indexed one element at a time with integer indices.skip_axes
should be a tuple of axes to skip or a list of tuples of axes to skip. If it is a single tuple, it applies to all shapes. Otherwise, each tuple applies to each shape respectively. It can use negative integers, e.g.,skip_axes=(-1,)
will skip the last axis. The order of the axes inskip_axes
does not matter. Mixing negative and nonnegative skip axes is supported, but the skip axes must refer to unique dimensions for each shape.The axes in
skip_axes
refer to the shapes before broadcasting (if you want to refer to the axes after broadcasting, either broadcast the shapes and arrays first, or refer to the axes using negative integers). For example,iter_indices((10, 2), (20, 1, 2), skip_axes=(0,))
will skip the size10
axis of(10, 2)
and the size20
axis of(20, 1, 2)
. The result is two sets of indices, one for each element of the non-skipped dimensions:>>> from ndindex import iter_indices >>> for idx1, idx2 in iter_indices((10, 2), (20, 1, 2), skip_axes=(0,)): ... print(idx1, idx2) Tuple(slice(None, None, None), 0) Tuple(slice(None, None, None), 0, 0) Tuple(slice(None, None, None), 1) Tuple(slice(None, None, None), 0, 1)
The skipped axes do not themselves need to be broadcast compatible, but the shapes with all the skipped axes removed should be broadcast compatible.
For example, suppose
a
is an array with shape(3, 2, 4, 4)
, which we wish to think of as a(3, 2)
stack of 4 x 4 matrices. We can generate an iterator for each matrix in the “stack” withiter_indices((3, 2, 4, 4), skip_axes=(-1, -2))
:>>> for idx in iter_indices((3, 2, 4, 4), skip_axes=(-1, -2)): ... print(idx) (Tuple(0, 0, slice(None, None, None), slice(None, None, None)),) (Tuple(0, 1, slice(None, None, None), slice(None, None, None)),) (Tuple(1, 0, slice(None, None, None), slice(None, None, None)),) (Tuple(1, 1, slice(None, None, None), slice(None, None, None)),) (Tuple(2, 0, slice(None, None, None), slice(None, None, None)),) (Tuple(2, 1, slice(None, None, None), slice(None, None, None)),)
Note
The iterates of
iter_indices
are always a tuple, even if only a single shape is provided (one could instead usefor idx, in iter_indices(...)
above).As another example, say
a
is shape(1, 3)
andb
is shape(2, 1)
, and we want to generate indices for every value of the broadcasted operationa + b
. We can do this by usinga[idx1.raw] + b[idx2.raw]
for everyidx1
andidx2
as below:>>> import numpy as np >>> a = np.arange(3).reshape((1, 3)) >>> b = np.arange(100, 111, 10).reshape((2, 1)) >>> a array([[0, 1, 2]]) >>> b array([[100], [110]]) >>> for idx1, idx2 in iter_indices((1, 3), (2, 1)): ... print(f"{idx1 = }; {idx2 = }; {(a[idx1.raw], b[idx2.raw]) = }") idx1 = Tuple(0, 0); idx2 = Tuple(0, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(0), np.int64(100)) idx1 = Tuple(0, 1); idx2 = Tuple(0, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(1), np.int64(100)) idx1 = Tuple(0, 2); idx2 = Tuple(0, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(2), np.int64(100)) idx1 = Tuple(0, 0); idx2 = Tuple(1, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(0), np.int64(110)) idx1 = Tuple(0, 1); idx2 = Tuple(1, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(1), np.int64(110)) idx1 = Tuple(0, 2); idx2 = Tuple(1, 0); (a[idx1.raw], b[idx2.raw]) = (np.int64(2), np.int64(110)) >>> a + b array([[100, 101, 102], [110, 111, 112]])
To include an index into the final broadcasted array, you can simply include the final broadcasted shape as one of the shapes (the function
broadcast_shapes()
is useful here).>>> np.broadcast_shapes((1, 3), (2, 1)) (2, 3) >>> for idx1, idx2, broadcasted_idx in iter_indices((1, 3), (2, 1), (2, 3)): ... print(broadcasted_idx) Tuple(0, 0) Tuple(0, 1) Tuple(0, 2) Tuple(1, 0) Tuple(1, 1) Tuple(1, 2)
- ndindex.broadcast_shapes(*shapes, skip_axes=())[source]¶
Broadcast the input shapes
shapes
to a single shape.This is the same as
np.broadcast_shapes()
, except is also supports skipping axes in the shape withskip_axes
.skip_axes
can be a tuple of integers which apply to all shapes, or a list of tuples of integers, one for each shape, which apply to each respective shape. Theskip_axes
argument works the same as initer_indices()
. See its docstring for more details.If the shapes are not broadcast compatible (excluding
skip_axes
),BroadcastError
is raised.>>> from ndindex import broadcast_shapes >>> broadcast_shapes((2, 3), (3,), (4, 2, 1)) (4, 2, 3) >>> broadcast_shapes((2, 3), (5,), (4, 2, 1)) Traceback (most recent call last): ... ndindex.shapetools.BroadcastError: shape mismatch: objects cannot be broadcast to a single shape. Mismatch is between arg 0 with shape (2, 3) and arg 1 with shape (5,).
Axes in
skip_axes
apply to each shape before being broadcasted. Each shape will be broadcasted together with these axes removed. The dimensions inskip_axes
do not need to be equal or broadcast compatible with one another. The final broadcasted shape be the result of broadcasting all the non-skip axes.>>> broadcast_shapes((10, 3, 2), (2, 20), skip_axes=[(0,), (1,)]) (3, 2)
Exceptions¶
These are some custom exceptions that are raised by the above functions. Note
that most of the other functions in ndindex will raise IndexError
(if the
index would be invalid), or TypeError
or ValueError
(if the input types or
values are incorrect).
- exception ndindex.BroadcastError(arg1, shape1, arg2, shape2)[source]¶
Exception raised by
iter_indices()
andbroadcast_shapes()
when the input shapes are not broadcast compatible.
- exception ndindex.AxisError(axis, ndim)[source]¶
Exception raised by
iter_indices()
andbroadcast_shapes()
when theskip_axes
argument is out of bounds.This is used instead of the NumPy exception of the same name so that
iter_indices
does not need to depend on NumPy.