Clarifying the mask propogation in np.ma.polyfit in documentation

I would like to add an example in the documentation of np.ma.polyfit to demonstrate that the 2D mask is collapsed on 1D mask before doing the polyfit of a masked array. The current documentation is not very clear on how np.ma.polyfit deals with 2D mask while fitting a 2D y array.

For example.

A = np.ma.array([[0, 0, 0],
                 [1, 1, 1],
                 [2, 2, 2],
                 [10, 14, 99]])

x = np.arange(A.shape[0])

# only the mA[3,2] entry is masked in mA below
mA = np.ma.masked_greater(A,90)

print(mA)
# polynomial fit of all three columns of A
np.ma.polyfit(x,A,1)
# Outputs: array([[  3.1,   4.3,  29.8],
                 [ -1.4,  -2.2, -19.2]])

# polynomial fit of all three columns of mA
np.ma.polyfit(x,mA,1)
# Outputs: array([[  1.00000000e+00,   1.00000000e+00,   1.00000000e+00],
                       [  4.10073934e-17,   4.10073934e-17,   4.10073934e-17]])

# Instead of expected:  array([[  3.1,   4.3,  1.00000000e+00],
                     [ -1.4,  -2.2, 4.10073934e-17]])

Important Note: Masking one element in the last column affected the polynomial fit of all columns.

[mhvk]

@indiajoe - yes, indeed, although I hadn’t specifically considered np.unique (which is quite suitable in this case for returning the indices, but a bit unhandy in almost certainly requiring a copy of the mask). The only thing would be to generalize the .T to a .swapaxis(axis, 0) so one can have an arbitrary axis.

As for backwards compatibility and a possible keyword argument, perhaps best would be to ask on the mailing list – I do agree that the current behaviour is surprising, so more likely that people got wrong results out of it than that they rely on it to get correct ones.

[indiajoe]

More I think about it, I feel np.ma.polyfit should not be combining all the masks into a single union mask. I can’t think of any application where one would just want to throw away data points from the fitting of independent columns just because some columns have certain rows masked.

I understand the compromise one has to make is in the computation time, since one will have to loop over all columns to get unique masks for each column.

Wouldn’t it be better to slice the columns containing similar mask pattern and pass them as a set to np.polyfit? That way, we will have to loop over only the number of times column masks are different.

If others also agree to that approach, I could implement that and submit a pull request.