Numpy взять с использованием индексов нескольких измерений

Question

Мне нужна общая функция f(array, axis, indices) для указания произвольной оси в массиве numpy. Здесь

• array - это массив numpy с произвольным числом измерений
• axis - это кортеж, который задает размеры массива
• indices - это кортеж, который определяет индексы вышеуказанной оси

Например, если у меня есть 6-мерный массив A, значение функции f(A, (0,3,4), (20, 70, 3)) будет равно

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

Я подозреваю, что можно использовать np.take для достижения этого следующим образом

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

Существует ли эта функция уже в numpy? Я мог бы использовать f_take, который я написал, но скорость - это проблема для меня, поэтому, если есть что-то чисто скомпилированное (нет python l oop), это было бы здорово

Answer 1

Это можно реализовать просто так:

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]