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.