ix_along_axis_to_global_ix

mdtools.numpy_helper_functions.ix_along_axis_to_global_ix(ix, axis)[source]

Construct a tuple of index arrays suitable to directly index an array a from an array of indices along an axis of a.

Parameters:

ix (array_like) – Array of indices along a given axis of an array a. The shape of ix must be equal to a.shape with the dimension along axis removed. Otherwise, the resulting index array cannot be used to index a.
axis (int) – The axis of a along which the indices ix were taken.

Returns:

ix_global (tuple of numpy.ndarrays) – A tuple of index arrays that can be used to directly index a.

See also

numpy.unravel_index(): Similar function for indices of a flattened version of a
numpy.take(): Take elements from an array along an axis
numpy.take_along_axis(): Take values from an array by matching 1d index and data slices
mdtools.numpy_helper_functions.take(): Take a slice from an array along an axis

Notes

Functions like numpy.argmin() or numpy.argmax() return an “array of indices into the array [a]. It has the same shape as a.shape with the dimension along axis removed”. This index array can generally not be used to get the minimum or maximum values of a by simply doing a[ix]. Rather, you have to do np.squeeze(numpy.take_along_axis(a, numpy.expand_dims(ix, axis=axis), axis=axis)) (in this example, you can of course also use numpy.amin() and numpy.amax()). But sometimes, the indices themselves, which would make a[ix] possible, are required. numpy.take_along_axis() and numpy.expand_dims() do not help in this case.

This is where mdtools.numpy_helper_functions.ix_along_axis_to_global_ix() is useful. It constructs a tuple of index arrays, ix_global, from ix such that a[ix_global] and np.squeeze(numpy.take_along_axis(a, numpy.expand_dims(ix, axis=axis), axis=axis)) give the same result.

Examples

If a is an 1-dimensional array, it can be directly indexed with ix. In this case, mdtools.numpy_helper_functions.ix_along_axis_to_global_ix() is overkill.

1-dimensional example:

>>> a = np.array([1, 0, 2])
>>> ax = 0

>>> ix = np.argmin(a, axis=ax)
>>> ix
1
>>> a[ix]
0

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
array(1)
>>> a[ix_global]
0

>>> np.squeeze(
...     np.take_along_axis(a, np.expand_dims(ix, axis=ax), axis=ax)
... )
array(0)

>>> np.min(a, axis=ax)
0

Only for for multidimensional arrays, mdtools.numpy_helper_functions.ix_along_axis_to_global_ix() is really useful, since in this case a[ix] yields a wrong result or an error.

2-dimensional example: First axis:

>>> b = np.array([[0, 1, 2],
...               [1, 2, 0]])
>>> ax = 0

>>> ix = np.argmin(b, axis=ax)
>>> ix
array([0, 0, 1])
>>> b[ix]
array([[0, 1, 2],
       [0, 1, 2],
       [1, 2, 0]])

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
(array([0, 0, 1]), array([0, 1, 2]))
>>> b[ix_global]
array([0, 1, 0])

>>> np.squeeze(
...     np.take_along_axis(b, np.expand_dims(ix, axis=ax), axis=ax)
... )
array([0, 1, 0])

>>> np.min(b, axis=0)
array([0, 1, 0])

Second axis:

>>> b
array([[0, 1, 2],
       [1, 2, 0]])
>>> ax = 1

>>> ix = np.argmin(b, axis=ax)
>>> ix
array([0, 2])
>>> b[ix]
Traceback (most recent call last):
...
IndexError: index 2 is out of bounds for axis 0 with size 2

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
(array([0, 1]), array([0, 2]))
>>> b[ix_global]
array([0, 0])

>>> np.squeeze(
...     np.take_along_axis(b, np.expand_dims(ix, axis=ax), axis=ax)
... )
array([0, 0])

>>> np.min(b, axis=ax)
array([0, 0])

3-dimensional example: First axis:

>>> c = np.array([[[0, 1, 2],
...                [1, 2, 0]],
...
...               [[2, 0, 1],
...                [0, 2, 1]]])
>>> ax = 0

