numpy: how to use masked_outside on certain columns of a numpy ndarray?

Question

What I am trying to do is to use masked_outside function to mask out value that is not in a range in a given ndarray, as

import numpy as np
import numpy.ma as ma

a = np.zeros((3,3))
a[1,1] = -1
a[2,1] = 1
a[0,2] = 1

b = ma.masked_outside(a, 0, 0)

then I get:

a = array([[ 0.,  0.,  1.],
           [ 0., -1.,  0.],
           [ 0.,  1.,  0.]])

b = masked_array(data =
  [[0.0 0.0 --]
   [0.0 -- 0.0]
   [0.0 -- 0.0]],
                 mask =
  [[False False  True]
   [False  True False]
   [False  True False]],
   fill_value = 1e+20)

However, I want to EXCLUDE a certain column from masking, something like:

 b = ma.masked_outside(a, 0, 0, exclude_cols=[2, ])

How can I achieve this?

What about when the array has a non-trivial dtype, ie array with named fields?

Answer 1

I don't think there's anything built into numpy, but it shouldn't be too hard to manually build up the mask. Does something like this work? The function builds a boolean array like masked_outside does, and fills excluded columns with False

In [66]: def make_mask_outside(data, lower, upper, exclude_cols=None):
    ...:     mask = (data < lower) | (data > upper)
    ...:     if exclude_cols is not None:
    ...:         for c in exclude_cols:
    ...:             mask[:,c] = False
    ...:     return mask

In [67]: b = ma.masked_array(a, make_mask_outside(a, 0, 0, [2,]))

In [68]: b
Out[68]: 
masked_array(data =
 [[0.0 0.0 1.0]
 [0.0 -- 0.0]
 [0.0 -- 0.0]],
             mask =
 [[False False False]
 [False  True False]
 [False  True False]],
       fill_value = 1e+20)

Answer 2

You can use the mask with multiple conditions, and in one of them you consider the column number(s) that you want to exclude:

rows, cols = np.indices(a.shape)
b = np.ma.array(a, mask=((a!=0) & (cols!=2)))

#masked_array(data =
# [[0.0 0.0 1.0]
# [0.0 -- 0.0]
# [0.0 -- 0.0]],
#             mask =
# [[False False False]
# [False  True False]
# [False  True False]],
#       fill_value = 1e+20)

This could be extended to more columns like: mask=((a!=0) & (cols!=2) & (cols!=3)) and so forth.