title: python/numpy: vectorize nested for loops
body: <p>I have been trying to shrug off my FORTRAN sensibilities over the last few days and embrace python's vecotrization to get rid of as many loops as possible and optimise my code.</p>
  
  <p>A number of posts on this site have been incredibly useful in achieving this, but I have hit a problem I’m not sure how to solve.</p>
  
  <p>Below is the "<code>for</code>-loop version" of the code which, I admit, uses some unnecessary array allocation, but it’s just to illustrate the issue:</p>
  
  <pre><code>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
  </code></pre>
  
  <p>So thats the snippet of code. When vecotrizing this, it seems straight forward to remove the 
  outer '<strong><code>i</code></strong>' and '<strong><code>j</code></strong>' for loops, leaving the inner '<strong><code>k</code></strong>' and '<strong><code>p</code></strong>' loops.</p>
  
  <p>So below is my attempt to do this:</p>
  
  <pre><code>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])
  </code></pre>
  
  <p>I took <strong><code>B</code></strong> array out of the <strong><code>k</code></strong> loop and just wrote:</p>
  
  <pre><code>B11 = BB( b1,b2,bs2,P1,P2,P3,S2_12,S2_13,S2_23,F2_12,F2_13,F2_23 )
  </code></pre>
  
  <p><strong>However,</strong> the thing I can't seem to get my head around is how to follow on from this and incorporate the <strong><code>p</code></strong> loop as this is nested within the <strong><code>k</code></strong> loop:</p>
  
  <pre><code>for p in range( Nmodes ):
        Bl[p] = Bl[p] + 2. * pi * LegMu[k,p] * dmu * B11
  </code></pre>
  
  <p><strong>Crucially</strong>, if you look at the first version I have to set <strong><code>Bl</code></strong> array to zero just before the <strong><code>k</code></strong> loop is called. There is stuff that happens after this, that uses <strong><code>Bl</code></strong>, but this is the place where I am stuck for now.</p>
  
  <p>Any help would be much appreciated!</p>
  
  <hr>
  
  <p>Ok, as requested I will simplify the above so it better illustrates the mechanics of the problem. You can ignore the array values I assign - its just an example:</p>
  
  <p>So beginning with the "<code>for</code>-loop version" ...</p>
  
  <pre><code>kbins  = linspace( -1, 1, 100 ) )
  mubins = linspace( -5, 5, 100 ) )
  
  nb     = 100
  
  BLB    = np.zeros( 10 )           # &lt;--------------------------------- 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 )      # &lt;---------------------------------- 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]
  </code></pre>
  
  <p>so my attempt at vectorising this gets as far as the <strong><code>p</code></strong> loop to look like this:</p>
  
  <pre><code>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]
  </code></pre>
  
  <p><strong>but</strong> how do i compute <strong><code>BL</code></strong> and <strong><code>BLB</code></strong> in the respective <strong><code>p</code></strong> and <strong><code>q</code></strong> loops? </p>
  
  <p>I hope this makes a bit more sense.</p>