fix: groupby-sum fast path reduces multi-dim DataArray grouper over all its dims

doc/release_notes.rst

@@ -38,6 +38,7 @@ Upcoming Version
 * LP file export now honors bounds tightened below ``[0, 1]`` on a binary variable via the ``.lower``/``.upper`` setters after creation (e.g. ``upper = 0``). Previously such bounds were written only by ``io_api="direct"`` and dropped by ``io_api="lp"``. (https://github.com/PyPSA/linopy/issues/776)
 * Freezing an empty constraint group (e.g. an empty ``isel`` slice) no longer raises ``ValueError: cannot reshape array of size 0``. ``Model(freeze_constraints=True)`` and ``Constraint.freeze()`` now round-trip zero-row constraints losslessly.
 * ``Variable.where`` no longer raises ``ValueError: exact match required for all data variable names`` once a solution is attached (after ``Model.solve``) or the variable is fixed. The fill value now covers auxiliary data variables (``solution``, stashed bounds) instead of only ``labels``/``lower``/``upper``.
+* ``LinearExpression.groupby(...).sum()`` with a multi-dimensional ``DataArray`` grouper now reduces over all of the grouper's dimensions on the default (fast) path, instead of leaking one of them into the result.
 * ``linopy.testing.assert_linequal`` now aligns dimension order before comparing, so mathematically identical expressions built in different orders (e.g. ``x + y`` versus ``y + x``, which inherit different dimension orders from xarray broadcasting) are correctly treated as equal. Genuinely different expressions still fail.
 
 Version 0.8.0

linopy/constants.py

@@ -76,6 +76,7 @@ class PerformanceWarning(UserWarning):
 LP_PIECE_DIM = f"{BREAKPOINT_DIM}_piece"
 PWL_LINK_DIM = "_pwl_var"
 GROUP_DIM = "_group"
+GROUP_STACK_DIM = "_group_stack"
 FACTOR_DIM = "_factor"
 CONCAT_DIM = "_concat"
 CV_DIM = "_cv"

linopy/expressions.py

@@ -74,6 +74,7 @@
     FACTOR_DIM,
     GREATER_EQUAL,
     GROUP_DIM,
+    GROUP_STACK_DIM,
     HELPER_DIMS,
     LESS_EQUAL,
     STACKED_TERM_DIM,
@@ -173,6 +174,24 @@ def _multikey_value_frame(group: Any, data: Dataset) -> pd.DataFrame | None:
     return data[names].to_dataframe()[names]
 
 
+def _flatten_multidim_group(
+    group: DataArray, data: Dataset
+) -> tuple[pd.Series, Dataset]:
+    """
+    Flatten an N-D DataArray grouper into a 1-D key over a stacked dimension.
+
+    A multi-dimensional grouper defines its groups jointly over all its dims, so
+    the data is stacked over those dims and the grouper's values are raveled into
+    a Series aligned with the stacked axis. This lets the fast path reduce over
+    every grouper dim at once instead of leaking the extra dims into the result.
+    """
+    group_dims = list(group.dims)
+    data = data.stack({GROUP_STACK_DIM: group_dims}, create_index=False)
+    series = pd.Series(group.transpose(*group_dims).values.reshape(-1), name=group.name)
+    series.index.name = GROUP_STACK_DIM
+    return series, data
+
+
 def _unstack_multikey(ds: Dataset, dim: str) -> Dataset:
     """
     Unstack a stacked multi-key group dimension into one dimension per key.
@@ -340,6 +359,15 @@ def sum(
         if multikey_frame is not None:
             group = multikey_frame
 
+        data = self.data
+        is_multidim_grouper = (
+            isinstance(group, DataArray)
+            and group.ndim > 1
+            and set(group.dims) <= set(data.dims)
+        )
+        if is_multidim_grouper and not use_fallback:
+            group, data = _flatten_multidim_group(group, data)
+
         fast_path_types = (pd.Series, pd.DataFrame, xr.DataArray)
         if isinstance(group, fast_path_types) and not use_fallback:
             final_group_name = (
@@ -358,7 +386,7 @@ def sum(
                 )
 
             assert isinstance(group, pd.Series)
-            ds = self._grouped_sum(group)
+            ds = self._grouped_sum(group, data)
 
             if multikey_decode is not None:
                 ds = _restore_multikey_index(ds, multikey_decode)
@@ -375,7 +403,7 @@ def func(ds: Dataset) -> Dataset:
 
         return self.map(func)
 
-    def _grouped_sum(self, group: pd.Series) -> Dataset:
+    def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset:
         """
         Sum groups by scattering all terms directly into the final padded arrays.
 
@@ -384,7 +412,7 @@ def _grouped_sum(self, group: pd.Series) -> Dataset:
         padded with fill values. Only the result arrays are allocated, keeping
         peak memory at input + result.
         """
-        data = self.data
+        data = self.data if data is None else data
         group_dim = group.index.name
         fill_value = LinearExpression._fill_value
 

test/test_linear_expression.py

@@ -1668,16 +1668,31 @@ def test_multiindex_level(
         assert grouped.vars.transpose(level, TERM_DIM).values.tolist() == vars_
 
 
-@pytest.mark.parametrize("use_fallback", [True])
+@pytest.mark.parametrize("use_fallback", [True, False])
 def test_linear_expression_groupby_ndim(z: Variable, use_fallback: bool) -> None:
-    # TODO: implement fallback for n-dim groupby, see https://github.com/PyPSA/linopy/issues/299
     expr = 1 * z
     groups = xr.DataArray([[1, 1, 2], [1, 3, 3]], coords=z.coords)
     grouped = expr.groupby(groups).sum(use_fallback=use_fallback)
     assert "group" in grouped.dims
     # there are three groups, 1, 2 and 3, the largest group has 3 elements
     assert (grouped.data.group == [1, 2, 3]).all()
     assert grouped.nterm == 3
+    assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
+
+
+def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None:
+    # a 2-D grouper reduces over both its dims jointly, leaving unrelated dims
+    # untouched, see https://github.com/PyPSA/linopy/issues/823
+    m = Model()
+    v = m.add_variables(coords=[[0, 1], [0, 1], [0, 1]], dims=["i", "j", "k"])
+    groups = xr.DataArray(
+        [[1, 1], [2, 2]], dims=["i", "j"], coords={"i": [0, 1], "j": [0, 1]}, name="g"
+    )
+    expr = 1 * v
+    grouped = expr.groupby(groups).sum()
+    assert set(grouped.dims) == {"g", "k", "_term"}
+    assert (grouped.data.g == [1, 2]).all()
+    assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
 
 
 @pytest.mark.parametrize("use_fallback", [True, False])