numpy · takanori-pskq · Oct 7, 2020 · Oct 7, 2020 · Oct 8, 2020 · Oct 12, 2020
diff --git a/doc/release/upcoming_changes/17480.new_function.rst b/doc/release/upcoming_changes/17480.new_function.rst
@@ -0,0 +1,5 @@
+Add some ``numpy.ma.*`` methods
+-------------------------------
+The methods ``numpy.ma.ones_like``, ``numpy.ma.zeros_like``, ``numpy.ma.dsplit``,
+``numpy.ma.split``, ``numpy.ma.vsplit`` are now available. These methods are for
+``MaskedArray``, and behave the same as the original version (``numpy.*``).
diff --git a/doc/source/reference/routines.ma.rst b/doc/source/reference/routines.ma.rst
@@ -44,7 +44,9 @@ Ones and zeros
    ma.masked_all
    ma.masked_all_like
    ma.ones
+   ma.ones_like
    ma.zeros
+   ma.zeros_like
 
 
 _____
@@ -129,6 +131,7 @@ Changing the number of dimensions
 
    ma.MaskedArray.squeeze
 
+   ma.dsplit
    ma.stack
    ma.column_stack
    ma.concatenate
@@ -137,7 +140,9 @@ Changing the number of dimensions
    ma.hsplit
    ma.mr_
    ma.row_stack
+   ma.split
    ma.vstack
+   ma.vsplit
 
 
 Joining arrays

diff --git a/numpy/ma/__init__.pyi b/numpy/ma/__init__.pyi
@@ -139,6 +139,7 @@ nomask: Any
 nonzero: Any
 not_equal: Any
 ones: Any
+ones_like: Any
 outer: Any
 outerproduct: Any
 power: Any
@@ -178,6 +179,7 @@ true_divide: Any
 var: Any
 where: Any
 zeros: Any
+zeros_like: Any
 apply_along_axis: Any
 apply_over_axes: Any
 atleast_1d: Any
@@ -196,6 +198,7 @@ corrcoef: Any
 cov: Any
 diagflat: Any
 dot: Any
+dsplit: Any
 dstack: Any
 ediff1d: Any
 flatnotmasked_contiguous: Any
@@ -216,10 +219,12 @@ notmasked_contiguous: Any
 notmasked_edges: Any
 polyfit: Any
 row_stack: Any
+split: Any
 setdiff1d: Any
 setxor1d: Any
 stack: Any
 unique: Any
 union1d: Any
 vander: Any
+vsplit: Any
 vstack: Any
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -70,13 +70,13 @@
     'masked_singleton', 'masked_values', 'masked_where', 'max', 'maximum',
     'maximum_fill_value', 'mean', 'min', 'minimum', 'minimum_fill_value',
     'mod', 'multiply', 'mvoid', 'ndim', 'negative', 'nomask', 'nonzero',
-    'not_equal', 'ones', 'outer', 'outerproduct', 'power', 'prod',
+    'not_equal', 'ones', 'ones_like', 'outer', 'outerproduct', 'power', 'prod',
     'product', 'ptp', 'put', 'putmask', 'ravel', 'remainder',
     'repeat', 'reshape', 'resize', 'right_shift', 'round', 'round_',
     'set_fill_value', 'shape', 'sin', 'sinh', 'size', 'soften_mask',
     'sometrue', 'sort', 'sqrt', 'squeeze', 'std', 'subtract', 'sum',
     'swapaxes', 'take', 'tan', 'tanh', 'trace', 'transpose', 'true_divide',
-    'var', 'where', 'zeros',
+    'var', 'where', 'zeros', 'zeros_like',
     ]
 
 MaskType = np.bool_
@@ -8127,10 +8127,10 @@ def __call__(self, *args, **params):
     'identity', params=dict(fill_value=None, hardmask=False))
 indices = np.indices
 ones = _convert2ma('ones', params=dict(fill_value=None, hardmask=False))
