Repository: dask/dask

Issue:
Mask preserving *_like functions
<!-- Please do a quick search of existing issues to make sure that this has not been asked before. -->
It would be useful to have versions of `ones_like`, `zeros_like` and `empty_like` that preserve masks when applied to masked dask arrays.  Currently (version 2022.7.1) we have

```python
import dask.array as da

array = da.ma.masked_array([2, 3, 4], mask=[0, 0, 1])
print(da.ones_like(array).compute())
```
```
[1 1 1]
```
whereas numpy's version preserves the mask
```python
import numpy as np

print(np.ones_like(array.compute()))
```
```
[1 1 --]
```

I notice there are several functions in `dask.array.ma` that just apply `map_blocks` to the `numpy.ma` version of the function.  So perhaps the simplest thing would be to implement `dask.array.ma.ones_like`, etc. that way.  If it really is that simple, I'd be happy to open a PR.


Discussion:
@rcomer Thanks for reporting!

> If it really is that simple, I'd be happy to open a PR.

I think this should mostly work.

However, I'd first like to confirm that this isn't related to  `https://github.com/numpy/numpy/issues/15200` -- what do you think about this? 

cc @jsignell for visibility
Thanks @pavithraes.  I'm afraid my understanding of the `__array_ufunc__` and `__array_function__` methods is not very deep.  However, looking at dask's current implementation of these `*_like` functions, they just pass straight to `ones`, etc, which I think can only ever return an unmasked array.
https://github.com/dask/dask/blob/50ab8af982a31b186a8e47597f0ad7e5b59bcab5/dask/array/creation.py#L124-L133

Edit: I am very new here so it's very possible that I am missing or misunderstanding something!
In principle this approach seems fine to me. I think masked arrays are kind of under-supported in Dask in general. So this kind of work is definitely appreciated!
Thanks @jsignell.  I'll put a PR up when I have some time.

Patch:
diff --git a/dask/array/ma.py b/dask/array/ma.py
--- 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)


Test patch:
diff --git a/dask/array/tests/test_masked.py b/dask/array/tests/test_masked.py
--- 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)