ix_along_axis_to_global_ix¶
-
mdtools.numpy_helper_functions.
ix_along_axis_to_global_ix
(ix, axis)[source]¶ Construnct 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 – A tuple of index arrays that can be used to directly index a.
- Return type
tuple of numpy.ndarrays
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
Notes
Functions like
numpy.argmin()
ornumpy.argmax()
return an “array of indices into the array [a]. It has the same shape asa.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 doinga[ix]
. Rather, you have to donp.squeeze(numpy.take_along_axis(a, numpy.expand_dims(ix, axis=axis), axis=axis))
(in this expample, you can of course also usenumpy.amin()
andnumpy.amax()
). But sometimes, the indices themselves, which would makea[ix]
possible, are required.numpy.take_along_axis()
andnumpy.expand_dims()
do not help in this case.This is where
ix_along_axis_to_global_ix()
is useful. It constructs a tuple of index arrays, ix_global, from ix such thata[ix_global]
andnp.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,
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,
ix_along_axis_to_global_ix()
is really useful, since in this casea[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)