[namedarray] split .set_dims() into .expand_dims() and broadcast_to() (#8380)

* add `.set_dims()`, `.transpose()` and `.T` to namedarray

* more typying fixes

* more typing fixes

* override set_dims for IndexVariable

* fix dims

* split `.set_dims()` into `.expand_dims()` and `broadcast_to()`

* more typing fixes

* update whats-new

* update tests

* doc fixes

* update whats-new

* keep `.set_dims()` on `Variable()`

* update docs

* revert to set_dims

* revert to .set_dims on Variable

* Update xarray/namedarray/core.py

Co-authored-by: Illviljan <[email protected]>

* restore .transpose on variable

* revert to set_dims in Variable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix docstring

* update test_namedarray

* update tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestions from code review

Co-authored-by: Deepak Cherian <[email protected]>

* fix formatting issue

* fix tests

* update expand_dims

* update tests

* update tests

* remove unnecessary guard conditions

* Update type hints in NamedArray class and test cases

* Refactor NamedArray T property to handle non-2D arrays

* Reverse the order of dimensions in x.T

* Refactor broadcasting and dimension expansion in NamedArray

* update docstring

* add todo item

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use comprehension

* use dim

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix imports

* formatting only

* Apply suggestions from code review

Co-authored-by: Deepak Cherian <[email protected]>

* [skip-rtd] fix indentation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor expand_dims to simplify API

simplified the `expand_dims` method of the NamedArray class by removing support for multiple and keyword arguments for new dimensions. the change updates the method signature to only accept a single dimension name, enhancing clarity and maintainability. this shift focuses on the common use case of adding a single dimension and moves extended functionality to a separate API.

* fix type hint for `dim` parameter in test.

* fix typing issues

* fix UnboundLocalError: local variable 'flattened_dims' referenced before assignment

* fix type hint

* ignore typing

* update whats-new

* adjust the `broadcast_to` method to prohibit adding new dimensions, allowing only the broadcasting of existing ones

---------

Co-authored-by: Illviljan <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 4 people

doc/api-hidden.rst

@@ -351,33 +351,35 @@
    IndexVariable.values
 
 
-   namedarray.core.NamedArray.all
-   namedarray.core.NamedArray.any
-   namedarray.core.NamedArray.attrs
-   namedarray.core.NamedArray.chunks
-   namedarray.core.NamedArray.chunksizes
-   namedarray.core.NamedArray.copy
-   namedarray.core.NamedArray.count
-   namedarray.core.NamedArray.cumprod
-   namedarray.core.NamedArray.cumsum
-   namedarray.core.NamedArray.data
-   namedarray.core.NamedArray.dims
-   namedarray.core.NamedArray.dtype
-   namedarray.core.NamedArray.get_axis_num
-   namedarray.core.NamedArray.max
-   namedarray.core.NamedArray.mean
-   namedarray.core.NamedArray.median
-   namedarray.core.NamedArray.min
-   namedarray.core.NamedArray.nbytes
-   namedarray.core.NamedArray.ndim
-   namedarray.core.NamedArray.prod
-   namedarray.core.NamedArray.reduce
-   namedarray.core.NamedArray.shape
-   namedarray.core.NamedArray.size
-   namedarray.core.NamedArray.sizes
-   namedarray.core.NamedArray.std
-   namedarray.core.NamedArray.sum
-   namedarray.core.NamedArray.var
+   NamedArray.all
+   NamedArray.any
+   NamedArray.attrs
+   NamedArray.broadcast_to
+   NamedArray.chunks
+   NamedArray.chunksizes
+   NamedArray.copy
+   NamedArray.count
+   NamedArray.cumprod
+   NamedArray.cumsum
+   NamedArray.data
+   NamedArray.dims
+   NamedArray.dtype
+   NamedArray.expand_dims
+   NamedArray.get_axis_num
+   NamedArray.max
+   NamedArray.mean
+   NamedArray.median
+   NamedArray.min
+   NamedArray.nbytes
+   NamedArray.ndim
+   NamedArray.prod
+   NamedArray.reduce
+   NamedArray.shape
+   NamedArray.size
+   NamedArray.sizes
+   NamedArray.std
+   NamedArray.sum
+   NamedArray.var
 
 
    plot.plot

doc/whats-new.rst

@@ -23,6 +23,9 @@ v2024.02.0 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Add :py:meth:`NamedArray.expand_dims`, :py:meth:`NamedArray.permute_dims` and :py:meth:`NamedArray.broadcast_to`
+  (:pull:`8380`) By `Anderson Banihirwe <https://github.com/andersy005>`_.
+
 - Xarray now defers to flox's `heuristics <https://flox.readthedocs.io/en/latest/implementation.html#heuristics>`_
   to set default `method` for groupby problems. This only applies to ``flox>=0.9``.
   By `Deepak Cherian <https://github.com/dcherian>`_.

xarray/__init__.py

@@ -41,6 +41,7 @@
 from xarray.core.options import get_options, set_options
 from xarray.core.parallel import map_blocks
 from xarray.core.variable import IndexVariable, Variable, as_variable
+from xarray.namedarray.core import NamedArray
 from xarray.util.print_versions import show_versions
 
 try:
@@ -104,6 +105,7 @@
     "IndexSelResult",
     "IndexVariable",
     "Variable",
+    "NamedArray",
     # Exceptions
     "MergeError",
     "SerializationWarning",

xarray/core/dataarray.py

@@ -71,6 +71,7 @@
     as_compatible_data,
     as_variable,
 )
