From 72f88cdc0deed59e99110daa084fa8e79f16db1f Mon Sep 17 00:00:00 2001
From: Eric Wieser <wieser.eric@gmail.com>
Date: Sat, 21 Jan 2017 00:04:17 +0000
Subject: [PATCH 1/2] MAINT: Rewrite np.ma.apply_along_axis in terms of
 apply_along_axis

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.
---
 numpy/ma/extras.py            | 82 +++--------------------------------
 numpy/ma/tests/test_extras.py | 15 +++++++
 2 files changed, 20 insertions(+), 77 deletions(-)

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 7149b525bad1..31c5c4d7dc06 100644
--- 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
index fb68bd261514..30bc230a5e58 100644
--- 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

From 39d105eb9e5f15cba5cf358c18f0e03b5c5e8fee Mon Sep 17 00:00:00 2001
From: Eric Wieser <wieser.eric@gmail.com>
Date: Mon, 28 May 2018 19:30:10 -0700
Subject: [PATCH 2/2] fixup last commit

---
 numpy/ma/extras.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 31c5c4d7dc06..24ae075654b7 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -381,7 +381,7 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     """
     def wrapped_func(a, *args, **kwargs):
         res = func1d(a, *args, **kwargs)
-        return np.asanyarray(res).view(masked_array)
+        return np.ma.asanyarray(res)
 
     return np.apply_along_axis(wrapped_func, axis, arr, *args, **kwargs)
 apply_along_axis.__doc__ = np.apply_along_axis.__doc__