Numpy take using multiple dimension indices

Question

I need a general function f(array, axis, indices) to specify arbitrary axis in a numpy array. Here

array is a numpy array of arbitrary number of dimensions
axis is a tuple that specifies the dimensions of the array
indices is a tuple that specifies the indices of the above axis

For example, if I have a 6 dimensional array A, the value of the function f(A, (0,3,4), (20, 70, 3)) would be

A[20, :, :, 70, 3, :]

I suspect that one can use np.take to achieve this the following way

def f_take(A, axis, indices):
    A1 = A.copy()

    # Make sure we iterate over axis in descending order
    descAxIdx = np.flip(np.argsort(axis))
    descAxis = np.array(axis)[descAxIdx]
    descIndices = np.array(indices)[descAxIdx]

    for ax, ind in zip(descAxis, descIndices):
        A1 = np.take(A1, ind, ax)
    return A1

Does this function already exist in numpy? I could use f_take I wrote, but speed is an issue for me, so if there is something purely compiled (no python loop), that would be great

jdehesa jdehesa · Accepted Answer · 2020-03-13T10:38:13

That can be implemented simply as this:

import numpy as np

def f(a, axes, indices):
    a = np.asarray(a)
    slices = tuple(indices[axes.index(i)] if i in axes else slice(None)
                   for i in range(a.ndim))
    return a[slices]

Numpy take using multiple dimension indices

1 Answers