+from xarray.namedarray.utils import infix_dims
 from xarray.plot.accessor import DataArrayPlotAccessor
 from xarray.plot.utils import _get_units_from_attrs
 from xarray.util.deprecation_helpers import _deprecate_positional_args, deprecate_dims
@@ -3017,7 +3018,7 @@ def transpose(
         Dataset.transpose
         """
         if dims:
-            dims = tuple(utils.infix_dims(dims, self.dims, missing_dims))
+            dims = tuple(infix_dims(dims, self.dims, missing_dims))
         variable = self.variable.transpose(*dims)
         if transpose_coords:
             coords: dict[Hashable, Variable] = {}

xarray/core/dataset.py

@@ -114,8 +114,6 @@
     drop_dims_from_indexers,
     either_dict_or_kwargs,
     emit_user_level_warning,
-    infix_dims,
-    is_dict_like,
     is_scalar,
     maybe_wrap_array,
 )
@@ -126,6 +124,7 @@
     broadcast_variables,
     calculate_dimensions,
 )
+from xarray.namedarray.utils import infix_dims, is_dict_like
 from xarray.plot.accessor import DatasetPlotAccessor
 from xarray.util.deprecation_helpers import _deprecate_positional_args
 

xarray/core/groupby.py

@@ -5,14 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Hashable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, field
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Generic,
-    Literal,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union
 
 import numpy as np
 import pandas as pd

xarray/core/utils.py

@@ -846,36 +846,6 @@ def __len__(self) -> int:
         return len(self._data) - num_hidden
 
 
-def infix_dims(
-    dims_supplied: Collection,
-    dims_all: Collection,
-    missing_dims: ErrorOptionsWithWarn = "raise",
-) -> Iterator:
-    """
-    Resolves a supplied list containing an ellipsis representing other items, to
-    a generator with the 'realized' list of all items
-    """
-    if ... in dims_supplied:
-        if len(set(dims_all)) != len(dims_all):
-            raise ValueError("Cannot use ellipsis with repeated dims")
-        if list(dims_supplied).count(...) > 1:
-            raise ValueError("More than one ellipsis supplied")
-        other_dims = [d for d in dims_all if d not in dims_supplied]
-        existing_dims = drop_missing_dims(dims_supplied, dims_all, missing_dims)
-        for d in existing_dims:
-            if d is ...:
-                yield from other_dims
-            else:
-                yield d
-    else:
-        existing_dims = drop_missing_dims(dims_supplied, dims_all, missing_dims)
-        if set(existing_dims) ^ set(dims_all):
-            raise ValueError(
-                f"{dims_supplied} must be a permuted list of {dims_all}, unless `...` is included"
-            )
-        yield from existing_dims
-
-
 def get_temp_dimname(dims: Container[Hashable], new_dim: Hashable) -> Hashable:
     """Get an new dimension name based on new_dim, that is not used in dims.
     If the same name exists, we add an underscore(s) in the head.

xarray/core/variable.py

@@ -42,11 +42,11 @@
     drop_dims_from_indexers,
     either_dict_or_kwargs,
     ensure_us_time_resolution,
-    infix_dims,
     is_duck_array,
     maybe_coerce_to_str,
 )
 from xarray.namedarray.core import NamedArray, _raise_if_any_duplicate_dimensions
+from xarray.namedarray.utils import infix_dims
 
 NON_NUMPY_SUPPORTED_ARRAY_TYPES = (
     indexing.ExplicitlyIndexed,

xarray/namedarray/_array_api.py

@@ -9,6 +9,7 @@
 from xarray.namedarray._typing import (
     Default,
     _arrayapi,
+    _Axes,
     _Axis,
     _default,
     _Dim,
@@ -196,3 +197,32 @@ def expand_dims(
     d.insert(axis, dim)
     out = x._new(dims=tuple(d), data=xp.expand_dims(x._data, axis=axis))
     return out
+
+
+def permute_dims(x: NamedArray[Any, _DType], axes: _Axes) -> NamedArray[Any, _DType]:
+    """
+    Permutes the dimensions of an array.
+
+    Parameters
+    ----------
+    x :
+        Array to permute.
+    axes :
+        Permutation of the dimensions of x.
+
+    Returns
+    -------
+    out :
+        An array with permuted dimensions. The returned array must have the same
+        data type as x.
+
+    """
+
+    dims = x.dims
+    new_dims = tuple(dims[i] for i in axes)
+    if isinstance(x._data, _arrayapi):
+        xp = _get_data_namespace(x)
+        out = x._new(dims=new_dims, data=xp.permute_dims(x._data, axes))
+    else:
+        out = x._new(dims=new_dims, data=x._data.transpose(axes))  # type: ignore[attr-defined]
+    return out

xarray/namedarray/_typing.py

@@ -7,6 +7,7 @@
     Any,
     Callable,
     Final,
+    Literal,
     Protocol,
     SupportsIndex,
     TypeVar,
@@ -311,8 +312,10 @@ def todense(self) -> np.ndarray[Any, _DType_co]:
 
 # NamedArray can most likely use both __array_function__ and __array_namespace__:
 _sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi)
-
 sparseduckarray = Union[
     _sparsearrayfunction[_ShapeType_co, _DType_co],
     _sparsearrayapi[_ShapeType_co, _DType_co],
 ]
+
+ErrorOptions = Literal["raise", "ignore"]
+ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]