沿特定维度从 ndarray 中减去矩阵而不重塑答案

【问题标题】：Subtract matrix from ndarray along particular dimension without reshape沿特定维度从 ndarray 中减去矩阵而不重塑
【发布时间】：2018-10-07 00:33:26
【问题描述】：

假设我有一个名为 ahat 的 2 x 2 x 100 ndarray 和一个名为 A 的 2 x 2 矩阵。从大小为 100 的维度中减去 2 x 2 矩阵而不循环或重塑的最 Pythonic 方法是什么？

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

我尝试使用np.apply_over_axes，但无法正常工作。更一般地说，如果我有两个兼容形状的 ndarray，就像上面的例子一样，沿着特定维度操作的首选方式是什么。例如，我可能想将 ahat 中的每个 2x2 矩阵与 A 相乘，或者我可能想沿大小为 100 的维度应用 np.linalg.inv。

【问题讨论】：

标签： python numpy

【解决方案1】：

将 A 提升为形状 (2,2,1) 的数组，然后 NumPy broadcasting 将完成剩下的工作：

ahat -= A[..., None]

A[..., None] 等价于A[..., np.newaxis]。它将长度为 1 的新轴添加到数组中。由于ahat 的形状为 (2,2,100)，而A[..., None] 的形状为 (2,2,1)，NumPy 广播会将两个数组提升为兼容的形状 (2,2,100)（但以一种内存有效的方式，没有实际上将值从 A 复制到更大的数组）。

更一般地，请注意 NumPy 广播会自动将新轴添加到任何 NumPy 数组形状的左侧。因此，例如，对于任何基本的 NumPy 算术运算（例如加法或乘法），如果涉及算术运算的其他数组是 3 维的，A 将自动广播到类似 (1,2,2) 的形状，（或甚至 (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)

这会将沿 ahat 的 1 轴的值与沿 A 的 0 轴的值相乘，然后求和。

【讨论】：

这通常是否适用于问题中的其他操作，如矩阵乘法或求逆？
@ITA：您关于将np.linalg.inv 应用于形状为(2,2,100) 的数组中的每个2x2 数组的问题是一个有趣的问题。据我所知，您必须为每个 2x2 数组调用一次 np.linalg.inv ......或者，如果您将数组重塑为形状为 (200, 200) 的大块矩阵，其中 2x2s 沿着 (block)对角线，您可以通过一次调用 np.linalg.inv 来计算所有这些的倒数。
我没想过函数调用的次数。我想到的一种方法是使用np.moveaxis 使其变为100 x 2 x 2，然后使用np.linalg.inv 调用map()，更pythonic，但仍然是100 个调用。
另一种计算倒数的方法是：np.apply_along_axis(lambda arr: np.linalg.inv(arr.reshape(2,2)), axis=0, arr=ahat.reshape(-1, 100))。但这只是语法糖——在后台 np.linalg.inv 在 for 循环中被调用了 100 次。

【解决方案2】：

除了接受的答案之外，如果有人想知道这两者中哪一个更快（来自 cmets）：

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]

【讨论】：