lovespeed - 1 year ago 63

Python Question

I have a code in which I need to handle some big numpy arrays. For example I have a 3D array

`A`

`B`

`A`

`B`

`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`

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 Source

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]
```