-ones_like = np.ones_like
+ones_like = _convert2ma('ones_like')
 squeeze = np.squeeze
 zeros = _convert2ma('zeros', params=dict(fill_value=None, hardmask=False))
-zeros_like = np.zeros_like
+zeros_like = _convert2ma('zeros_like')
 
 
 def append(a, b, axis=None):

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
@@ -13,11 +13,12 @@
     'atleast_3d', 'average', 'clump_masked', 'clump_unmasked',
     'column_stack', 'compress_cols', 'compress_nd', 'compress_rowcols',
     'compress_rows', 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot',
-    'dstack', 'ediff1d', 'flatnotmasked_contiguous', 'flatnotmasked_edges',
+    'dsplit', 'dstack', 'ediff1d', 'flatnotmasked_contiguous', 'flatnotmasked_edges',
     'hsplit', 'hstack', 'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols',
     'mask_rows', 'masked_all', 'masked_all_like', 'median', 'mr_',
     'notmasked_contiguous', 'notmasked_edges', 'polyfit', 'row_stack',
-    'setdiff1d', 'setxor1d', 'stack', 'unique', 'union1d', 'vander', 'vstack',
+    'setdiff1d', 'setxor1d', 'split', 'stack', 'unique', 'union1d', 'vander', 'vsplit',
+    'vstack',
     ]
 
 import itertools
@@ -352,6 +353,9 @@ def __call__(self, *args, **params):
 stack = _fromnxfunction_seq('stack')
 
 hsplit = _fromnxfunction_single('hsplit')
+vsplit = _fromnxfunction_single('vsplit')
+dsplit = _fromnxfunction_single('dsplit')
+split = _fromnxfunction_args('split')
 
 diagflat = _fromnxfunction_single('diagflat')
 

diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
@@ -43,9 +43,9 @@
     masked_less, masked_less_equal, masked_not_equal, masked_outside,
     masked_print_option, masked_values, masked_where, max, maximum,
     maximum_fill_value, min, minimum, minimum_fill_value, mod, multiply,
-    mvoid, nomask, not_equal, ones, outer, power, product, put, putmask,
+    mvoid, nomask, not_equal, ones, ones_like, outer, power, product, put, putmask,
     ravel, repeat, reshape, resize, shape, sin, sinh, sometrue, sort, sqrt,
-    subtract, sum, take, tan, tanh, transpose, where, zeros,
+    subtract, sum, take, tan, tanh, transpose, where, zeros, zeros_like,
     )
 from numpy.compat import pickle
 
@@ -3186,6 +3186,72 @@ def test_empty(self):
         b = a.view(masked_array)
         assert_(np.may_share_memory(a.mask, b.mask))
 
    def test_zeros(self):
