在两个 numpy 数组上沿轴应用函数 - 形状未对齐答案

【问题标题】：Apply function along axis over two numpy arrays - shapes not aligned在两个 numpy 数组上沿轴应用函数 - 形状未对齐
【发布时间】：2017-05-30 01:55:28
【问题描述】：

我可能在这里看不到明显的东西，但不要相信 np.apply_along_axis 或 np.apply_over_axes 是我正在寻找的东西。假设我有以下两个数组：

arr1 = np.random.randn(10, 5)
arr2 = np.random.randn(10, )

还有以下功能：

def coefs(x, y):
    return np.dot(np.linalg.inv(np.dot(x.T, x)), np.dot(x.T, y))
    # the vector of coefficients in a multiple linear regression

在arr1 和arr2 上调用它可以正常工作：

coefs(arr1, arr2)
Out[111]: array([-0.19474836, -0.50797551,  0.82903805,  0.06332607, -0.26985597])

但是，假设我有两个 3d 数组，而不是 1 维或 2 维数组：

arr3 = np.array([arr1[:-1], arr1[1:]])
arr4 = np.array([arr2[:-1], arr2[1:]])

正如预期的那样，如果我在这里应用该功能，我会得到

coefs(arr3, arr4)
Traceback (most recent call last):

  File "<ipython-input-127-4a3e7df02cda>", line 1, in <module>
    coefs(arr3, arr4)

  File "<ipython-input-124-7532b8516784>", line 2, in coefs
    return np.dot(np.linalg.inv(np.dot(x.T, x)), np.dot(x.T, y))

ValueError: shapes (5,9,2) and (2,9,5) not aligned: 2 (dim 2) != 9 (dim 1)

...因为 NumPy 将每个数组视为应有的对象。我想要做的是将coefs() 函数应用于沿数组0 轴的2 个元素中的每一个元素，逐个元素。这是一种粗略的做法：

tgt = []
for i, j in zip(arr3, arr4):
    tgt.append(coefs(i, j))

np.array(tgt) 
Out[136]: 
array([[-0.34328006, -0.99116672,  1.42757897, -0.06687851, -0.44669182],
       [ 0.44494495, -0.58017705,  0.75825944,  0.18795889,  0.4560851 ]])

我的问题是，有没有比使用 zip 和迭代更有效和 Pythonic 的方法，如上所述？ 基本上，给定两个形状为 (2, n, k) 的输入数组和 (2, n)，我希望返回的数组具有 (2, k) 的形状。谢谢。

【问题讨论】：

为什么你认为arr2 的形状是 10x5？
3D 数组的第一个轴长度是否会像 arr3 和 2D 数组 arr4 一样始终为 2？
@user2357112 你说得对，是我的错字。
@Divakar 简短的回答是否定的，我想概括一下，以便生成的数组的第一个轴长度保持 arr3 和 arr4 的第一个轴长度。
@BradSolomon Kool。查看已发布的实现相同的解决方案？

标签： python python-3.x numpy

【解决方案1】：

对于通用形状的 3D 和 2D 数组 - arr3 和 arr4，我们可以使用一些 np.einsum 魔法来获得矢量化解决方案，就像这样 -

dot1 = np.einsum('ijk,ijl->ikl',arr3,arr3)
dot2 = np.einsum('ijk,ij->ik',arr3,arr4)
inv1 = np.linalg.inv(dot1)
tgt_out = np.einsum('ijk,ij->ik',inv1, dot2)

运行时测试

方法-

def org_app(arr3, arr4):
    tgt = []
    for i, j in zip(arr3, arr4):
        tgt.append(coefs(i, j))
    return np.array(tgt)

def einsum_app(arr3, arr4):
    dot1 = np.einsum('ijk,ijl->ikl',arr3,arr3)
    dot2 = np.einsum('ijk,ij->ik',arr3,arr4)
    inv1 = np.linalg.inv(dot1)
    return np.einsum('ijk,ij->ik',inv1, dot2)

时间和验证 -

In [215]: arr3 = np.random.rand(50,50,50)
     ...: arr4 = np.random.rand(50,50)
     ...: 

In [216]: np.allclose(org_app(arr3, arr4), einsum_app(arr3, arr4))
Out[216]: True

In [217]: %timeit org_app(arr3, arr4)
100 loops, best of 3: 4.82 ms per loop

In [218]: %timeit einsum_app(arr3, arr4)
100 loops, best of 3: 19.7 ms per loop

看起来einsum 并没有给我们带来任何好处。这是意料之中的，因为基本上einsum 正在与np.dot 对抗，这在sum-reduction 上要好得多，即使我们在循环中使用它。我们可以与np.dot 抗争的唯一情况/情况是，当我们循环足够多并且应该使einsum 具有竞争力时。我们循环的时间等于输入数组的第一个轴的长度。让我们增加它并再次测试-

In [219]: arr3 = np.random.rand(1000,10,10)
     ...: arr4 = np.random.rand(1000,10)
     ...: 

In [220]: %timeit org_app(arr3, arr4)
10 loops, best of 3: 23 ms per loop

In [221]: %timeit einsum_app(arr3, arr4)
100 loops, best of 3: 9.1 ms per loop

einsum 绝对赢了！

This related postnp.einsum 和 np.dot 之间的战斗值得一看。

另外，请注意，如果我们需要使用基于循环的方法，我们应该初始化输出数组，然后将来自coefs 的输出值分配到其中而不是追加，因为后者是一个缓慢的过程。

【讨论】：