python/numpy: vectorize inested for loops

Question

Я пытаюсь избавиться от своих чувств FORTRAN в течение последних нескольких дней и использовать vecotrization 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 

Так вот фрагмент кода. Когда vecotrizing это, кажется, прямо вперед, чтобы удалить внешнее «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:

for p in range(Nmodes): 
     Bl[p] = Bl[p] + 2. * pi * LegMu[k,p] * dmu * B11 

Кардинально, если вы посмотрите на первую версию я должен установить Bl массив к нулю только перед циклом k называется. После этого происходит материал, который использует Bl, но это то место, где я застрял.

Любая помощь была бы высоко оценена!

ОК, в соответствии с просьбой, я упрощу это, чтобы лучше показать механику проблемы. Вы можете игнорировать значения массива, которые я присваиваю, - это всего лишь пример:

Итак, начиная с «for -loop version» ...

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] 

так что моя попытка vectorising это становится насколько p петли выглядеть следующим образом:

kbins = linspace(-1, 1, 100))     # as before 

# i now define k1 and k2 here as vectors, and not scalars as was above example 
kk  = np.zeros(nbk) 
kk[:-1] = (kbins[:-1] + kbins[1:])/2.0 

k1  = kk[np.newaxis].T      # equivalent to k1 in above 
k2  = kk          # equivalent to k2 in above 

mu[:-1] = (mubins[:-1] + mubins[1:]) /2.  # mu is now an array as well 

for k in range(nbk - 1):      # i,j-loops removed,left with k 
    k3[:-1,:-1,k] = np.sqrt(k1[:-1] + k2[:-1] - 2 * mu[k])  # k3 is now an array 
    x[ :-1,:-1,k] = some_function(k1[:-1],k2[:-1],k3[:-1,:-1,k]) # x,y,z now arrays to allow looping over elements 
    y[ :-1,:-1,k] = some_function(k1[:-1],k3[:-1,:-1,k],k2[:-1]) 
    z[ :-1,:-1,k] = some_function(k2[:-1],k3[:-1,:-1,k],k1[:-1]) 

    B11[:-1,:-1,k]= x[:-1,:-1,k] + y[:-1,:-1,k] + z[:-1,:-1,k] 

но как я вычислить BL и BLB в соответствующие p и q Петли?

Надеюсь, это будет иметь больше смысла.

Answer 1

Похож на то, что some_function принимает 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. Как написано, все 10 значений BL будут такими же, как и 10 значений BLB.

Когда я пытаюсь запустить скрипт (используя простой some_function как 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 сложений, так что они сделали какой-то смысл.

Если этот внутренний скалярный calc может быть обобщен для принятия 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 

Оглядываясь назад 1-й сценарий цикла p является немного более сложным

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

, который должен работать (при условии X своего рода 2d массив коэффициентов):

Bl += X[k,:] * B