+        # Tests zeros
+        datatype = [('a', int), ('b', float), ('c', '|S8')]
+        a = masked_array([(1, 1.1, '1.1'), (2, 2.2, '2.2'), (3, 3.3, '3.3')],
+                         dtype=datatype)
+        assert_equal(len(a.fill_value.item()), len(datatype))
+
+        b = zeros(len(a), dtype=datatype)
+        assert_equal(b.shape, a.shape)
+        assert_equal(b.fill_value, a.fill_value)
+        assert_equal(b.data, array([(0, 0.0, b''), (0, 0.0, b''), (0, 0.0, b'')],
+            dtype=datatype))
+
+    def test_zeros_like(self):
+        # Tests zeros_like
+        datatype = [('a', int), ('b', float), ('c', '|S8')]
+        a = masked_array([(1, 1.1, '1.1'), (2, 2.2, '2.2'), (3, 3.3, '3.3')],
+                         dtype=datatype)
+        assert_equal(len(a.fill_value.item()), len(datatype))
+
+        b = zeros_like(a)
+        assert_equal(b.shape, a.shape)
+        assert_equal(b.fill_value, a.fill_value)
+        assert_equal(b.data, array([(0, 0.0, b''), (0, 0.0, b''), (0, 0.0, b'')],
+            dtype=datatype))
+
+        # check zeros_like mask handling
+        a = masked_array([1, 2, 3], mask=[False, True, False])
+        b = zeros_like(a)
+        assert_(not np.may_share_memory(a.mask, b.mask))
+        b = a.view(masked_array)
+        assert_(np.may_share_memory(a.mask, b.mask))
+
+    def test_ones(self):
+        # Tests ones
+        datatype = [('a', int), ('b', float), ('c', '|S8')]
+        a = masked_array([(1, 1.1, '1.1'), (2, 2.2, '2.2'), (3, 3.3, '3.3')],
+                         dtype=datatype)
+        assert_equal(len(a.fill_value.item()), len(datatype))
+
+        b = ones(len(a), dtype=datatype)
+        assert_equal(b.shape, a.shape)
+        assert_equal(b.fill_value, a.fill_value)
+        assert_equal(b.data, array([(1, 1.0, b'1'), (1, 1.0, b'1'), (1, 1.0, b'1')],
+            dtype=datatype))
+
+    def test_ones_like(self):
+        # Tests ones_like
+        datatype = [('a', int), ('b', float), ('c', '|S8')]
+        a = masked_array([(1, 1.1, '1.1'), (2, 2.2, '2.2'), (3, 3.3, '3.3')],
+                         dtype=datatype)
+        assert_equal(len(a.fill_value.item()), len(datatype))
+
+        b = ones_like(a)
+        assert_equal(b.shape, a.shape)
+        assert_equal(b.fill_value, a.fill_value)
+        assert_equal(b.data, array([(1, 1.0, b'1'), (1, 1.0, b'1'), (1, 1.0, b'1')],
+            dtype=datatype))
+
+        # check ones_like mask handling
+        a = masked_array([1, 2, 3], mask=[False, True, False])
+        b = ones_like(a)
+        assert_(not np.may_share_memory(a.mask, b.mask))
+        b = a.view(masked_array)
+        assert_(np.may_share_memory(a.mask, b.mask))
+
     @suppress_copy_mask_on_assignment
     def test_put(self):
         # Tests put.

diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
@@ -16,19 +16,19 @@
     assert_warns, suppress_warnings
     )
 from numpy.ma.testutils import (
-    assert_, assert_array_equal, assert_equal, assert_almost_equal
+    assert_, assert_array_equal, assert_equal, assert_almost_equal, assert_raises
     )
 from numpy.ma.core import (
     array, arange, masked, MaskedArray, masked_array, getmaskarray, shape,
     nomask, ones, zeros, count
     )
 from numpy.ma.extras import (
-    atleast_1d, atleast_2d, atleast_3d, mr_, dot, polyfit, cov, corrcoef,
+    atleast_1d, atleast_2d, atleast_3d, mr_, dot, dsplit, polyfit, cov, corrcoef,
     median, average, unique, setxor1d, setdiff1d, union1d, intersect1d, in1d,
     ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols,
     mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous,
     notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, isin,
-    diagflat, stack, vstack
+    hsplit, diagflat, split, stack, vsplit, vstack
     )
 
 
@@ -1686,3 +1686,94 @@ def test_stack_nd(self):
         assert_equal(c.shape, c_shp)
         assert_array_equal(a1.mask, c[..., 0].mask)
         assert_array_equal(a2.mask, c[..., 1].mask)
