TST: Verify apply_along_axis now works on masked arrays

Note that this is not a full subsitute for np.ma.apply_along_axis,
as that allows functions to return a mix of np.ma.masked and scalars

1 files changed, 14 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index 111f302aa..530472a9b 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -120,6 +120,20 @@ class TestApplyAlongAxis(TestCase):
         self.test_0d_array(MinimalSubclass)
         self.test_axis_insertion(MinimalSubclass)
 
+    def test_axis_insertion_ma(self):
+        def f1to2(x):
+            """produces an assymmetric non-square matrix from x"""
+            assert_equal(x.ndim, 1)
+            res = x[::-1] * x[1:,None]
+            return np.ma.masked_where(res%5==0, res)
+        a = np.arange(6*3).reshape((6, 3))
+        res = apply_along_axis(f1to2, 0, a)
+        assert_(isinstance(res, np.ma.masked_array))
+        assert_equal(res.ndim, 3)
+        assert_array_equal(res[:,:,0].mask, f1to2(a[:,0]).mask)
+        assert_array_equal(res[:,:,1].mask, f1to2(a[:,1]).mask)
+        assert_array_equal(res[:,:,2].mask, f1to2(a[:,2]).mask)
+
     def test_tuple_func1d(self):
         def sample_1d(x):
             return x[1], x[0]