diff --git a/cunumeric/__init__.py b/cunumeric/__init__.py index 13c8504b8..3ad86dd9f 100644 --- a/cunumeric/__init__.py +++ b/cunumeric/__init__.py @@ -28,7 +28,7 @@ import numpy as _np -from cunumeric import linalg, random, fft +from cunumeric import linalg, random, fft, ma from cunumeric.array import maybe_convert_to_np_ndarray, ndarray from cunumeric.bits import packbits, unpackbits from cunumeric.module import * diff --git a/cunumeric/array.py b/cunumeric/array.py index a219e416a..59635309c 100644 --- a/cunumeric/array.py +++ b/cunumeric/array.py @@ -178,7 +178,9 @@ def maybe_convert_to_np_ndarray(obj: Any) -> Any: """ Converts cuNumeric arrays into NumPy arrays, otherwise has no effect. """ - if isinstance(obj, ndarray): + from .ma import MaskedArray + + if isinstance(obj, (ndarray, MaskedArray)): return obj.__array__() return obj diff --git a/cunumeric/ma/__init__.py b/cunumeric/ma/__init__.py new file mode 100644 index 000000000..14a9e0d46 --- /dev/null +++ b/cunumeric/ma/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2023 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +import numpy.ma as _ma + +from cunumeric.array import maybe_convert_to_np_ndarray +from cunumeric.coverage import clone_module +from cunumeric.ma._masked_array import MaskedArray + +masked_array = MaskedArray + +clone_module(_ma, globals(), maybe_convert_to_np_ndarray) + +del maybe_convert_to_np_ndarray +del clone_module +del _ma diff --git a/cunumeric/ma/_masked_array.py b/cunumeric/ma/_masked_array.py new file mode 100644 index 000000000..d3e9f3449 --- /dev/null +++ b/cunumeric/ma/_masked_array.py @@ -0,0 +1,79 @@ +# Copyright 2023 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +import numpy as _np + +from ..array import maybe_convert_to_np_ndarray +from ..coverage import clone_class + +NDARRAY_INTERNAL = { + "__array_finalize__", + "__array_function__", + "__array_interface__", + "__array_prepare__", + "__array_priority__", + "__array_struct__", + "__array_ufunc__", + "__array_wrap__", +} + +MaskType = _np.bool_ +nomask = MaskType(0) + + +@clone_class(_np.ma.MaskedArray, NDARRAY_INTERNAL, maybe_convert_to_np_ndarray) +class MaskedArray: + def __new__(cls, *args, **kw): + return super().__new__(cls) + + def __init__( + self, + data=None, + mask=nomask, + dtype=None, + copy=False, + subok=True, + ndmin=0, + fill_value=None, + keep_mask=True, + hard_mask=None, + shrink=True, + order=None, + ): + self._internal_ma = _np.ma.MaskedArray( + data=maybe_convert_to_np_ndarray(data), + mask=maybe_convert_to_np_ndarray(mask), + dtype=dtype, + copy=copy, + subok=subok, + ndmin=ndmin, + fill_value=fill_value, + keep_mask=keep_mask, + hard_mask=hard_mask, + shrink=shrink, + order=order, + ) + + def __array__(self, _dtype=None): + return self._internal_ma + + @property + def size(self): + return self._internal_ma.size + + @property + def shape(self): + return self._internal_ma.shape