From d3edc964bc16a02e9ac5382008c4139390c6b8ba Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Mon, 23 Oct 2023 19:55:06 -0700 Subject: [PATCH 1/7] add `.imag` and `.real` properties and `.astype` method to `NamedArray` --- xarray/namedarray/core.py | 50 +++++++++++++++++++++++++++++---- xarray/tests/test_namedarray.py | 16 +++++++---- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 4eed572cbd2..83a64790b22 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -24,16 +24,11 @@ from xarray.namedarray._typing import ( _arrayfunction_or_api, _chunkedarray, - _DType, _DType_co, _ScalarType_co, _ShapeType_co, ) -from xarray.namedarray.utils import ( - _default, - is_duck_dask_array, - to_0d_object_array, -) +from xarray.namedarray.utils import _default, is_duck_dask_array, to_0d_object_array if TYPE_CHECKING: from numpy.typing import ArrayLike, NDArray @@ -46,6 +41,7 @@ _Dim, _Dims, _DimsLike, + _DType, _IntOrUnknown, _ScalarType, _Shape, @@ -516,6 +512,48 @@ def data(self, data: duckarray[Any, _DType_co]) -> None: self._check_shape(data) self._data = data + @property + def imag(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + """ + The imaginary part of the array. + + See Also + -------- + numpy.ndarray.imag + """ + + from xarray.namedarray._array_api import imag + + return imag(self) + + @property + def real(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + """ + The real part of the array. + + See Also + -------- + numpy.ndarray.real + """ + + from xarray.namedarray._array_api import real + + return real(self) + + def astype( + self, dtype: _DType, copy: bool = True + ) -> NamedArray[_ShapeType, _DType]: + """ + Copy of the array, cast to a specified type. + + See Also + -------- + numpy.ndarray.astype + """ + from xarray.namedarray._array_api import astype + + return astype(self, dtype, copy=copy) + def __dask_tokenize__(self) -> Hashable: # Use v.data, instead of v._data, in order to cope with the wrappers # around NetCDF and the like diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 93bff4d6a05..fe41b097c92 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -8,11 +8,7 @@ import pytest from xarray.core.indexing import ExplicitlyIndexed -from xarray.namedarray._typing import ( - _arrayfunction_or_api, - _DType_co, - _ShapeType_co, -) +from xarray.namedarray._typing import _arrayfunction_or_api, _DType_co, _ShapeType_co from xarray.namedarray.core import NamedArray, from_array from xarray.namedarray.utils import _default @@ -171,6 +167,16 @@ def test_data(random_inputs: np.ndarray[Any, Any]) -> None: named_array.data = np.random.random((3, 4)).astype(np.float64) +def test_real_and_imag() -> None: + named_array: NamedArray[Any, Any] + named_array = NamedArray(["x"], np.arange(3) - 1j * np.arange(3)) + expected_real = np.arange(3) + assert np.array_equal(named_array.real.data, expected_real) + + expected_imag = -np.arange(3) + assert np.array_equal(named_array.imag.data, expected_imag) + + # Additional tests as per your original class-based code @pytest.mark.parametrize( "data, dtype", From 1991fb3265afe4622228c694432ae9798a4d4c29 Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Mon, 23 Oct 2023 20:02:56 -0700 Subject: [PATCH 2/7] remove `.imag` and `.real` --- xarray/core/variable.py | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 9601ad49d05..ce368dc445a 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -2365,28 +2365,6 @@ def notnull(self, keep_attrs: bool | None = None): keep_attrs=keep_attrs, ) - @property - def real(self): - """ - The real part of the variable. - - See Also - -------- - numpy.ndarray.real - """ - return self._replace(data=self.data.real) - - @property - def imag(self): - """ - The imaginary part of the variable. - - See Also - -------- - numpy.ndarray.imag - """ - return self._replace(data=self.data.imag) - def __array_wrap__(self, obj, context=None): return Variable(self.dims, obj) From 01396c8a779ae13ad109e0adfcc3b7299f0c82e4 Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Mon, 23 Oct 2023 23:29:31 -0700 Subject: [PATCH 3/7] revert variable changes --- xarray/core/variable.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index ce368dc445a..53f74b64522 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -2365,6 +2365,30 @@ def notnull(self, keep_attrs: bool | None = None): keep_attrs=keep_attrs, ) + @property + def real(self): + """ + The real part of the variable. + See Also + -------- + numpy.ndarray.real + """ + # TODO: this should be removed (in favor of namedarray/_array_api) once there's wide support for Python Array API + # https://github.com/pydata/xarray/pull/8365#issuecomment-1776593216 + return self._replace(data=self.data.real) + + @property + def imag(self): + """ + The imaginary part of the variable. + See Also + -------- + numpy.ndarray.imag + """ + # TODO: this should be removed (in favor of namedarray/_array_api) once there's wide support for Python Array API + # https://github.com/pydata/xarray/pull/8365#issuecomment-1776593216 + return self._replace(data=self.data.imag) + def __array_wrap__(self, obj, context=None): return Variable(self.dims, obj) From 2193130a3d97717d2ffd9f3b9a0da0c54ae2e9ee Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Tue, 24 Oct 2023 03:03:17 -0700 Subject: [PATCH 4/7] move .imag and .real --- xarray/core/variable.py | 24 ------------------------ xarray/namedarray/core.py | 18 +++++------------- 2 files changed, 5 insertions(+), 37 deletions(-) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 53f74b64522..ce368dc445a 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -2365,30 +2365,6 @@ def notnull(self, keep_attrs: bool | None = None): keep_attrs=keep_attrs, ) - @property - def real(self): - """ - The real part of the variable. - See Also - -------- - numpy.ndarray.real - """ - # TODO: this should be removed (in favor of namedarray/_array_api) once there's wide support for Python Array API - # https://github.com/pydata/xarray/pull/8365#issuecomment-1776593216 - return self._replace(data=self.data.real) - - @property - def imag(self): - """ - The imaginary part of the variable. - See Also - -------- - numpy.ndarray.imag - """ - # TODO: this should be removed (in favor of namedarray/_array_api) once there's wide support for Python Array API - # https://github.com/pydata/xarray/pull/8365#issuecomment-1776593216 - return self._replace(data=self.data.imag) - def __array_wrap__(self, obj, context=None): return Variable(self.dims, obj) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 83a64790b22..f939196a30d 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -513,7 +513,7 @@ def data(self, data: duckarray[Any, _DType_co]) -> None: self._data = data @property - def imag(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + def imag(self) -> Self: """ The imaginary part of the array. @@ -521,13 +521,10 @@ def imag(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: -------- numpy.ndarray.imag """ - - from xarray.namedarray._array_api import imag - - return imag(self) + return self._replace(data=self._data.imag) @property - def real(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + def real(self) -> Self: """ The real part of the array. @@ -535,14 +532,9 @@ def real(self) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: -------- numpy.ndarray.real """ + return self._replace(data=self._data.real) - from xarray.namedarray._array_api import real - - return real(self) - - def astype( - self, dtype: _DType, copy: bool = True - ) -> NamedArray[_ShapeType, _DType]: + def astype(self, dtype: _DType, copy: bool = True) -> Self: """ Copy of the array, cast to a specified type. From 5358670ed49aa134b1056ee3cebd3cdb41771856 Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Tue, 24 Oct 2023 03:07:04 -0700 Subject: [PATCH 5/7] use .data --- xarray/namedarray/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index f939196a30d..8f5b7a6416c 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -521,7 +521,7 @@ def imag(self) -> Self: -------- numpy.ndarray.imag """ - return self._replace(data=self._data.imag) + return self._replace(data=self.data.imag) @property def real(self) -> Self: @@ -532,7 +532,7 @@ def real(self) -> Self: -------- numpy.ndarray.real """ - return self._replace(data=self._data.real) + return self._replace(data=self.data.real) def astype(self, dtype: _DType, copy: bool = True) -> Self: """ From 8d77be1d069a1d259ef2897db61c789626cee249 Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Tue, 24 Oct 2023 03:13:00 -0700 Subject: [PATCH 6/7] remove astype for the time being --- xarray/namedarray/core.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 8f5b7a6416c..d5c56dd1947 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -534,18 +534,6 @@ def real(self) -> Self: """ return self._replace(data=self.data.real) - def astype(self, dtype: _DType, copy: bool = True) -> Self: - """ - Copy of the array, cast to a specified type. - - See Also - -------- - numpy.ndarray.astype - """ - from xarray.namedarray._array_api import astype - - return astype(self, dtype, copy=copy) - def __dask_tokenize__(self) -> Hashable: # Use v.data, instead of v._data, in order to cope with the wrappers # around NetCDF and the like From c2b12cfda864673b1a3b60c79f8c013653ed9e15 Mon Sep 17 00:00:00 2001 From: Anderson Banihirwe Date: Tue, 24 Oct 2023 03:26:19 -0700 Subject: [PATCH 7/7] silence mypy --- xarray/namedarray/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index d5c56dd1947..726a3eaf8cb 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -521,7 +521,7 @@ def imag(self) -> Self: -------- numpy.ndarray.imag """ - return self._replace(data=self.data.imag) + return self._replace(data=self.data.imag) # type: ignore @property def real(self) -> Self: @@ -532,7 +532,7 @@ def real(self) -> Self: -------- numpy.ndarray.real """ - return self._replace(data=self.data.real) + return self._replace(data=self.data.real) # type: ignore def __dask_tokenize__(self) -> Hashable: # Use v.data, instead of v._data, in order to cope with the wrappers