Implement ma.*_like functions (#9378)

rcomer · web-flow · commit 0a7b58b134ba · 2022-08-17T09:46:04.000-04:00
Implement `ones_like`, `zeros_like` and `empty_like` within `dask.array.ma`.
diff --git a/dask/array/ma.py b/dask/array/ma.py
@@ -190,3 +190,21 @@ def count(a, axis=None, keepdims=False, split_every=None):
         split_every=split_every,
         out=None,
     )
+
+
+@derived_from(np.ma.core)
+def ones_like(a, **kwargs):
+    a = asanyarray(a)
+    return a.map_blocks(np.ma.core.ones_like, **kwargs)
+
+
+@derived_from(np.ma.core)
+def zeros_like(a, **kwargs):
+    a = asanyarray(a)
+    return a.map_blocks(np.ma.core.zeros_like, **kwargs)
+
+
+@derived_from(np.ma.core)
+def empty_like(a, **kwargs):
+    a = asanyarray(a)
+    return a.map_blocks(np.ma.core.empty_like, **kwargs)
diff --git a/dask/array/tests/test_masked.py b/dask/array/tests/test_masked.py
@@ -427,3 +427,22 @@ def test_count():
         res = da.ma.count(dx, axis=axis)
         sol = np.ma.count(x, axis=axis)
         assert_eq(res, sol, check_dtype=sys.platform != "win32")
+
+
+@pytest.mark.parametrize("funcname", ["ones_like", "zeros_like", "empty_like"])
+def test_like_funcs(funcname):
+    mask = np.array([[True, False], [True, True], [False, True]])
+    data = np.arange(6).reshape((3, 2))
+    a = np.ma.array(data, mask=mask)
+    d_a = da.ma.masked_array(data=data, mask=mask, chunks=2)
+
+    da_func = getattr(da.ma, funcname)
+    np_func = getattr(np.ma.core, funcname)
+
+    res = da_func(d_a)
+    sol = np_func(a)
+
+    if "empty" in funcname:
+        assert_eq(da.ma.getmaskarray(res), np.ma.getmaskarray(sol))
+    else:
+        assert_eq(res, sol)
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
@@ -356,6 +356,7 @@ Masked Arrays
    :toctree: generated/
 
    ma.average
+   ma.empty_like
    ma.filled
    ma.fix_invalid
    ma.getdata
@@ -372,7 +373,9 @@ Masked Arrays
    ma.masked_outside
    ma.masked_values
    ma.masked_where
+   ma.ones_like
    ma.set_fill_value
+   ma.zeros_like
 
 Random
 ~~~~~~
