python - 将函数沿轴应用于两个 numpy 数组 - 形状未对齐

Question

我可能没有在这里看到明显的东西，但不相信 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 维或 2d 数组:

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 和迭代(如上所述)更高效、更Python的方法？基本上，给定两个形状为 (2, n, k) 的输入数组和 (2, n)，我希望返回的数组的形状为 (2, k)。谢谢。

Answer 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 post np.einsum 和 np.dot 之间的斗争值得一看。

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