如何将多层递归向量化？

Question

我是python和numpy（以及一般编程）的新手。我正在尽可能加快我的代码速度。这个数学包括几个数组的多个轴上的几个求和。我已经达到了向量化的一个层次，但我似乎不能再深入到这个层次，所以不得不求助于for循环（我相信有三个层次的递归，M、N和I，其中一个我已经消除了，I）。以下是我在相关部分的代码（这段代码可以工作，但我想加快速度）：

def B1(n, i):
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1)

for n in N:
    B[n, :] = B1(n, I)

for m in M:
    for n in N:
        C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :])

    Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :])

A = np.linalg.solve(C[1:, 1:], (0.25)*Y[1:])

dmaxi只是一个浮点数，m，n和i是整数。阵列具有以下形状：

^{pr2}$

对于第二个轴的计算，我还是看不出来。似乎当我试图像我对B的第一个轴做同样的矢量化时（定义一个函数，然后给数组作为参数），我得到了一个广播错误，因为它似乎试图同时计算两个轴，而不是第一个，然后第二个，这就是为什么我不得不强迫它进入for循环。同样的问题发生在C和Y上，这就是为什么它们都在for循环中。以防让人困惑，我所做的基本上是：

>>> B[:, :] = B1(N, I)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "sasrec_v6.py", line 155, in B1
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1)
ValueError: operands could not be broadcast together with shapes (18) (551) 

矢量化B的第二个轴极大地提高了代码的速度，所以我假设同样的方法也适用于进一步的矢量化（顺便说一下，我希望我正确地使用了这个术语）。在

Answer 1

可以使用broadcasting从一维索引向量生成二维数组。我还没有测试这些，但它们应该可以工作：

如果将N重塑为列向量，则B1将返回一个二维数组：

B[N] = B1(N[:, None], I)

对于Y和C，我会使用np.einsum更好地控制哪些轴是多重的（这可能也可以用np.dot来实现，但我不确定如何实现。

^{pr2}$

要了解索引技巧的作用：

In [1]: a = np.arange(3)

In [2]: a
Out[2]: array([0, 1, 2])

In [3]: a[:, None]
Out[3]: 
array([[0],
       [1],
       [2]])

In [4]: b = np.arange(4,1,-1)

In [5]: b
Out[5]: array([4, 3, 2])

In [6]: a[:, None] * b
Out[6]: 
array([[0, 0, 0],
       [4, 3, 2],
       [8, 6, 4]])

它在时间上节省了两个数量级：

In [92]: %%timeit
   ....: B = np.zeros((18, 551))
   ....: C = np.zeros((18, 18))
   ....: Y = np.zeros((18))
   ....: for n in N:
   ....:     B[n, :] = B1(n, I)
   ....: for m in M:
   ....:     for n in N:
   ....:         C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :])
   ....:     Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :])
   ....: 
100 loops, best of 3: 15.8 ms per loop

In [93]: %%timeit
   ....: Bv = np.zeros((18, 551))
   ....: Cv = np.zeros((18, 18))
   ....: Yv = np.zeros((18))
   ....: Bv[N] = B1(N[:, None], I)
   ....: Cv[M[:, None], N] = np.einsum('ij,kj->ik', B[M]/np.square(qi*Iq[0, :, 2]), B[N])
   ....: Yv[M] = np.einsum('i, ki->k', U[0, :, 1]/np.square(qi*Iq[0, :, 2]), B[M])
   ....: 
1000 loops, best of 3: 1.34 ms per loop

这是我的测试：

import numpy as np

# make fake data:
np.random.seed(5)

qi = np.random.rand(551)
N = np.random.randint(0,18,18)#np.arange(18)
M = np.random.randint(0,18,18)#np.arange(18)
I = np.arange(551)
Iq = np.random.rand(1, 551, 3)
U = np.random.rand(1, 551, 3)

B = np.zeros((18, 551))
C = np.zeros((18, 18))
Y = np.zeros((18))
Bv = np.zeros((18, 551))
Cv = np.zeros((18, 18))
Yv = np.zeros((18))

dmaxi = 1.

def B1(n, i):
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1)

for n in N:
    B[n, :] = B1(n, I)

for m in M:
    for n in N:
        C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :])
    Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :])

Bv[N] = B1(N[:, None], I)
print "B correct?", np.allclose(Bv, B)

# np.einsum test case:
n, m = 2, 3
a = np.arange(n*m).reshape(n,m)*8 + 2
b = np.arange(n*m)[::-1].reshape(n,m)
c = np.empty((n,n))
for i in range(n):
    for j in range(n):
        c[i,j] = np.dot(a[i],b[j])
cv = np.einsum('ij,kj->ik', a, b)
print "einsum test successful?", np.allclose(c,cv)

Cv[M[:, None], N] = np.einsum('ij,kj->ik',
        B[M]/np.square(qi*Iq[0, :, 2]),
        B[N])
print "C correct?", np.allclose(Cv, C)

Yv[M] = np.einsum('i, ki->k',
        U[0, :, 1]/np.square(qi*Iq[0, :, 2]),
        B[M])
print "Y correct?", np.allclose(Yv, Y)

输出：D

B correct? True
einsum test successful? True
C correct? True
Y correct? True