Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add .imag and .real properties to NamedArray #8365

Merged
merged 7 commits into from
Oct 24, 2023
Merged

add .imag and .real properties to NamedArray #8365

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d3edc96
add `.imag` and `.real` properties and `.astype` method to `NamedArray`
andersy005 Oct 24, 2023
1991fb3
remove `.imag` and `.real`
andersy005 Oct 24, 2023
01396c8
revert variable changes
andersy005 Oct 24, 2023
2193130
move .imag and .real
andersy005 Oct 24, 2023
5358670
use .data
andersy005 Oct 24, 2023
8d77be1
remove astype for the time being
andersy005 Oct 24, 2023
c2b12cf
silence mypy
andersy005 Oct 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 0 additions & 22 deletions xarray/core/variable.py
View file
Open in desktop
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
View file
Open in desktop
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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When you use imag you're changing the dtype of the class, so it's not Self you're returning but rather a new NamedArray. Which won't match with self.

This should similar to the function, like:

def imag(
    self: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], /  # type: ignore[type-var]
) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:

Here's how numpy does it:
https://github.com/numpy/numpy/blob/14bb214bca49b167abc375fa873466a811e62102/numpy/__init__.pyi#L1495-L1500

Variable will also need this style of typing.

"""
The imaginary part of the array.

See Also
--------
numpy.ndarray.imag
"""
return self._replace(data=self.data.imag) # type: ignore
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only use ._replace if the shape or dtype wont change.

return self._new(data=self.data.imag)


@property
def real(self) -> Self:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same idea here. Follow the same idea but use _SupportsReal instead.

"""
The real part of the array.

See Also
--------
numpy.ndarray.real
"""
return self._replace(data=self.data.real) # type: ignore
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same idea, dtype or shape has changed in this method so use ._new.


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
View file
Open in desktop
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)
Comment on lines +170 to +177
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're not testing dtypes here. They don't match.

  • expected_imag is an int dtype.
  • named_array.imag.data is a float dtype.

The test should look something like this:

def test_real_and_imag() -> None:
    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: np.ndarray[Any, np.dtype[np.float64]] = named_array.real.data
    assert np.array_equal(actual_real, expected_real)
    assert actual_real.dtype == expected_real.dtype

    actual_imag: np.ndarray[Any, np.dtype[np.float64]] = named_array.imag.data
    assert np.array_equal(actual_imag, expected_imag)
    assert actual_imag.dtype == expected_imag.dtype

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thank you for the review... i'm going to address these in a follow up PR.



# Additional tests as per your original class-based code
@pytest.mark.parametrize(
"data, dtype",
Expand Down