pythonnumpymasked-array

How to create a numpy masked array using lower dim array as a mask?


Assume

a = [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ]

mask = [1, 0, 1]

I want

a[mask] == [
           [1, 2, 3],
           [False, False, False],
           [7, 8, 9],
           ]

or equivalent.

Meaning, I want to access a with mask, where mask is of lower dimension, and have auto broadcasting. I want to use that in the constructor of np.ma.array in the mask= argument.


Solution

  • This should work. Note that your mask has the opposite meaning of np.ma.masked_array, where 1 means removed, so I inverted your mask:

    >>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    >>> mask = ~np.array([1, 0, 1], dtype=np.bool)  # Note - inverted mask.
    >>> masked_a = np.ma.masked_array(
    ...     a,
    ...     np.repeat(mask, a.shape[1]).reshape(a.shape)
    ... )
    >>> masked_a
    masked_array(
      data=[[1, 2, 3],
            [--, --, --],
            [7, 8, 9]],
      mask=[[False, False, False],
            [ True,  True,  True],
            [False, False, False]],
      fill_value=999999)