Add support for Array API in NamedArray

xarray/core/common.py

@@ -198,6 +198,14 @@ def __iter__(self: Any) -> Iterator[Any]:
             raise TypeError("iteration over a 0-d array")
         return self._iter()
 
+    @overload
+    def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]:
+        ...
+
+    @overload
+    def get_axis_num(self, dim: Hashable) -> int:
+        ...
+
     def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, ...]:
         """Return axis number(s) corresponding to dimension(s) in this array.
 

xarray/namedarray/_array_api.py

@@ -8,13 +8,15 @@
 
 from xarray.namedarray._typing import (
     _arrayapi,
+    _AxisLike,
+    _Dims,
     _DType,
     _ScalarType,
     _ShapeType,
     _SupportsImag,
     _SupportsReal,
 )
-from xarray.namedarray.core import NamedArray
+from xarray.namedarray.core import NamedArray, _dims_to_axis, _get_remaining_dims
 
 with warnings.catch_warnings():
     warnings.filterwarnings(
@@ -144,3 +146,79 @@ def real(
     xp = _get_data_namespace(x)
     out = x._new(data=xp.real(x._data))
     return out
+
+
+# %% Statistical Functions
+
+
+def mean(
+    x: NamedArray[Any, _DType],
+    /,
+    *,
+    dims: _Dims | None = None,
+    keepdims: bool = False,
+    axis: _AxisLike | None = None,
+) -> NamedArray[Any, _DType]:
+    """
+    Calculates the arithmetic mean of the input array x.
+
+    Parameters
+    ----------
+    x :
+        Should have a real-valued floating-point data type.
+    dims :
+        Dim or dims along which arithmetic means must be computed. By default,
+        the mean must be computed over the entire array. If a tuple of hashables,
+        arithmetic means must be computed over multiple axes.
+        Default: None.
+    keepdims :
+        if True, the reduced axes (dimensions) must be included in the result
+        as singleton dimensions, and, accordingly, the result must be compatible
+        with the input array (see Broadcasting). Otherwise, if False, the
+        reduced axes (dimensions) must not be included in the result.
+        Default: False.
+    axis :
+        Axis or axes along which arithmetic means must be computed. By default,
+        the mean must be computed over the entire array. If a tuple of integers,
+        arithmetic means must be computed over multiple axes.
+        Default: None.
+
+    Returns
+    -------
+    out :
+        If the arithmetic mean was computed over the entire array,
+        a zero-dimensional array containing the arithmetic mean; otherwise,
+        a non-zero-dimensional array containing the arithmetic means.
+        The returned array must have the same data type as x.
+
+    Examples
+    --------
+    >>> x = NamedArray(("x", "y"), nxp.asarray([[1.0, 2.0], [3.0, 4.0]]))
+    >>> mean(x).data
+    Array(2.5, dtype=float64)
+    >>> mean(x, dims=("x",)).data
+    Array([2., 3.], dtype=float64)
+
+    Using keepdims:
+
+    >>> mean(x, dims=("x",), keepdims=True)
+    <xarray.NamedArray (x: 1, y: 2)>
+    Array([[2., 3.]], dtype=float64)
+    >>> mean(x, dims=("y",), keepdims=True)
+    <xarray.NamedArray (x: 2, y: 1)>
+    Array([[1.5],
+           [3.5]], dtype=float64)
+    """
+    xp = _get_data_namespace(x)
+    axis_ = _dims_to_axis(x, dims, axis)
+    d = xp.mean(x._data, axis=axis_, keepdims=False)  # We fix keepdims later
+    # TODO: Why do we need to do the keepdims ourselves?
+    dims_, data_ = _get_remaining_dims(x, d, axis_, keepdims=keepdims)
+    out = x._new(dims=dims_, data=data_)
+    return out
+
+
+# if __name__ == "__main__":
+#     import doctest
+
+#     doctest.testmod()

xarray/namedarray/_typing.py

@@ -17,14 +17,15 @@
 import numpy as np
 
 if TYPE_CHECKING:
-    from numpy.typing import NDArray
+    pass
 
 
 # https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
 
 
+_dtype = np.dtype
 _DType = TypeVar("_DType", bound=np.dtype[Any])
 _DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any])
 # A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic`
@@ -48,21 +49,32 @@ def dtype(self) -> _DType_co:
 ]
 
 # For unknown shapes Dask uses np.nan, array_api uses None:
-_IntOrUnknown = int
+_IntOrUnknown = int  # Union[int, _Unknown]
 _Shape = tuple[_IntOrUnknown, ...]
 _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]
 _ShapeType = TypeVar("_ShapeType", bound=Any)
 _ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True)
 
+_Axis = int
+_Axes = tuple[_Axis, ...]
+_AxisLike = Union[_Axis, _Axes]
+
 _Chunks = tuple[_Shape, ...]
 
 _Dim = Hashable
 _Dims = tuple[_Dim, ...]
-
 _DimsLike = Union[str, Iterable[_Dim]]
-_AttrsLike = Union[Mapping[Any, Any], None]
+_DimsLikeAgg = Union[_DimsLike, "ellipsis", None]
 
-_dtype = np.dtype
+
+# https://data-apis.org/array-api/latest/API_specification/indexing.html
+# TODO: np.array_api doesn't allow None for some reason, maybe they're
+# recommending to use expand_dims?
+_IndexKey = Union[int, slice, "ellipsis"]
+_IndexKeys = tuple[Union[_IndexKey, None], ...]
+_IndexKeyLike = Union[_IndexKey, _IndexKeys]
+
+_AttrsLike = Union[Mapping[Any, Any], None]
 
 
 class _SupportsReal(Protocol[_T_co]):
@@ -101,6 +113,11 @@ def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]:
     def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]:
         ...
 
+    def __array__(
+        self, dtype: _DType | None = ..., /
+    ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]:
+        ...
+
 
 # Corresponds to np.typing.NDArray:
 _Array = _array[Any, np.dtype[_ScalarType_co]]
@@ -116,6 +133,23 @@ class _arrayfunction(
     Corresponds to np.ndarray.
     """
 
