How can I improve this nested for loops using numpy arrays

Question

I imported a code from matlab. I want to use numpy arrays to avoid the nested loops. But I am failing, especially at the last part. The snippet is from a larger code. This part takes with large files about 5-10 min. For simplicity, I reduced the size of the loops.

import numpy as np
i = np.sqrt(-1+0j)
omega = np.array([0.0, 392.6, 785.3, 1178.0])
P = np.array([[ 1.,  2.,  3.],
           [ 4.,  5.,  6.],
           [ 7.,  8.,  9.],
           [ 10., 11., 12.]])
cT = np.array([ 100,  600, 1100, 1600, 2100, 2600, 3100])
x = np.array([2.3,2.92,3.55])
c = np.zeros((4,7))
f = np.zeros((4,7))
A = np.zeros((4,7))
for j in range(4):
        for k in range(7): 
            f[j,k] = omega[j]/(2*np.pi)
            c[j,k] = cT[k]
            delta = omega[j]/cT[k]
            temp = 0
            for l in range(3):
                temp = temp + np.exp(-i*delta*x[l])*P[j,l]
            A[j,k] = abs(temp)/3 

What I did so far:

f = omega/(2*np.pi)+f.T).T
c = cT+c
delta = omega.T/(A+cT).T).T

I cannot figure out how to implement it for temp and A and I assume there are more elegant ways to get f and delta. For temp I think, I need one more dimension like temp.shape=(4,7,3) and then use np.sum() to sum the columns. Can someone help? Do you need more information?

Answer 1

Yes, you can do it with third dimension like this

f = np.zeros((4, 7))
f += omega.reshape(-1, 1) / (2 * np.pi)
c = np.zeros((4, 7))
c += cT
A = (
    np.abs(
        (
            np.exp(
                -i
                * omega.reshape(-1, 1, 1)
                / cT.reshape(1, -1, 1)
                * x.reshape(1, 1, -1)
            )
            * P.reshape(P.shape[0], 1, -1)
        ).sum(axis=2)
    )
    / 3
)

and

delta = np.divide.outer(omega, cT)

if you need it separate from A .

Answer 2

When I run your code in a ipython session, I get (for some of the variables):

In [2]: f                                                                                              
Out[2]: 
array([[  0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ],
       [ 62.48423066,  62.48423066,  62.48423066,  62.48423066,
         62.48423066,  62.48423066,  62.48423066],
       [124.98437681, 124.98437681, 124.98437681, 124.98437681,
        124.98437681, 124.98437681, 124.98437681],
       [187.48452296, 187.48452296, 187.48452296, 187.48452296,
        187.48452296, 187.48452296, 187.48452296]])
In [3]: c                                                                                              
Out[3]: 
array([[ 100.,  600., 1100., 1600., 2100., 2600., 3100.],
       [ 100.,  600., 1100., 1600., 2100., 2600., 3100.],
       [ 100.,  600., 1100., 1600., 2100., 2600., 3100.],
       [ 100.,  600., 1100., 1600., 2100., 2600., 3100.]])
In [4]: A                                                                                              
Out[4]: 
array([[ 2.        ,  2.        ,  2.        ,  2.        ,  2.        ,
         2.        ,  2.        ],
       [ 0.98938405,  4.73223819,  4.91953117,  4.96188066,  4.97785314,
         4.98554628,  4.98983047],
       [ 3.7788147 ,  6.33023676,  7.48300263,  7.75346598,  7.85641105,
         7.9061776 ,  7.93394347],
       [ 7.13085317,  6.16616817,  9.42598136, 10.2410835 , 10.55615294,
        10.70940983, 10.79518136]])

omega has 4 values, so your f just replicates those to a (4,7) array. Is that intentional?

In [5]: omega                                                                                          
Out[5]: array([   0. ,  392.6,  785.3, 1178. ])
In [6]: omega/(2*np.pi)                                                                                
Out[6]: array([  0.        ,  62.48423066, 124.98437681, 187.48452296])

c is just the 7 values of cT replicated to a (4,7) array

In [7]: cT                                                                                             
Out[7]: array([ 100,  600, 1100, 1600, 2100, 2600, 3100])

In [9]: cT[None,:].repeat(4,axis=0)                                                                    
Out[9]: 
array([[ 100,  600, 1100, 1600, 2100, 2600, 3100],
       [ 100,  600, 1100, 1600, 2100, 2600, 3100],
       [ 100,  600, 1100, 1600, 2100, 2600, 3100],
       [ 100,  600, 1100, 1600, 2100, 2600, 3100]])

But you could make a (4,7) delta from omega and cT with:

In [64]: delta = omega[:,None]/cT                                                                      
In [65]: delta                                                                                         
Out[65]: 
array([[ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
         0.        ,  0.        ],
       [ 3.926     ,  0.65433333,  0.35690909,  0.245375  ,  0.18695238,
         0.151     ,  0.12664516],
       [ 7.853     ,  1.30883333,  0.71390909,  0.4908125 ,  0.37395238,
         0.30203846,  0.25332258],
       [11.78      ,  1.96333333,  1.07090909,  0.73625   ,  0.56095238,
         0.45307692,  0.38      ]])

Then the temp calc is just a (4,7,3) broadcasted product, followed by a sum:

In [66]: t = np.sum(np.exp(-i*delta[:,:,None]*x)*P[:,None,:], axis=2)                                  
In [67]: np.allclose(A, abs(t)/3)                                                                      
Out[67]: True

Oh, and minor point

In [68]: i      # i = np.sqrt(-1+0j)                                                                                       
Out[68]: 1j
In [69]: 1j                                                                                            
Out[69]: 1j