Question

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 3×2 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.

Answer 1

As HYRY mentioned, I believe what you want to avoid is creating the full arrays. mgrid creates a full array, however if you use:

x, y = np.broadcast_arrays(*np.ogrid[0:2,0:3])

x and y take up no more memory then np.arange(0,2) (and np.arange(0,3)), while acting as if each was a full array. If you require a single large result array, you should probably slice these arrays individually and then concatenate them. (np.broadcast_arrays returns a tuple of arrays instead of an array)