Вычесть матрицу из ndarray вдоль определенного измерения без изменения формы

Question

Предположим, у меня есть 2 x 2 x 100 ndarray с именем ahat и матрица 2 x 2 с именем A.Какой самый питонный способ вычесть матрицу 2 x 2 из размера 100 без зацикливания или изменения формы?

for k in range(ahat.shape[2]):
    ahat[:,:,k] = ahat[:,:, k] - A

Я пытался использовать np.apply_over_axes, но не смог заставить его работать.В более общем случае, если у меня есть два ndarrays совместимых форм, как в примере выше, что является предпочтительным способом работы вдоль определенного измерения.Для примера я мог бы хотеть умножить каждую матрицу 2x2 в ahat на A, или я мог бы хотеть применить np.linalg.inv вдоль измерения размера 100.

Answer 1

Помимо принятого ответа, на случай, если кому-то интересно, какой из двух вариантов для инвертирования быстрее (из комментариев):

import timeit

setup='''import numpy as np
x = np.random.randint(0,100,(3,3,10000))'''


stm1 = 'xinv1 = np.moveaxis(np.asarray(list(map(np.linalg.inv, np.moveaxis(x,2,0)))),0,2)'
stm2 = 'np.apply_along_axis(lambda arr: np.linalg.inv(arr.reshape(3,3)), axis=0, arr=x.reshape(-1, 10000))'

times = timeit.repeat(setup = setup, stmt = stm1, number=100)
print('Using map: ', times)
times = timeit.repeat(setup = setup, stmt = stm2, number=100)
print('Using apply_along_axis: ', times)

дал мне:

Using map:  [5.976081462009461, 6.024182428998756, 6.218410155008314]
Using apply_along_axis:  [8.279263457996421, 7.926949607004644, 7.928437952010427]

Answer 2

Повышение A до массива формы (2,2,1), затем NumPy широковещание сделает все остальное:

ahat -= A[..., None]

A[..., None] эквивалентно A[..., np.newaxis].Добавляет новую ось длины 1 в массив.Поскольку ahat имеет форму (2,2,100), а A[..., None] имеет форму (2,2,1), трансляция NumPy приведет оба массива к совместимой форме (2,2100) (но с эффективным использованием памяти, без фактического копирования значений из A в больший массив).

В более общем смысле, обратите внимание, что трансляция NumPy автоматически добавит новые оси к левой стороне формы любого массива NumPy.Так, например, с любой базовой арифметической операцией NumPy, такой как сложение или умножение, A будет автоматически транслироваться в форму, подобную (1,2,2), если другой массив, участвующий в арифметической операции, будет 3-мерным (иличетный (1,1,2,2), если другой массив был 4-мерным).Выше нам потребовалось A[..., None], чтобы явно добавить новую ось, поскольку мы хотели, чтобы новая ось находилась с правой стороны фигуры.

---

Для умножения матриц вы обычно используете np.dot, или np.einsum, или np.tensordot.Некоторые из этих функций, такие как np.einsum и np.tensordot, позволяют вам указать, какие оси вы хотите использовать в умножении матриц (поэтому вам не нужно явно добавлять новые оси).Например, чтобы умножить матрицу ahat и A, вы можете использовать

np.tensordot(ahat, A, axes=[[1], [0]])

или

np.einsum('ijk,jl->ikl', ahat, A)

Это умножает значения вдоль оси 1 ahat на значениявдоль оси 0 A, а затем суммы для продуктов.