MAINT: Rewrite np.ma.apply_along_axis in terms of apply_along_axis

eric-wieser · eric-wieser · commit 72f88cdc0dee · 2017-02-21T21:17:55.000Z
This changes the behavior somewhat, but not in ways related to masked arrays
It would previously try to find a suitable dtype for all values, now it just uses
the first dtype.
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
@@ -379,83 +379,11 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     """
     (This docstring should be overwritten)
     """
-    arr = array(arr, copy=False, subok=True)
-    nd = arr.ndim
-    axis = normalize_axis_index(axis, nd)
-    ind = [0] * (nd - 1)
-    i = np.zeros(nd, 'O')
-    indlist = list(range(nd))
-    indlist.remove(axis)
-    i[axis] = slice(None, None)
-    outshape = np.asarray(arr.shape).take(indlist)
-    i.put(indlist, ind)
-    j = i.copy()
-    res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
-    #  if res is a number, then we have a smaller output array
-    asscalar = np.isscalar(res)
-    if not asscalar:
-        try:
-            len(res)
-        except TypeError:
-            asscalar = True
-    # Note: we shouldn't set the dtype of the output from the first result
-    # so we force the type to object, and build a list of dtypes.  We'll
-    # just take the largest, to avoid some downcasting
-    dtypes = []
-    if asscalar:
-        dtypes.append(np.asarray(res).dtype)
-        outarr = zeros(outshape, object)
-        outarr[tuple(ind)] = res
-        Ntot = np.product(outshape)
-        k = 1
-        while k < Ntot:
-            # increment the index
-            ind[-1] += 1
-            n = -1
-            while (ind[n] >= outshape[n]) and (n > (1 - nd)):
-                ind[n - 1] += 1
-                ind[n] = 0
-                n -= 1
-            i.put(indlist, ind)
-            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
-            outarr[tuple(ind)] = res
-            dtypes.append(asarray(res).dtype)
-            k += 1
-    else:
-        res = array(res, copy=False, subok=True)
-        j = i.copy()
-        j[axis] = ([slice(None, None)] * res.ndim)
-        j.put(indlist, ind)
-        Ntot = np.product(outshape)
-        holdshape = outshape
-        outshape = list(arr.shape)
-        outshape[axis] = res.shape
-        dtypes.append(asarray(res).dtype)
-        outshape = flatten_inplace(outshape)
-        outarr = zeros(outshape, object)
-        outarr[tuple(flatten_inplace(j.tolist()))] = res
-        k = 1
-        while k < Ntot:
-            # increment the index
-            ind[-1] += 1
-            n = -1
-            while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
-                ind[n - 1] += 1
-                ind[n] = 0
-                n -= 1
-            i.put(indlist, ind)
-            j.put(indlist, ind)
-            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
-            outarr[tuple(flatten_inplace(j.tolist()))] = res
-            dtypes.append(asarray(res).dtype)
-            k += 1
-    max_dtypes = np.dtype(np.asarray(dtypes).max())
-    if not hasattr(arr, '_mask'):
-        result = np.asarray(outarr, dtype=max_dtypes)
-    else:
-        result = asarray(outarr, dtype=max_dtypes)
-        result.fill_value = ma.default_fill_value(result)
-    return result
+    def wrapped_func(a, *args, **kwargs):
+        res = func1d(a, *args, **kwargs)
+        return np.asanyarray(res).view(masked_array)
+
+    return np.apply_along_axis(wrapped_func, axis, arr, *args, **kwargs)
 apply_along_axis.__doc__ = np.apply_along_axis.__doc__
 
 
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
@@ -640,6 +640,21 @@ def myfunc(b, offset=0):
         xa = apply_along_axis(myfunc, 2, a, offset=1)
         assert_equal(xa, [[2, 5], [8, 11]])
 
+    def test_mask(self):
+        x = np.arange(20).reshape(4, 5)
+        m = np.ma.masked_where(x%2 == 0, x)
+
+        # note that this lambda sometime returns np.ma.masked
+        # works fine with normal apply_along_axis
+        row0 = np.ma.apply_along_axis(lambda x: x[0], 0, m)
+
+        # fails with normal apply_along_axis, which doesn't use a masked_array
+        row1 = np.ma.apply_along_axis(lambda x: x[1], 0, m)
+
+        # note the first of these is a float, because that's the type of np.ma.masked...
+        assert_equal(row0, np.ma.masked_array([-1, 1.0, -1, 3.0, -1], [1, 0, 1, 0, 1]))
+        assert_equal(row1, np.ma.masked_array([5, -1, 7, -1, 9], [0, 1, 0, 1, 0]))
+
 
 class TestApplyOverAxes(TestCase):
     # Tests apply_over_axes
