How do I outline this double loop for Numpy?

Question

I have Python / Numpy code that runs slowly, and I think this is due to using 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 double for a 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 does not work. I get the following error:

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

So, somewhere in an attempt to vectorize Z, I messed up. Could you help me identify my mistake?

Answer 1

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

 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 dimensions, you should see significant speed.

Answer 2

You cannot delete both for loops, because the calculating column i depends on the i-1 column, which (in your second bit of code) is only zeros, except for the first column.

What can you do:

 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 return to your code: You fill in Z [1,1: -1] s (only the first term) Z [1: -1 ,: - 1]. Form mismatch is easily seen here.

Excluding the second index (since you still have to loop), 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 ,:]), and in the vectorized solution you are trying to set Z [-1 ,:] to something other than 0, filling Z [1 :, i]. What situation are you trying to simulate ?!