From ff9c363bb729a37f09c63937e45d5870501bfbad Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Fri, 20 Jan 2017 18:56:41 +0000 Subject: 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 --- numpy/lib/tests/test_shape_base.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) (limited to 'numpy/lib/tests/test_shape_base.py') 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] -- cgit v1.2.1