diff options
| -rw-r--r-- | numpy/lib/tests/test_shape_base.py | 14 |
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] |