diff --git a/xarray/core/variable.py b/xarray/core/variable.py index ce368dc445a..f18c4044f40 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -2365,6 +2365,28 @@ def notnull(self, keep_attrs: bool | None = None): keep_attrs=keep_attrs, ) + @property + def imag(self) -> Variable: + """ + The imaginary part of the variable. + + See Also + -------- + numpy.ndarray.imag + """ + return self._new(data=self.data.imag) + + @property + def real(self) -> Variable: + """ + The real part of the variable. + + See Also + -------- + numpy.ndarray.real + """ + return self._new(data=self.data.real) + def __array_wrap__(self, obj, context=None): return Variable(self.dims, obj) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 820371a7463..b9173c49003 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -138,6 +138,14 @@ def __array_function__( ) -> Any: ... + @property + def imag(self) -> _arrayfunction[_ShapeType_co, Any]: + ... + + @property + def real(self) -> _arrayfunction[_ShapeType_co, Any]: + ... + # Corresponds to np.typing.NDArray: _ArrayFunction = _arrayfunction[Any, np.dtype[_ScalarType_co]] diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 726a3eaf8cb..eba3f2f714b 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -22,11 +22,15 @@ from xarray.core import dtypes, formatting, formatting_html from xarray.namedarray._aggregations import NamedArrayAggregations from xarray.namedarray._typing import ( + _arrayapi, _arrayfunction_or_api, _chunkedarray, + _dtype, _DType_co, _ScalarType_co, _ShapeType_co, + _SupportsImag, + _SupportsReal, ) from xarray.namedarray.utils import _default, is_duck_dask_array, to_0d_object_array @@ -513,7 +517,9 @@ def data(self, data: duckarray[Any, _DType_co]) -> None: self._data = data @property - def imag(self) -> Self: + def imag( + self: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] + ) -> NamedArray[_ShapeType, _dtype[_ScalarType]]: """ The imaginary part of the array. @@ -521,10 +527,17 @@ def imag(self) -> Self: -------- numpy.ndarray.imag """ - return self._replace(data=self.data.imag) # type: ignore + if isinstance(self._data, _arrayapi): + from xarray.namedarray._array_api import imag + + return imag(self) + + return self._new(data=self._data.imag) @property - def real(self) -> Self: + def real( + self: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] + ) -> NamedArray[_ShapeType, _dtype[_ScalarType]]: """ The real part of the array. @@ -532,7 +545,11 @@ def real(self) -> Self: -------- numpy.ndarray.real """ - return self._replace(data=self.data.real) # type: ignore + if isinstance(self._data, _arrayapi): + from xarray.namedarray._array_api import real + + return real(self) + return self._new(data=self._data.real) def __dask_tokenize__(self) -> Hashable: # Use v.data, instead of v._data, in order to cope with the wrappers diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index fe41b097c92..6e39a3aa94f 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -168,13 +168,25 @@ def test_data(random_inputs: np.ndarray[Any, Any]) -> None: 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_real: np.ndarray[Any, np.dtype[np.float64]] + expected_real = np.arange(3, dtype=np.float64) + + expected_imag: np.ndarray[Any, np.dtype[np.float64]] + expected_imag = -np.arange(3, dtype=np.float64) + + arr: np.ndarray[Any, np.dtype[np.complex128]] + arr = expected_real + 1j * expected_imag + + named_array: NamedArray[Any, np.dtype[np.complex128]] + named_array = NamedArray(["x"], arr) + + actual_real: duckarray[Any, np.dtype[np.float64]] = named_array.real.data + assert np.array_equal(actual_real, expected_real) + assert actual_real.dtype == expected_real.dtype - expected_imag = -np.arange(3) - assert np.array_equal(named_array.imag.data, expected_imag) + actual_imag: duckarray[Any, np.dtype[np.float64]] = named_array.imag.data + assert np.array_equal(actual_imag, expected_imag) + assert actual_imag.dtype == expected_imag.dtype # Additional tests as per your original class-based code @@ -347,7 +359,9 @@ def _new( def test_replace_namedarray() -> None: dtype_float = np.dtype(np.float32) + np_val: np.ndarray[Any, np.dtype[np.float32]] np_val = np.array([1.5, 3.2], dtype=dtype_float) + np_val2: np.ndarray[Any, np.dtype[np.float32]] np_val2 = 2 * np_val narr_float: NamedArray[Any, np.dtype[np.float32]]