I use numpy.mgrid
to generate "coordinate index arrays"
y, x = np.mgrid[0:3, 0:2]
print x
array([[0, 1],
[0, 1],
[0, 1]])
In many situations, I take some slice through these arrays (e.g. x[0, :]
) and discard the rest of the data. Sometimes, these slices are much smaller than the original arrays, which are expensive to compute (i.e. np.mgrid[0:512, 0:512, 0:512]
). Does numpy provide an equivalent to [coord[view] for coord in np.mgrid[0:512, 0:512, 0:512]
that doesn't generate large intermediate arrays?
I realize the solution is trivial for the slice [0,:]
, but I'm looking for a general solution that handles any valid way to index numpy arrays
Edit
Some have asked for specific examples for what view
might look like. Ideally, I'm hoping for a general solution that handles any valid way to index a ndarray. Here are a few specific examples for the 3x2 array above:
1) view = (1, slice(None, None, 2))
2) view = (np.array([0,1]), np.array([0, 1]))
3) view = np.array([[False, False], [False, True], [False, False]])
And I'm looking for a function like
def mgrid_with_view(array_shape, view)
...
That returns the equivalent of [o[view] for o in np.indices(array_shape)]
without unnecessary computation or memory.