python/numpy：对嵌套循环进行矢量化

在过去的几天里，我一直试图摆脱对FORTRAN的敏感，接受python的矢量化，以尽可能多地消除循环并优化代码

这个网站上的许多帖子在实现这一点上非常有用，但我遇到了一个我不知道如何解决的问题

下面是代码的for-loop版本，我承认它使用了一些不必要的数组分配，但这只是为了说明问题：

mu       = np.zeros( nbk )
mubins   = np.linspace( -1, 1, nbk )
mu[:-1]  = ( mubins[:-1] + mubins[1:] ) / 2.

kbins    = 10**( np.linspace( kmin, kmax, nb ) )
k1       = np.zeros( nbk )
k1[:-1]  = ( kbins[:-1] + kbins[1:] ) / 2.0

nb       = 100
for i in range(         nb - 1 ):
    for j in range(     nb - 1 ):
        Bl = np.zeros( Nmodes )                 # ( will be important later ) initialising array here 
        for k in range( nb - 1 ):

            k33[i,j,k] = np.sqrt(  k1[i] * k1[i] + k1[j] * k1[j] - 2 * k1[i] * k1[j] * mu[k] )
            P11[i,j,k] = pkspline( k1[i] )      # just using intrep1d from earlier in the code - not important
            P22[i,j,k] = pkspline( k1[j] )
            P33[i,j,k] = pkspline( k33[i,j,k] )

            f212[i,j,k],s212[i,j,k] = S2F2_SLOW(k1[i],k1[j],k33[i,j,k]) # just calling some function - not important
            f213[i,j,k],s213[i,j,k] = S2F2_SLOW(k1[i],k33[i,j,k],k1[j])
            f223[i,j,k],s223[i,j,k] = S2F2_SLOW(k1[j],k33[i,j,k],k1[i])

             # computing B11 to be used in following ‘p’ loop
            B11=b1*b1*b1*P11[i,j,k]*P22[i,j,k]*2.*f212[i,j,k] + b1**2*b2*P11[i,j,k]*P22[i,j,k] + b1**2*bs2*P11[i,j,k]*P22[i,j,k]*s212[i,j,k] + b1*b1*b1*P11[i,j,k]*P33[i,j,k]*2.*f213[i,j,k] + b1**2*b2*P11[i,j,k]*P33[i,j,k] + b1**2*bs2*P11[i,j,k]*P33[i,j,k]*s213[i,j,k] + b1*b1*b1*P22[i,j,k]*P33[i,j,k]*2.*f223[i,j,k] + b1**2*b2*P22[i,j,k]*P33[i,j,k] + b1**2*bs2*P22[i,j,k]*P33[i,j,k]*s223[i,j,k]

            # new loop ( this is where my issue is ) v-v-v-v-v-v-v-v-v-v-v-v-v-v-v
            for p in range( Nmodes ):
                Bl[p] = Bl[p] + 2. * pi * LegMu[k,p] * dmu * B11

这就是代码片段。将其矢量化时，删除 外部“i”和“j”表示循环，内部“k”和“p”表示循环

下面是我的尝试：

kbins = 10**(np.linspace(kmin,kmax,nb))
kk  = np.zeros(nbk)
kk[:-1]  = (kbins[:-1]+kbins[1:])/2.0

# so from above i now create 2 new arrays that will replace k1[i] and k1[j] in the previous version 
k1 = kk[np.newaxis].T #equivalent to k1[i]
k2 = kk               #equivalent to  k1[j]

#i and j loops now removed and left with k ( i may be able to get rid of the 'k' loop as well but i can't see how)
for k in range(nbk-1):
    k3[:-1,:-1,k]=   np.sqrt(np.square(k2[:-1]) + np.square(k1[:-1]) -2*k1[:-1]*k2[:-1]*mu[k])
    print k
    P1[:-1,:-1,k]=pkspline(k1[:-1])
    P2[:-1,:-1,k]=pkspline(k2[:-1])
    P3[:-1,:-1,k]=pkspline(k3[:-1,:-1,k])

    F2_12[:-1,:-1,k],S2_12[:-1,:-1,k]=S2F2(k1[:-1],k2[:-1],k3[:-1,:-1,k])
    F2_13[:-1,:-1,k],S2_13[:-1,:-1,k]=S2F2(k1[:-1],k3[:-1,:-1,k],k2[:-1])
    F2_23[:-1,:-1,k],S2_23[:-1,:-1,k]=S2F2(k2[:-1],k3[:-1,:-1,k],k1[:-1])

    #i've now put BB into a function. 
    B11[:-1,:-1,k] = BB(b1,b2,bs2,P1[:-1,:-1,k],P2[:-1,:-1,k],P3[:-1,:-1,k],S2_12[:-1,:-1,k],S2_13[:-1,:-1,k],S2_23[:-1,:-1,k],F2_12[:-1,:-1,k],F2_13[:-1,:-1,k],F2_23[:-1,:-1,k])

