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..4420bdf1c --- /dev/null +++ b/cunumeric/ma/_masked_array.py @@ -0,0 +1,88 @@ +# 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 + +from typing import TYPE_CHECKING, Any, Type, Union + +if TYPE_CHECKING: + import numpy.typing as npt + from ..types import NdShape + + +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: + _internal_ma: _np.ma.MaskedArray[Any, Any] + + def __new__(cls: Type[Any], *args: Any, **kw: Any) -> MaskedArray: + return super().__new__(cls) + + def __init__( + self, + data: Any = None, + mask: _np.bool_ = nomask, + dtype: Union[npt.DTypeLike, None] = None, + copy: bool = False, + subok: bool = True, + ndmin: int = 0, + fill_value: Any = None, + keep_mask: Any = True, + hard_mask: Any = None, + shrink: bool = True, + order: Union[str, None] = None, + ) -> None: + self._internal_ma = _np.ma.MaskedArray( # type: ignore + 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: Any = None) -> _np.ma.MaskedArray[Any, Any]: + return self._internal_ma + + @property + def size(self) -> int: + return self._internal_ma.size + + @property + def shape(self) -> NdShape: + return self._internal_ma.shape