Python 在两个numpy阵列上沿轴应用函数-形状未对齐

我可能在这里没有看到明显的东西，但不相信

np。沿轴应用
或
np。在轴上应用
是我要找的。假设我有以下两个数组：

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)

我的问题是，有没有一种比上面使用zip和迭代更有效、更符合python风格的方法？基本上，给定两个形状（2，n，k）和（2，n）的输入数组，我希望返回的数组是形状（2，k）。谢谢。
对于通用形状的
3D
和
2D
阵列-
arr3
和
arr4
，我们可以使用一些魔法来获得矢量化解决方案，就像这样-

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
进行斗争，后者在
求和方面做得更好，即使我们在循环中使用它。我们唯一能给
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
这一次肯定会赢

关于
np.einsum
和
np.dot
之间的斗争值得一看

另外，请注意，如果我们需要使用基于循环的方法，我们应该初始化输出数组，然后将
coefs
中的输出值赋给它，而不是附加，因为后者是一个缓慢的过程。
你为什么认为
arr2
的形状是10x5？3D
数组的第一轴长度是否总是
2
像
arr3
和
2D
数组
arr4
一样？@user2357112你是对的，我是错的。@Divakar简短的回答是否定的，我想概括一下，这样得到的数组的第一轴长度就保持了arr3和arr4的第一轴长度。查看发布的解决方案，以实现相同的效果？
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