fix NamedArray.imag and NamedArray.real typing info

[andersy005]

• add .imag and .real properties to NamedArray #8365 follow up
• Tests added

[Illviljan]

xarray/namedarray/core.py: note: In member "imag" of class "NamedArray":
xarray/namedarray/core.py:528: error: Item "_arrayfunction[Any, dtype[_SupportsImag[_ScalarType]]]" of "_arrayfunction[Any, dtype[_SupportsImag[_ScalarType]]] | _arrayapi[Any, dtype[_SupportsImag[_ScalarType]]]" has no attribute "imag"  [union-attr]
xarray/namedarray/core.py:528: error: Item "_arrayapi[Any, dtype[_SupportsImag[_ScalarType]]]" of "_arrayfunction[Any, dtype[_SupportsImag[_ScalarType]]] | _arrayapi[Any, dtype[_SupportsImag[_ScalarType]]]" has no attribute "imag"  [union-attr]
xarray/namedarray/core.py: note: In member "real" of class "NamedArray":
xarray/namedarray/core.py:541: error: Item "_arrayfunction[Any, dtype[_SupportsReal[_ScalarType]]]" of "_arrayfunction[Any, dtype[_SupportsReal[_ScalarType]]] | _arrayapi[Any, dtype[_SupportsReal[_ScalarType]]]" has no attribute "real"  [union-attr]
xarray/namedarray/core.py:541: error: Item "_arrayapi[Any, dtype[_SupportsReal[_ScalarType]]]" of "_arrayfunction[Any, dtype[_SupportsReal[_ScalarType]]] | _arrayapi[Any, dtype[_SupportsReal[_ScalarType]]]" has no attribute "real"  [union-attr]

Because the duckarray protocols are missing .imag and .real. You can add them to _arrayfunction, but not _arrayapi. because the array api standard doesn't include them as methods. I wonder if it's better to use the array_api functions?

xarray/core/dataarray.py: note: In member "real" of class "DataArray":
xarray/core/dataarray.py:4901: error: Argument 1 to "_replace" of "DataArray" has incompatible type "NamedArray[Any, dtype[Any]]"; expected "Variable | None"  [arg-type]
xarray/core/dataarray.py: note: In member "imag" of class "DataArray":
xarray/core/dataarray.py:4912: error: Argument 1 to "_replace" of "DataArray" has incompatible type "NamedArray[Any, dtype[Any]]"; expected "Variable | None"  [arg-type]

Because imag and real only return NamedArray now, rather than a wrong subclass of it. Add imag and real to xr.Variable with similar typing as NamedArray.
Any method using _new directly or indirectly will likely need this.

xarray/tests/test_namedarray.py: note: In function "test_real_and_imag":
xarray/tests/test_namedarray.py:183: error: Incompatible types in assignment (expression has type "_arrayfunction[Any, dtype[floating[_64Bit]]] | _arrayapi[Any, dtype[floating[_64Bit]]]", variable has type "ndarray[Any, dtype[floating[_64Bit]]]")  [assignment]
xarray/tests/test_namedarray.py:187: error: Incompatible types in assignment (expression has type "_arrayfunction[Any, dtype[floating[_64Bit]]] | _arrayapi[Any, dtype[floating[_64Bit]]]", variable has type "ndarray[Any, dtype[floating[_64Bit]]]")  [assignment]

Not sure about this one, maybe will self heal once the other issues are fixed.

[andersy005]

I wonder if it's better to use the array_api functions?

last night i tried, and it appears it won't work until there's wide support of the array API in upstream dependencies (e.g. pint).

In [14]: import pint, xarray as xr, numpy as np

In [15]: unit_registry = pint.UnitRegistry(force_ndarray_like=False)

In [16]: Quantity = unit_registry.Quantity

In [17]: dtype = float

In [18]: array = (
    ...:             np.arange(5 * 10).astype(dtype)
    ...:             + 1j * np.linspace(-1, 0, 5 * 10).astype(dtype)
    ...:         ).reshape(5, 10) * unit_registry.s

In [23]: np.real(array)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[23], line 1
----> 1 np.real(array)

File <__array_function__ internals>:200, in real(*args, **kwargs)

TypeError: no implementation found for 'numpy.real' on types that implement __array_function__: [<class 'pint.Quantity'>]

curious to hear if there is a consistent way of using the array api functions as they are today without needing to special case upstream libraries with limited or without array api compliance.

[Illviljan]

Hmm, I feel like pint is just lagging behind here. np.real and np.imag has been around for a long time, dask.array handles it just fine.

I like the idea of keeping the _array_api modules clean of strange workarounds, so I think you should do the workaround in the NamedArray methods with maybe a check like this:

    @property
    def imag(
        self: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]],  # type: ignore[type-var]
    ) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:
        if isinstance(self._data, _arrayapi):
            from xarray.namedarray._array_api import real

            return real(self)
        else:
            return self._new(data=self.data.imag) # _arrayfunction might need .imag though

I figured these out:

xarray/tests/test_namedarray.py: note: In function "test_real_and_imag":
xarray/tests/test_namedarray.py:183: error: Incompatible types in assignment (expression has type "_arrayfunction[Any, dtype[floating[_64Bit]]] | _arrayapi[Any, dtype[floating[_64Bit]]]", variable has type "ndarray[Any, dtype[floating[_64Bit]]]")  [assignment]
xarray/tests/test_namedarray.py:187: error: Incompatible types in assignment (expression has type "_arrayfunction[Any, dtype[floating[_64Bit]]] | _arrayapi[Any, dtype[floating[_64Bit]]]", variable has type "ndarray[Any, dtype[floating[_64Bit]]]")  [assignment]

NamedArray.data returns the wide duckarray and we try to assign it to a more narrow version:

a: wide_type = ...
b: narrow_type = a  # Not OK

# So this is not ok because np.array is narrower type than duckarray:
actual_imag: np.ndarray[Any, np.dtype[np.float64]] = named_array.imag.data

# But this is ok and is still relevant for our dtype test:
actual_real: duckarray[Any, np.dtype[np.float64]] = named_array.real.data

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: 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

    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

[andersy005]

thank you for all your help, @Illviljan! i believe this is ready for a last round of reviews.

[Illviljan]

Well done, looking good to me!

[Illviljan]

Not sure why the numpy strategy of replacing self with a new type doesn't work here.
Oh well, NamedArray has the correct dtype so it's good enough for me.