+    @overload
+    def __getitem__(self, key: _IndexKeyLike) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...]
+    ) -> _arrayfunction[Any, _DType_co]:
+        ...
+
+    def __getitem__(
+        self,
+        key: _IndexKeyLike | _arrayfunction[Any, Any],
+        /,
+    ) -> _arrayfunction[Any, _DType_co] | Any:
+        ...
+
     # TODO: Should return the same subclass but with a new dtype generic.
     # https://github.com/python/typing/issues/548
     def __array_ufunc__(
@@ -151,6 +185,27 @@ class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType
     Corresponds to np.ndarray.
     """
 
+    # TODO: Only integer _arrayapi:
+    # def __getitem__(
+    #     self,
+    #     key: Union[
+    #         int,
+    #         slice,
+    #         "ellipsis",
+    #         tuple[Union[int, slice, "ellipsis", None], ...],
+    #         _arrayapi[Any, Any],
+    #     ],
+    #     /,
+    # ) -> _arrayapi[Any, _DType_co]:
+    #     ...
+
+    def __getitem__(
+        self,
+        key: Any,
+        /,
+    ) -> _arrayapi[Any, _DType_co]:
+        ...
+
     def __array_namespace__(self) -> ModuleType:
         ...
 
@@ -238,7 +293,7 @@ class _sparsearray(
     Corresponds to np.ndarray.
     """
 
-    def todense(self) -> NDArray[_ScalarType_co]:
+    def todense(self) -> np.ndarray[Any, _DType_co]:
         ...
 
 
@@ -256,7 +311,7 @@ class _sparsearrayfunction(
     Corresponds to np.ndarray.
     """
 
-    def todense(self) -> NDArray[_ScalarType_co]:
+    def todense(self) -> np.ndarray[Any, _DType_co]:
         ...
 
 
@@ -274,7 +329,7 @@ class _sparsearrayapi(
     Corresponds to np.ndarray.
     """
 
-    def todense(self) -> NDArray[_ScalarType_co]:
+    def todense(self) -> np.ndarray[Any, _DType_co]:
         ...