Shape Tools

ndindex contains several helper functions for working with and manipulating array shapes.

Functions

ndindex.iter_indices(*shapes, skip_axes=(), _debug=False)

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, whereas iter_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 in skip_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 size 10 axis of (10, 2) and the size 20 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” with iter_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 use for idx, in iter_indices(...) above).

As another example, say a is shape (1, 3) and b is shape (2, 1), and we want to generate indices for every value of the broadcasted operation a + b. We can do this by using a[idx1.raw] + b[idx2.raw] for every idx1 and idx2 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=())

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 with skip_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. The skip_axes argument works the same as in iter_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 in skip_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)

Exception raised by iter_indices() and broadcast_shapes() when the input shapes are not broadcast compatible.

exception ndindex.AxisError(axis, ndim)

Exception raised by iter_indices() and broadcast_shapes() when the skip_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.