Skip to content

Commit

Permalink
add .imag and .real properties to NamedArray (#8365)
Browse files Browse the repository at this point in the history
  • Loading branch information
andersy005 authored Oct 24, 2023
1 parent eb74944 commit ccc8f99
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 33 deletions.
22 changes: 0 additions & 22 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
30 changes: 24 additions & 6 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -46,6 +41,7 @@
_Dim,
_Dims,
_DimsLike,
_DType,
_IntOrUnknown,
_ScalarType,
_Shape,
Expand Down Expand Up @@ -516,6 +512,28 @@ def data(self, data: duckarray[Any, _DType_co]) -> None:
self._check_shape(data)
self._data = data

@property
def imag(self) -> Self:
"""
The imaginary part of the array.
See Also
--------
numpy.ndarray.imag
"""
return self._replace(data=self.data.imag) # type: ignore

@property
def real(self) -> Self:
"""
The real part of the array.
See Also
--------
numpy.ndarray.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
# around NetCDF and the like
Expand Down
16 changes: 11 additions & 5 deletions xarray/tests/test_namedarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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",
Expand Down

0 comments on commit ccc8f99

Please sign in to comment.