Efficient memory usage with numpy masked arrays

Question

I have a large ndarray X (roughly (1e3, 1e3, 1e3)), where I want to do manipulations of X including and not including particular elements of the 0th axis (for each element of the 1st and 2nd axes). i.e. there are (1e3, 1e3) elements which I want to (at times) mask in or out.

The simplest thing to do would be to construct a masked array like,

Z = np.zeros_like(X, dtype=bool)
# assume `inds` is some indexing array which will target
#    the particular (1e3 x 1e3) elements I'm interested in
Z[inds] = True
Y = np.ma.masked_array(X, mask=Z)

But this uses an extra gigabyte of memory just for the masking array. Is there any way to do this without constructing a second 10^9 element array of masks? For example, is it possible to construct a sparse-matrix for the mask?

Answer 1

If you just want to take "clean" slices, as opposed to only taking some elements from some "rows", then you could use numeric indices instead of a mask.

E.g.:

arr = np.array([[[1,2,3,4], [5,6,7,8]], [[9,8,9,8], [7,6,7,6]]])
sub_idx = np.array([0,2])
sub_arr = arr[:, :, sub_idx]

This is a copy of a subset of arr, namely the 0th and 2nd "slices" in the last dimension:

array([[[1, 3],
        [5, 7]],

       [[9, 9],
        [7, 7]]])

Note that the array that defines which indexes to use is only one-dimensional, severely reducing its memory requirements. (Though of course the copy still takes up a significant chunk of memory in your case.)

Also note that this gives you a copy, so any changes you make to the result (sub_arr) do not manifest in the original array. To do that, you'd have to copy the array back over:

sub_arr[:] = 0 # Manipulate the values
arr[sub_idx] = sub_arr