How do I vectorize this double for loop in Numpy?

Question

I have some Python / Numpy code that is running slow and I think it is because of the use of a double for loop. Here is the code.

def heat(D,u0,q,tdim):
    xdim = np.size(u0)
    Z = np.zeros([xdim,tdim])
    Z[:,0]=u0;
    for i in range(1,tdim):
        for j in range (1,xdim-1):
            Z[j,i]=Z[j,i-1]+D*q*(Z[j-1,i-1]-2*Z[j,i-1]+Z[j+1,i-1])
    return Z

I am trying to remove the double for loop and vectorize Z. Here is my attempt.

def heat(D,u0,q,tdim):
    xdim = np.size(u0)
    Z = np.zeros([xdim,tdim])
    Z[:,0]=u0;
    Z[1:,1:-1]=Z[1:-1,:-1]+D*q*(Z[:-2,:-1]-2*Z[1:-1,:-1]+Z[2:,:-1])
    return Z

This doesn't work - I get the following error:

operands could not be broadcast together with shapes (24,73) (23,74)

So somewhere in trying to vectorize Z, I messed up. Can you please help me spot my error?

Answer 1

You cannot vectorize that diffusion calculation in the time dimension of the problem, that still requires a loop. The only obvious optimization here is to replace the Laplacian calculation with a call to the numpy.diff function (which is precompiled C), so your heat equation solver becomes:

def heat(D,u0,q,tdim): 
    xdim = np.size(u0) 
    Z = np.zeros([xdim,tdim]) 
    Z[:,0]=u0; 

    for i in range(1,tdim): 
        Z[1:-1,i]=Z[1:-1,i-1] + D*q*np.diff(Z[:,i-1], 2)

    return Z

For non-trivial spatial sizes you should see considerable speed up.

Answer 2

You won't be able to remove both for loops as calculating column i depends on column i-1 which (in your second bit of code) is just zeros except for the first column.

What you could do is:

def heat(D,u0,q,tdim):
    xdim = np.size(u0)
    Z = np.zeros([xdim,tdim])
    Z[:,0]=u0;
    for i in range(1,tdim):
        Z[1:-1,i] = Z[1:-1,i-1] + D*q*(Z[:-2,i-1] - 2*Z[1:-1,i-1] + Z[2:,i-1])
    return Z

To get back to your code: You are filling Z[1,1:-1] with (first term only) Z[1:-1,:-1]. The mismatch in shape is readily apparent here.

Disregarding the second index (as you will have to loop anyway), your vectorized solution uses a different assumption than the non-vectorized approach: in the first version you have one side u0 (Z[:,0]) and two sides 0 (Z[0,:] and Z[-1,:]) while in the vectorized solution you do try to set Z[-1,:] to something other than 0 by filling Z[1:,i]. Which situation are you trying to simulate?!