CODEWITHSUNDEEP
pythonarraysperformancenumpyoptimization
velenos14
velenos14

Reputation: 566

How to efficiently populate a numpy 2D array?

I want to create a 2D numpy array of size (N_r * N_z).

Across columns, the elements for 1 specific column (say j) shall be created based on the value r_thresh[j]. So 1 column (say j) out of the total of N_z columns in the numpy 2D array is created as:

(np.arange(N_r) + 0.5) * r_thresh[j] # this gives an array of size (1, N_r)

Of course, the column j + 1 shall be created as:

(np.arange(N_r) + 0.5) * r_thresh[j+1] # this gives an array of size (1, N_r)

r_thresh is a numpy array of size (1, N_z), already populated with values before I want to create the 2D array.

I want to ask you how do I go further and use this ''rule'' of creating each element of the numpy 2D array and actually create the whole array, in the most efficient way possible (speed-wise).

I initially wrote all the code using 2 nested for loops and plain python lists and the code worked, but took forever to run.

More experienced programmers told me to avoid for loops and use numpy because it's the best.

I now understand how to create 1D arrays using numpy np.arange() instruction, but I lack the knowledge on how to extrapolate this to 2 Dimensions.

Thanks!

Upvotes: 0

Views: 845

Answers (1)

Alexander Korovin
Alexander Korovin

Reputation: 1475

The easiest way is to use einsum. In the case of r_thresh with the shape of (N_z,), you can use this code:

res = np.einsum("i,j->ij", np.arange(N_r) + 0.5, r_thresh)

Also, you can reshape np.arange(N_r) + 0.5 to the shape (N_r,1) and r_thresh to the shape (1,N_z). Thus, you can use the dot product (for Python version > 3.5):

res = (np.arange(N_r) + 0.5).reshape(N_r,1) @ r_thresh.reshape(1,N_z)

or following to the comment of hpaulj:

res = (np.arange(N_r) + 0.5)[:,None] @ r_thresh[None,:]

EDIT1

The comment of hpaulj is also very helpful (I pasted this into my answer to see better):

res = (np.arange(N_r) + 0.5)[:,None] * r_thresh
res = np.outer(np.arange(N_r) + 0.5, r_thresh)

IN ADDITION

You can also use tensordot:

res = np.tensordot((np.arange(N_r) + 0.5)[:,None], r_thresh[:,None], axes=[[-1],[-1]])

Upvotes: 2

Related Questions