+
+class TestSplit:
+
+    def test_equal_split(self):
+        a = masked_array(np.arange(10), mask=[i%2 for i in range(10)]);
+        res = split(a, 2)
+        desired = masked_array([np.arange(5), np.arange(5, 10)],
+                mask=[[i%2 for i in range(5)], [i%2 for i in range(5, 10)]])
+        assert_equal(res.mask, desired.mask)
+        assert_equal(res.data, desired.data)
+
+    def test_unequal_split(self):
+        a = masked_array(np.arange(10), mask=[i%2 for i in range(10)])
+        assert_raises(ValueError, split, a, 3)
+
+class TestHsplit:
+    """Only testing for integer splits.
+
+    """
+    def test_0D_array(self):
+        a = masked_array(np.array(1), mask=[True])
+        assert_raises(ValueError, hsplit, a, 2)
+
+    def test_1D_array(self):
+        a = masked_array(np.array([1, 2, 3, 4]), mask=[1, 0, 1, 0])
+        res = hsplit(a, 2)
+        desired = masked_array(np.array([[1, 2], [3, 4]]), mask=[[1, 0], [1, 0]])
+        assert_equal(res.mask, desired.mask)
+        assert_equal(res.data, desired.data)
+
+    def test_2D_array(self):
+        a = masked_array(np.array([[1, 2, 3, 4], [1, 2, 3, 4]]),
+                mask=[[0, 1, 0, 1], [1, 0, 1, 0]]);
+        res = hsplit(a, 2)
+        desired = masked_array(np.array([[[1, 2], [1, 2]], [[3, 4], [3, 4]]]),
+                [[[0, 1], [1, 0]], [[0, 1], [1, 0]]]);
+        assert_equal(res.mask, desired.mask)
+        assert_equal(res.data, desired.data)
+
+
+class TestVsplit:
+    """Only testing for integer splits.
+
+    """
+    def test_0D_array(self):
+        a = masked_array(np.array(1), mask=[True])
+        assert_raises(ValueError, vsplit, a, 2)

+    def test_1D_array(self):
+        a = masked_array(np.array([1, 2, 3, 4]), mask=[1, 0, 1, 0])
+        assert_raises(ValueError, vsplit, a, 2)
+
+    def test_2D_array(self):
+        a = masked_array(np.array([[1, 2, 3, 4], [1, 2, 3, 4]]),
+                mask=[[0, 1, 0, 1], [1, 0, 1, 0]]);
+        res = vsplit(a, 2)
+        desired = masked_array(np.array([[[1, 2, 3, 4]], [[1, 2, 3, 4]]]),
+                mask=[[[0, 1, 0, 1]], [[1, 0, 1, 0]]]);
+        assert_equal(res.mask, desired.mask)
+        assert_equal(res.data, desired.data)
+
+
+class TestDsplit:
+    # Only testing for integer splits.
+    def test_0D_array(self):
+        a = masked_array(np.array(1), mask=[True])
+        assert_raises(ValueError, dsplit, a, 2)
+
+    def test_1D_array(self):
+        a = masked_array(np.array([1, 2, 3, 4]), mask=[1, 0, 1, 0])
+        assert_raises(ValueError, dsplit, a, 2)
+
+    def test_2D_array(self):
+        a = masked_array(np.array([[1, 2, 3, 4], [1, 2, 3, 4]]),
+                mask=[[0, 1, 0, 1], [1, 0, 1, 0]]);
+        assert_raises(ValueError, dsplit, a, 2)
+
+    def test_3D_array(self):
+        a = masked_array(
+                np.array([[[1, 2, 3, 4], [1, 2, 3, 4]],
+                    [[1, 2, 3, 4], [1, 2, 3, 4]]]),
+                mask=[[[1, 0 ,1, 0], [0, 1 ,0, 1]],
+                    [[0, 1, 0, 1], [1, 0, 1, 0]]]);
+        res = dsplit(a, 2)
+        desired = masked_array(
+                np.array([[[[1, 2], [1, 2]], [[1, 2], [1, 2]]],
+                    [[[3, 4], [3, 4]], [[3, 4], [3, 4]]]]),
+                mask=[[[[1, 0], [0, 1]], [[0, 1], [1, 0]]],
+                    [[[1, 0], [0, 1]], [[0, 1], [1, 0]]]]);
+        assert_equal(res.mask, desired.mask)
+        assert_equal(res.data, desired.data)