indexing into numpy's mgrid

Viewed 2519

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.

1 Answers
Related