Is there a better way to reference the indices (coordinates) of a numpy array?

Viewed 81

I'm looking to find a vectorised way to create ndarrays from formulae, using the indices (or coordinates) of the value being calculated.

For example, if I want a 4x5x3 array filled by the formula 3x+y^z, I have no current way to reference x, y, or z directly. The closest I can come to that is wasting memory by creating incrementing arrays through arange().

Here are my current methods:

import numpy as np


# for loops

arr = np.empty((4, 5, 3))
for x in range(4):
    for y in range(5):
        for z in range(3):
            arr[x, y, z] = 3 * x + y ** z
print(arr)


# separate arange()-created arrays for x, y, and z

x = np.arange(4)[:, None, None]
y = np.arange(5)[None, :, None]
z = np.arange(3)[None, None, :]
arr = 3 * x + y ** z
print(arr)


# same thing but written in a different way (I think)

all = np.arange(5)
arr = 3 * all[:4, None, None] + all[None, :, None] ** all[None, None, :3]
print(arr)

Is there any more efficient way to do this? I'm assuming there is, because numpy's underlying C code must use an iterator to find the memory addresses of each item, so surely that iterator could be repurposed as a faster alternative to using indices stored in memory?

0 Answers
Related