>>> ix = np.argmin(c, axis=ax)
>>> ix
array([[0, 1, 1],
       [1, 0, 0]])
>>> c[ix]
array([[[[0, 1, 2],
         [1, 2, 0]],

        [[2, 0, 1],
         [0, 2, 1]],

        [[2, 0, 1],
         [0, 2, 1]]],


       [[[2, 0, 1],
         [0, 2, 1]],

        [[0, 1, 2],
         [1, 2, 0]],

        [[0, 1, 2],
         [1, 2, 0]]]])

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
(array([[0, 1, 1],
       [1, 0, 0]]), array([[0, 0, 0],
       [1, 1, 1]]), array([[0, 1, 2],
       [0, 1, 2]]))
>>> c[ix_global]
array([[0, 0, 1],
       [0, 2, 0]])

>>> np.squeeze(
...     np.take_along_axis(c, np.expand_dims(ix, axis=ax), axis=ax)
... )
array([[0, 0, 1],
       [0, 2, 0]])

>>> np.min(c, axis=ax)
array([[0, 0, 1],
       [0, 2, 0]])

Second axis:

>>> c
array([[[0, 1, 2],
        [1, 2, 0]],

       [[2, 0, 1],
        [0, 2, 1]]])
>>> ax = 1

>>> ix = np.argmin(c, axis=ax)
>>> ix
array([[0, 0, 1],
       [1, 0, 0]])
>>> c[ix]
array([[[[0, 1, 2],
         [1, 2, 0]],

        [[0, 1, 2],
         [1, 2, 0]],

        [[2, 0, 1],
         [0, 2, 1]]],


       [[[2, 0, 1],
         [0, 2, 1]],

        [[0, 1, 2],
         [1, 2, 0]],

        [[0, 1, 2],
         [1, 2, 0]]]])

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
(array([[0, 0, 0],
       [1, 1, 1]]), array([[0, 0, 1],
       [1, 0, 0]]), array([[0, 1, 2],
       [0, 1, 2]]))
>>> c[ix_global]
array([[0, 1, 0],
       [0, 0, 1]])

>>> np.squeeze(
...     np.take_along_axis(c, np.expand_dims(ix, axis=ax), axis=ax)
... )
array([[0, 1, 0],
       [0, 0, 1]])
>>>
>>> np.min(c, axis=ax)
array([[0, 1, 0],
       [0, 0, 1]])

Third axis:

>>> c
array([[[0, 1, 2],
        [1, 2, 0]],

       [[2, 0, 1],
        [0, 2, 1]]])
>>> ax = 2

>>> ix = np.argmin(c, axis=ax)
>>> ix
array([[0, 2],
       [1, 0]])
>>> c[ix]
Traceback (most recent call last):
...
IndexError: index 2 is out of bounds for axis 0 with size 2

>>> ix_global = mdt.nph.ix_along_axis_to_global_ix(ix, ax)
>>> ix_global
(array([[0, 0],
       [1, 1]]), array([[0, 1],
       [0, 1]]), array([[0, 2],
       [1, 0]]))
>>> c[ix_global]
array([[0, 0],
       [0, 0]])

>>> np.squeeze(
...     np.take_along_axis(c, np.expand_dims(ix, axis=ax), axis=ax)
... )
array([[0, 0],
       [0, 0]])

>>> np.min(c, axis=ax)
array([[0, 0],
       [0, 0]])

Edge cases:

>>> ix = np.array([], dtype=int)
>>> for ax in range(3):
...     mdt.nph.ix_along_axis_to_global_ix(ix, ax)
(array([], dtype=int64), array([], dtype=int64))
(array([], dtype=int64), array([], dtype=int64))
(array([], dtype=int64), array([], dtype=int64))

>>> ix = np.array(0)
>>> for ax in range(3):
...     mdt.nph.ix_along_axis_to_global_ix(ix, ax)
array(0)
array(0)
array(0)