Получить индексы элементов numpy.argmax по оси

Question

У меня есть N-мерная матрица, которая содержит значения для функции с N параметрами. Каждый параметр имеет дискретное количество значений. Мне нужно максимизировать функцию по всем параметрам, кроме одного, в результате чего получается одномерный вектор с размером, равным количеству значений не максимизированного параметра. Мне также нужно сохранить, какие значения принимаются другими параметрами.

Для этого я хотел бы применить итеративно numpy.max по разным осям, чтобы уменьшить размерность матрицы, чтобы найти то, что мне нужно. Окончательный вектор будет зависеть только от того параметра, который я пропустил.

Однако у меня возникают проблемы с поиском исходных индексов конечных элементов (которые содержат информацию о значениях, взятых другими параметрами). Я хоть об использовании numpy.argmax так же, как numpy.max но я не могу получить исходные показатели.

Пример того, что я пытаюсь это:

x = [[[1,2],[0,1]],[[3,4],[6,7]]]
args = np.argmax(x, 0)

Это возвращает

[[1 1]
 [1 1]]

Это означает, что argmax выбирает элементы (2,1,4,7) в исходной матрице. Но как получить свои показатели? Я старался unravel_index, с использованием args непосредственно в качестве индекса для матрицы x, куча функций от NumPy до индекса без успеха.

С помощью numpy.where не является решением, поскольку входная матрица может иметь одинаковые значения внутри, поэтому я не смогу отличить разные исходные значения.

Answer 1

x.argmax(0) дает индексы вдоль 1-й оси для максимальных значений. использование np.indices генерировать индексы для другой оси.

x = np.array([[[1,2],[0,1]],[[3,4],[6,7]]])
x.argmax(0)
    array([[1, 1],
           [1, 1]])
a1, a2 = np.indices((2,2))
(x.argmax(0),a1,a2)
    (array([[1, 1],
            [1, 1]]),
     array([[0, 0],
            [1, 1]]),
     array([[0, 1],
            [0, 1]]))


x[x.argmax(0),a1,a2]
    array([[3, 4],
           [6, 7]])

x[a1,x.argmax(1),a2] 
    array([[1, 2],
           [6, 7]])

x[a1,a2,x.argmax(2)] 
    array([[2, 1],
           [4, 7]])

Если x имеет другие размеры, генерировать a1, а также a2 соответственно.

Официальная документация мало говорит о том, как использовать argmax, но ранее ТАК обсуждали это. Я получил эту общую идею от использования numpy.argmax() на многомерных массивах