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..1cff0a0ac 100644 --- a/cunumeric/array.py +++ b/cunumeric/array.py @@ -178,7 +178,8 @@ 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..54be15d3c --- /dev/null +++ b/cunumeric/ma/_masked_array.py @@ -0,0 +1,60 @@ +# 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 \ No newline at end of file