Question

I have a code in which I need to handle some big numpy arrays. For example I have a 3D array A and I need to construct another 3d array B using the elements of A. However all the elements of B are independent of each other. Example:

for i in np.arange(Nx):
  for j in np.arange(Ny):
   for k in np.arange(Nz):
       B[i][j][k] = A[i+1][j][k]*np.sqrt(A[i][j-1][k-1])

So it will speed up immensely if I can construct the B array parallely. What is the simplest way to do this in python?

I also have similar matrix operations like normalizing each row of a 2D array. Example

for i in np.arange(Nx):
   f[i,:] = f[i,:]/np.linalg.norm(f[i,:])

This will also speed up if it runs parallely for each row. How can it be done?

Answer 1

You should look into Numpy’s roll function. I think this is equivalent to your first block of code (though you need to decide what happens at the edges – roll “wraps around”):

B = np.roll(A,1,axis=0) * np.sqrt(np.roll(np.roll(A,-1,axis=1),-1,axis=2))

Another fairly horrible one-liner for your second case is:

f /= np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]

Explanation of this line:

We are first going to calculate the norm of each row. Let’s

f = np.random.rand(5,6)

Square each element of f

f**2

Sum the squares along axis 1, which “flattens” out that axis.

np.sum(f**2, axis=1)

Take the square root of the sum of the squares.

np.sqrt(np.sum(f**2, axis=1))

We now have the norm of each row.

To divide each original row of f by this correctly we need to make use of the Numpy broadcasting rules to effectively add a dimension:

np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]

And finally we calculate our result

f /= np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]