我将B数组从k循环中取出，然后写下：

B11 = BB( b1,b2,bs2,P1,P2,P3,S2_12,S2_13,S2_23,F2_12,F2_13,F2_23 )

然而，我似乎无法理解的是，如何从这一点出发，合并p循环，因为它嵌套在k循环中：

关键的是，如果你看第一个版本，我必须在调用k循环之前将Bl数组设置为零。在这之后会发生一些事情，使用Bl，但这就是我现在被困的地方

任何帮助都将不胜感激

好的，根据要求，我将简化以上内容，以便更好地说明问题的机理。您可以忽略我指定的数组值-这只是一个示例：

所以从for循环版本开始

kbins  = linspace( -1, 1, 100 ) )
mubins = linspace( -5, 5, 100 ) )

nb     = 100

BLB    = np.zeros( 10 )           # <--------------------------------- see q loop

for i in range(         nb - 1 ):
    k1 = ( kbins[i] + kbins[i+1] ) / 2.0

    for j in range(     nb - 1 ): 
        k2 = ( kbins[j] + kbins[j+1] ) / 2.0

        BL = np.zeros( 10 )      # <---------------------------------- see 'p' loop

        for k in range( nb - 1 ):
            mu = ( mubins[k] + mubins[k+1] ) / 2.
            k3 =  np.sqrt( k1 + k2 - 2 * mu )

            x = some_function(k1,k2,k3)
            y = some_function(k1,k3,k2)
            z = some_function(k2,k3,k1)

            B = x + y + z

            for p in range( 10 ):
                BL[p] = BL[p] + 2. * B

            for q in range( 10 ):
                BLB[q] = BLB[q] + BL[q]

但是我如何在各自的p和q循环中计算BL和BLB呢

我希望这更有意义。

看起来有些函数接受3个标量并返回一个标量。B也是标量。所以

        B = x + y + z

        for p in range( 10 ):
            BL[p] = BL[p] + 2. * B

        for q in range( 10 ):
            BLB[q] = BLB[q] + BL[q]

可简化为：

BL += 2*B
BLB += BL

我看不到在p和q上迭代的点。如所述，BL的所有10个值与BLB的10个值相同

当我尝试使用一个简单的some_函数（如k1+k2+k3）运行脚本时，我在k3=np.sqrt k1+k2-2*mu处得到了一个间歇性错误，可能是因为k1+k2-2*mu可能变为负值。但忽略这一点，我认为您的脚本简化为：

nb     = 11
kbins  = np.linspace( -1, 1, nb )
mubins = np.linspace( -5, 5, nb )

kx = (kbins[:-1]+kbins[1:])/2.0
kmx = (mubins[:-1]+mubins[1:])/2.0

for k1 in kx:
    BLB = 0
    for k2 in kx:
        BL = 0
        for mu in kmx:
            B = k1+k2-2*mu
            BL += 2*B
        BLB += BL
    print BLB

在这里，我处理了BL和BLB总和的嵌套，所以它们有一定的意义

如果可以将内部标量计算推广到3D数组，我们就可以将整个过程矢量化

k1 = kx[:, None, None]
k2 = kx[None, :, None]
mu = kmx[None, None,:]
B = 2*(k1 + k2 - 2*mu)
# (nb-1, nb-1, nb-1)
BL = np.sum(B, axis=2)
BLB = np.sum(BL, axis=1)
print BLB

回顾第一个脚本，p循环有点复杂

for p in range( Nmodes ):
            Bl[p] += X[k,p] * B

假设X是某种2d系数数组，则其工作原理如下：

Bl += X[k,:] * B

---

你能写一个简单的例子吗？这是相当复杂的——但是如果你写一个简单的例子，抓住你试图解决的问题的所有基本特征，它会更容易帮助你。你甚至可能会发现你可以自己解决这个问题。现在就看吧，谢谢。好的，在原始帖子的底部添加了一个精简版。我希望这一点现在更清楚一些。感谢for-loop版本，Bl[p]应该是Bl[p]吗？@russelljohnston你能澄清一下你的优先级吗？您是在寻求如何提高代码执行速度，还是理解矢量化作为一种代码设计方法？另外，请确认p&q循环的缩进级别是否正确，是否只是外部k循环内的一个串行代码块。谢谢你的回答！因此，为了澄清，我在示例中为kbins和mubin提供的值只是为了简单起见。我将在我最初的帖子底部澄清p和q实际上在做什么。

for p in range( Nmodes ):
            Bl[p] += X[k,p] * B

Bl += X[k,:] * B