Initialise a NumPy array based on its index

I think there are a couple of things to consider:

  • is there a reason that pos_data has to be a list?
  • don't have another variable (d) that you have to hard code, when it is always supposed to be the length of some other tuple.

With these in mind, you can solve your problem of variable numbers of for loops using itertools.product, which basically is just a shorthand for nested for loops. The positional args for product are the ranges of the loops.

My implementation is:

from itertools import product

vol_ext = (1000, 500)  # If d = 3, this will have another entry
ratio = [5.0, 8.0]  # Again, if d = 3, it will have another entry

pos_data_new = np.zeros((len(ratio), *vol_ext))

# now loop over each dimension in `vol_ext`. Since `product` expects
# positional arguments, we have to unpack a tuple of `range(vol)`s.
for inds in product(*(range(vol) for vol in vol_ext)):
    # inds is now a tuple, and we have to combine it with a slice in 
    # in the first dimension, and use it as an array on the right hand 
    # side to do the computation. 
    pos_data_new[(slice(None),) + inds] = (np.array(inds) - 1) * ratio

I don't think this will be any faster, but it certainly looks nicer.

Note that pos_data_new is now an array, to get it as a list in the first dimension, as per the original example, is simple enough.


It's easy with np.meshgrid:

pos_data = np.meshgrid(*(r * (np.arange(s) - 1.0)
                         for s, r in zip(vol_ext, ratio)), indexing='ij')

I would generate a two or three dimensional numpy.meshgrid of data, then scale each entry by the ratio per slice.

For the 2D case:

(X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

For the 3D case:

(X, Y, Z) = np.meshgrid(np.arange(vol_ext[2]), np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Z - 1) * ratio[0], (Y - 1) * ratio[1], (X - 1) * ratio[2]]

Example using your 2D data

pos_data has been generated by your code. I've created a new list pos_data2 that stores the equivalent list using the above solution:

In [40]: vol_ext = (1000, 500)

In [41]: (X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))

In [42]: pos_data2 = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

In [43]: np.allclose(pos_data[0], pos_data2[0])
Out[43]: True

In [44]: np.allclose(pos_data[1], pos_data2[1])
Out[44]: True

Making this adaptive based on vol_ext

We can combine this with a list comprehension where we can take advantage of the fact that the output of numpy.meshgrid is a tuple:

pts = [np.arange(v) for v in reversed(vol_ext)]
pos_data = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

The first line of code generates the range of points per desired dimension as a list. We then use a list comprehension to compute the desired calculations per slice by iterating over each desired grid of points in the desired dimension combined with the correct ratio to apply.

Example Run

In [49]: pts = [np.arange(v) for v in reversed(vol_ext)]

In [50]:  pos_data2 = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

In [51]: np.all([np.allclose(p, p2) for (p, p2) in zip(pos_data, pos_data2)])
Out[51]: True

The last line goes through each slice and ensures both lists align.