Equivalent of np.array split for masked array

Question

numpy.ma.split_array does not exist.

Consequently, does the following code works as intended if arr is a masked array?

np.array_split(arr, multiprocessing.cpu_count())

If not, how should I define a function split_masked_array to achieve similar behavior?

Answer 1

Not sure why you didn't just try. Seems to work, but it's hard to tell for sure since you didn't provide a minimal reproducible example so I don't know what "works as intended" means.

In [5]: x = np.ma.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])

In [6]: x
Out[6]:
masked_array(
  data=[[--, 2, 3],
        [4, --, 6],
        [7, 8, --]],
  mask=[[ True, False, False],
        [False,  True, False],
        [False, False,  True]],
  fill_value=999999)

In [7]: np.array_split(x, 3)
Out[7]:
[masked_array(data=[[--, 2, 3]],
              mask=[[ True, False, False]],
        fill_value=999999),
 masked_array(data=[[4, --, 6]],
              mask=[[False,  True, False]],
        fill_value=999999),
 masked_array(data=[[7, 8, --]],
              mask=[[False, False,  True]],
        fill_value=999999)]