Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add .imag and .real properties and .astype method to NamedArray
Browse files Browse the repository at this point in the history
andersy005 committed Oct 24, 2023

Verified

This commit was signed with the committer’s verified signature.
andersy005 Anderson Banihirwe
1 parent eb74944 commit d3edc96
Showing 2 changed files with 55 additions and 11 deletions.
50 changes: 44 additions & 6 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
@@ -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
16 changes: 11 additions & 5 deletions xarray/tests/test_namedarray.py
Original file line number Diff line number Diff line change
@@ -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",

0 comments on commit d3edc96

Please sign in to comment.