TST: test/fix astype/clip

marray/__init__.py

@@ -30,7 +30,8 @@ def __init__(self, data, mask=None):
             data = xp.asarray(getattr(data, '_data', data))
             mask = (xp.zeros(data.shape, dtype=xp.bool) if mask is None
                     else xp.asarray(mask, dtype=xp.bool))
-            mask = xp.asarray(xp.broadcast_to(mask, data.shape), copy=True)
+            if mask.shape != data.shape:  # avoid copy if possible
+                mask = xp.asarray(xp.broadcast_to(mask, data.shape), copy=True)
             self._data = data
             self._dtype = data.dtype
             self._device = data.device
@@ -203,7 +204,7 @@ def asarray(obj, /, *, mask=None, dtype=None, device=None, copy=None):
                 if mask is None else mask)
 
         data = xp.asarray(data, dtype=dtype, device=device, copy=copy)
-        mask = xp.asarray(mask, dtype=dtype, device=device, copy=copy)
+        mask = xp.asarray(mask, dtype=xp.bool, device=device, copy=copy)
         return MaskedArray(data, mask=mask)
     mod.asarray = asarray
 
@@ -233,7 +234,7 @@ def fun(*args, name=name, **kwargs):
     elementwise_names = ['abs', 'acos', 'acosh', 'add', 'asin', 'asinh', 'atan',
                          'atan2', 'atanh', 'bitwise_and', 'bitwise_left_shift',
                          'bitwise_invert', 'bitwise_or', 'bitwise_right_shift',
-                         'bitwise_xor', 'ceil', 'clip', 'conj', 'copysign', 'cos',
+                         'bitwise_xor', 'ceil', 'conj', 'copysign', 'cos',
                          'cosh', 'divide', 'equal', 'exp', 'expm1', 'floor',
                          'floor_divide', 'greater', 'greater_equal', 'hypot',
                          'imag', 'isfinite', 'isinf', 'isnan', 'less', 'less_equal',
@@ -252,6 +253,17 @@ def fun(*args, name=name, **kwargs):
             return MaskedArray(out, mask=xp.any(masks, axis=0))
         setattr(mod, name, fun)
 
+
+    def clip(x, /, min=None, max=None):
+        args = [x, min, max]
+        masks = [arg.mask for arg in args if hasattr(arg, 'mask')]
+        masks = xp.broadcast_arrays(*masks)
+        mask = xp.any(masks, axis=0)
+        datas = [getattr(arg, 'data', arg) for arg in args]
+        data = xp.clip(datas[0], min=datas[1], max=datas[2])
+        return MaskedArray(data, mask)
+    mod.clip = clip
+
     ## Indexing Functions
     # To be written:
     # take

marray/tests/test_marray.py

@@ -54,8 +54,8 @@ def get_arrays(n_arrays, *, dtype='float64', xp=np, seed=None):
 def assert_comparison(res, ref, seed, comparison, **kwargs):
     ref_mask = np.broadcast_to(ref.mask, ref.data.shape)
     try:
-        comparison(res.data[~res.mask], ref.data[~ref_mask], **kwargs)
-        comparison(res.mask, ref_mask, **kwargs)
+        comparison(res.data[~res.mask], ref.data[~ref_mask], strict=True, **kwargs)
+        comparison(res.mask, ref_mask, strict=True, **kwargs)
     except AssertionError as e:
         raise AssertionError(seed) from e
 
@@ -417,11 +417,41 @@ def test_statistical_array(f_name, keepdims, xp=np, dtype='float64', seed=None):
 
 # Use Array API tests to test the following:
 # Creation Functions (same behavior but with all-False mask)
-# Data Type Functions (only `astype` remains to be tested)
-# Elementwise function `clip` (all others are tested above)
 # Indexing (same behavior as indexing data and mask separately)
 # Manipulation functions (apply to data and mask separately)
 
+
+@pytest.mark.filterwarnings('ignore::numpy.exceptions.ComplexWarning')
+@pytest.mark.parametrize('dtype_in', dtypes_all)
+@pytest.mark.parametrize('dtype_out', dtypes_all)
+@pytest.mark.parametrize('copy', [False, True])
+def test_astype(dtype_in, dtype_out, copy, xp=np, seed=None):
+    mxp = marray.masked_array(xp)
+    marrays, masked_arrays, seed = get_arrays(1, dtype=dtype_in, seed=seed)
+
+    if dtype_in != dtype_out and not copy:
+        pytest.mark.skip("Can't change type without copy.")
+
+    res = mxp.astype(marrays[0], dtype_out, copy=copy)
+    if dtype_in == dtype_out:
+        if copy:
+            assert res.data is not marrays[0].data
+            assert res.mask is not marrays[0].mask
+        else:
+            assert res.data is marrays[0].data
+            assert res.mask is marrays[0].mask
+    ref = masked_arrays[0].astype(dtype_out, copy=copy)
+    assert_equal(res, ref, seed)
+
+
+@pytest.mark.parametrize('dtype', dtypes_real)
+def test_clip(dtype, xp=np, seed=None):
+    mxp = marray.masked_array(xp)
+    marrays, masked_arrays, seed = get_arrays(3, dtype=dtype, seed=seed)
+    res = mxp.clip(marrays[0], min=marrays[1], max=marrays[2])
+    ref = np.ma.clip(*masked_arrays)
+    assert_equal(res, ref, seed)
+
 #?
 # Searching functions - would test argmin/argmax with statistical functions,
 #                       but NumPy masked version isn't correct