lovespeed
lovespeed

Reputation: 5045

Parallel array manipulations in numpy

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?

Upvotes: 1

Views: 2376

Answers (3)

Charles Brunet
Charles Brunet

Reputation: 23120

To perform parallel processing in numpy, you should look at mpi4py. It's an MPI binding for Python. It allows distributed processing.

Upvotes: 0

Jaime
Jaime

Reputation: 67427

If you are taking good care of the edges, the standard way of going about your first vectorization would be something like this:

B = np.zeros(A.shape)
B[:-1, 1:, 1:] = A[1:, 1:, 1:] * np.sqrt(A[:-1, :-1, :-1])

You would then need to fill B[-1, :, :], B[:, 0, :] and B[:, :, 0] with appropriate values.

Extending this to other indices should be pretty straightforward.

Upvotes: 0

YXD
YXD

Reputation: 32511

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]

Upvotes: 2

Related Questions