From 9ca701f0707fe0b18453c02d95212a6e6b4d7f29 Mon Sep 17 00:00:00 2001
From: Fabian <fab.hof@gmx.de>
Date: Tue, 21 Jul 2026 11:31:33 +0200
Subject: [PATCH 1/9] test: pin groupby-sum to preserve grouped-dim position
 (xfail)

groupby-sum moves the group dim to the trailing position instead of
keeping the grouped dim's slot, diverging from xarray's native
groupby-reduce and linopy <= 0.8.0 (regressed in #802). Existing tests
use assert_linequal, which aligns dims and misses the reorder.

Strict xfail: remove the marker once grouping preserves dim position.
---
 test/test_linear_expression.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py
index 92e2510f9..42c348015 100644
--- a/test/test_linear_expression.py
+++ b/test/test_linear_expression.py
@@ -1695,6 +1695,39 @@ def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None:
     assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
 
 
+@pytest.mark.xfail(
+    strict=True,
+    reason="groupby-sum moves the group dim to the trailing position instead of "
+    "keeping the grouped dim's slot, diverging from xarray's native groupby-reduce "
+    "and linopy <= 0.8.0. Regressed in #802 (apply_ufunc scatter kernel); "
+    "assert_linequal aligns dims so existing tests miss it.",
+)
+@pytest.mark.parametrize("use_fallback", [False, True])
+def test_linear_expression_groupby_preserves_dim_position(use_fallback: bool) -> None:
+    # Grouping must replace the grouped dim in place, like xarray. `name` sits
+    # between `scenario` and `snapshot`, so the group dim must stay in the middle.
+    m = Model()
+    coords = [
+        pd.Index(["s1", "s2"], name="scenario"),
+        pd.Index(["a", "b", "c"], name="name"),
+        pd.RangeIndex(2, name="snapshot"),
+    ]
+    v = m.add_variables(coords=coords, name="x")
+    groups = xr.DataArray(
+        ["g1", "g1", "g2"], coords={"name": ["a", "b", "c"]}, name="bus"
+    )
+
+    reference = (
+        xr.DataArray(np.zeros((2, 3, 2)), coords=coords, dims=[c.name for c in coords])
+        .groupby(groups)
+        .sum()
+    )
+    grouped = (1 * v).groupby(groups).sum(use_fallback=use_fallback)
+
+    result_dims = tuple(d for d in grouped.dims if d != TERM_DIM)
+    assert result_dims == reference.dims
+
+
 class TestGroupbyGrouperAlignment:
     """
     A ``pd.Series``/``DataArray`` grouper whose labels are reordered or a

From 43829158185570ac196eba31c7268603e8145445 Mon Sep 17 00:00:00 2001
From: Fabian <fab.hof@gmx.de>
Date: Tue, 21 Jul 2026 12:12:17 +0200
Subject: [PATCH 2/9] fix(groupby): preserve grouped-dim position in
 groupby-sum

groupby-sum moved the group dim to the trailing position; both the
apply_ufunc scatter kernel and the fallback append it there, and the
LinearExpression constructor's broadcast then locks that order.

Restore the grouped dim's original slot via _restore_group_dim_position:
transpose the terms and reorder the dimension coordinates so the
constructor's canonical dim order matches. MultiIndex group coords are
left untouched (multikey/observed results stay order-agnostic).

Drops the strict xfail; adds a release note.
---
 doc/release_notes.rst          |  1 +
 linopy/expressions.py          | 47 +++++++++++++++++++++++++++++-----
 test/test_linear_expression.py |  7 -----
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index 8ba3a0ee6..c7d3f5c7e 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -39,6 +39,7 @@ Upcoming Version
 * 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.
+* ``LinearExpression.groupby(...).sum()`` now keeps the grouped dimension's position, replacing it in place like xarray's native groupby-reduce, instead of moving the group dimension to the trailing position (regressed in 0.8.0). Both the default and ``use_fallback=True`` paths are fixed.
 * ``LinearExpression.groupby(...).sum()`` now raises when a grouper's labels are reordered or a different set relative to the expression, instead of silently regrouping by position. Reorder the grouper to match the expression's coordinates before grouping. (https://github.com/PyPSA/linopy/issues/827)
 * ``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.
 * ``Solver.close()`` (also triggered by ``model.solver = None`` and the next ``solve()`` call) now explicitly disposes the ``gurobipy`` model before the environment. Previously the model was only dereferenced, so a user-held ``model.solver_model`` reference silently kept the Gurobi license acquired after ``close()``. (https://github.com/PyPSA/linopy/issues/459)
diff --git a/linopy/expressions.py b/linopy/expressions.py
index f13fc6f26..07e4ea686 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -192,6 +192,37 @@ def _flatten_multidim_group(
     return series, data
 
 
+def _restore_group_dim_position(result: Dataset, original_dims: tuple[str, ...]) -> Dataset:
+    """
+    Move the new group dimension into the slot the grouped dimension occupied.
+
+    xarray's groupby-reduce and linopy <= 0.8.0 replace the grouped dimension in
+    place, but both the scatter kernel and the fallback append the group dim at
+    the trailing position. Surviving dims keep their order; the group dim(s) are
+    reinserted where the consumed dim(s) used to be.
+    """
+    consumed = [d for d in original_dims if d not in result.dims]
+    if not consumed:
+        return result
+    new_dims = [str(d) for d in result.dims if d not in original_dims]
+    order: list[str] = []
+    inserted = False
+    for d in original_dims:
+        if d in consumed:
+            if not inserted:
+                order.extend(new_dims)
+                inserted = True
+        else:
+            order.append(d)
+    result = result.transpose(*order, missing_dims="ignore")
+    reordered_coords = {
+        d: result[d]
+        for d in order
+        if d in result.coords and not isinstance(result.get_index(d), pd.MultiIndex)
+    }
+    return result.assign_coords(reordered_coords)
+
+
 def _unstack_multikey(ds: Dataset, dim: str) -> Dataset:
     """
     Unstack a stacked multi-key group dimension into one dimension per key.
@@ -397,6 +428,7 @@ def sum(
             group = multikey_frame
 
         data = self.data
+        original_dims = self.data.coeffs.dims
         is_multidim_grouper = (
             isinstance(group, DataArray)
             and group.ndim > 1
@@ -429,14 +461,17 @@ def sum(
             ds = ds.rename({GROUP_DIM: final_group_name})
             if multikey_frame is not None and not observed:
                 ds = _unstack_multikey(ds, final_group_name)
-            return LinearExpression(ds, self.model)
+        else:
 
-        def func(ds: Dataset) -> Dataset:
-            ds = LinearExpression._sum(ds, str(self.groupby._group_dim))
-            ds = ds.assign_coords({TERM_DIM: np.arange(len(ds._term))})
-            return ds
+            def func(ds: Dataset) -> Dataset:
+                ds = LinearExpression._sum(ds, str(self.groupby._group_dim))
+                ds = ds.assign_coords({TERM_DIM: np.arange(len(ds._term))})
+                return ds
 
-        return self.map(func)
+            ds = self.groupby.map(func)
+
+        ds = _restore_group_dim_position(ds, original_dims)
+        return LinearExpression(ds, self.model)
 
     def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset:
         """
diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py
index 42c348015..f16516149 100644
--- a/test/test_linear_expression.py
+++ b/test/test_linear_expression.py
@@ -1695,13 +1695,6 @@ def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None:
     assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
 
 
-@pytest.mark.xfail(
-    strict=True,
-    reason="groupby-sum moves the group dim to the trailing position instead of "
-    "keeping the grouped dim's slot, diverging from xarray's native groupby-reduce "
-    "and linopy <= 0.8.0. Regressed in #802 (apply_ufunc scatter kernel); "
-    "assert_linequal aligns dims so existing tests miss it.",
-)
 @pytest.mark.parametrize("use_fallback", [False, True])
 def test_linear_expression_groupby_preserves_dim_position(use_fallback: bool) -> None:
     # Grouping must replace the grouped dim in place, like xarray. `name` sits

From 16661a61fef2bd3b3971be63d54684efc7ec7f4f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 21 Jul 2026 10:12:50 +0000
Subject: [PATCH 3/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 linopy/expressions.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index 07e4ea686..13108cdbc 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -192,7 +192,9 @@ def _flatten_multidim_group(
     return series, data
 
 
-def _restore_group_dim_position(result: Dataset, original_dims: tuple[str, ...]) -> Dataset:
+def _restore_group_dim_position(
+    result: Dataset, original_dims: tuple[str, ...]
+) -> Dataset:
     """
     Move the new group dimension into the slot the grouped dimension occupied.
 

From fd41d26d8b163689775eb23616664fab54d46a36 Mon Sep 17 00:00:00 2001
From: Fabian <fab.hof@gmx.de>
Date: Tue, 21 Jul 2026 12:17:07 +0200
Subject: [PATCH 4/9] refactor(groupby): fold insertion flag into new_dims
 sentinel

---
 linopy/expressions.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index 13108cdbc..5a2a673c7 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -208,14 +208,12 @@ def _restore_group_dim_position(
         return result
     new_dims = [str(d) for d in result.dims if d not in original_dims]
     order: list[str] = []
-    inserted = False
     for d in original_dims:
-        if d in consumed:
-            if not inserted:
-                order.extend(new_dims)
-                inserted = True
-        else:
+        if d not in consumed:
             order.append(d)
+        elif new_dims:
+            order.extend(new_dims)
+            new_dims = []
     result = result.transpose(*order, missing_dims="ignore")
     reordered_coords = {
         d: result[d]

From 649d86e4874cef19205d30e50eea784666b9add1 Mon Sep 17 00:00:00 2001
From: Fabian <fab.hof@gmx.de>
Date: Tue, 21 Jul 2026 12:35:30 +0200
Subject: [PATCH 5/9] perf(groupby): make reordered terms C-contiguous to keep
 export cheap

The dim-position transpose left coeffs/vars as non-contiguous views, so
LP/solver export had to materialise a contiguous copy on ravel (CodSpeed
flagged ~4x export memory on the nodal_balance padding benchmark).

Force numpy-backed terms C-contiguous after the reorder so downstream
flattening ravels a view. ascontiguousarray is a no-op when no reorder
moved data; dask arrays are left lazy (chunks preserved).
---
 linopy/expressions.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index 5a2a673c7..d1ca4afd3 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -202,6 +202,10 @@ def _restore_group_dim_position(
     place, but both the scatter kernel and the fallback append the group dim at
     the trailing position. Surviving dims keep their order; the group dim(s) are
     reinserted where the consumed dim(s) used to be.
+
+    The transposed terms are a non-contiguous view; numpy-backed arrays are made
+    C-contiguous so downstream flattening (LP/solver export) ravels a view
+    instead of copying. Dask arrays are left lazy.
     """
     consumed = [d for d in original_dims if d not in result.dims]
     if not consumed:
@@ -220,7 +224,12 @@ def _restore_group_dim_position(
         for d in order
         if d in result.coords and not isinstance(result.get_index(d), pd.MultiIndex)
     }
-    return result.assign_coords(reordered_coords)
+    result = result.assign_coords(reordered_coords)
+    for name in ("coeffs", "vars"):
+        term = result[name]
+        if term.chunks is None:
+            result[name] = (term.dims, np.ascontiguousarray(term.values))
+    return result
 
 
 def _unstack_multikey(ds: Dataset, dim: str) -> Dataset:

From cc7b7f1717e86b3b89c90aac83e7139871b9b68b Mon Sep 17 00:00:00 2001
From: Fabian <fab.hof@gmx.de>
Date: Tue, 21 Jul 2026 12:53:46 +0200
Subject: [PATCH 6/9] perf(groupby): scatter into target axis order in one
 contiguous alloc

Make _grouped_sum emit the grouped dimension in place directly: pass all
surviving dims as core dims so apply_ufunc's output order is fully chosen,
and build the padded array with the group axis in its target slot. The
result is contiguous in a single allocation, so the downstream reorder and
contiguity coercion become no-ops (no copy).

Removes the build-memory regression from the previous force-contiguous
approach: build and export peak now match master (~10MB / ~3MB) on
nodal_balance instead of 19MB / 3MB.

Trade-off: every dim is now a core dim, so dask inputs collapse to a
single chunk (no per-surviving-dim parallelism). Kept lazy.
---
 doc/release_notes.rst |   2 +-
 linopy/expressions.py | 113 +++++++++++++++++++++++++-----------------
 2 files changed, 69 insertions(+), 46 deletions(-)

diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index c7d3f5c7e..29d75df35 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -25,7 +25,7 @@ Upcoming Version
 
 **Performance**
 
-* ``LinearExpression.groupby(...).sum()`` now scatters terms directly into the padded result arrays via ``xarray.apply_ufunc``, avoiding intermediate copies and speeding up the grouping. A single kernel covers both numpy and chunked (dask) data, the latter staying lazy. On representative models this lowers build and export peak memory by up to ~3x.
+* ``LinearExpression.groupby(...).sum()`` now scatters terms directly into the padded result arrays via ``xarray.apply_ufunc``, avoiding intermediate copies and speeding up the grouping. A single kernel covers both numpy and chunked (dask) data, the latter staying lazy. On representative models this lowers build and export peak memory by up to ~3x. The kernel emits the grouped result in its final axis order in one contiguous allocation; on dask inputs the reduction now runs over a single chunk (it no longer parallelises over the surviving dimensions).
 
 **Deprecations**
 
diff --git a/linopy/expressions.py b/linopy/expressions.py
index d1ca4afd3..17f600865 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -196,16 +196,19 @@ def _restore_group_dim_position(
     result: Dataset, original_dims: tuple[str, ...]
 ) -> Dataset:
     """
-    Move the new group dimension into the slot the grouped dimension occupied.
-
-    xarray's groupby-reduce and linopy <= 0.8.0 replace the grouped dimension in
-    place, but both the scatter kernel and the fallback append the group dim at
-    the trailing position. Surviving dims keep their order; the group dim(s) are
-    reinserted where the consumed dim(s) used to be.
-
-    The transposed terms are a non-contiguous view; numpy-backed arrays are made
-    C-contiguous so downstream flattening (LP/solver export) ravels a view
-    instead of copying. Dask arrays are left lazy.
+    Move the new group dimension into the slot the grouped dimension occupied,
+    matching xarray's groupby-reduce and linopy <= 0.8.0.
+
+    Surviving dims keep their order; the group dim(s) are reinserted where the
+    consumed dim(s) used to be. The coordinate variables are reordered too so the
+    dataset's canonical dim order (which the constructor's ``broadcast`` re-derives
+    from coordinate-insertion order) matches, instead of trailing the group dim.
+    This covers the fallback, which appends the group dim, and normalises the
+    scatter kernel's target-ordered output.
+
+    Numpy-backed terms are made C-contiguous so downstream flattening (LP/solver
+    export) ravels a view instead of copying; a no-op when already contiguous, as
+    the scatter kernel's output is. Dask arrays are left lazy.
     """
     consumed = [d for d in original_dims if d not in result.dims]
     if not consumed:
@@ -488,16 +491,18 @@ def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset
 
         Every group member keeps its block of ``nterm`` terms, so the resulting
         term dimension has size ``max_group_size * nterm`` and smaller groups are
-        padded with fill values. Only the result arrays are allocated, keeping
-        peak memory at input + result.
+        padded with fill values. The group dimension replaces the grouped one in
+        place: all surviving dimensions are passed as core dims, so ``apply_ufunc``
+        emits the target axis order and the padded arrays are allocated once,
+        contiguous. Dask arrays are collapsed to a single chunk (no per-dim
+        parallelism) since every dimension is now a core dim.
         """
         data = self.data if data is None else data
         group_dim = group.index.name
         fill_value = LinearExpression._fill_value
 
-        _scatter_core_dims = {group_dim: -1, TERM_DIM: -1}
         if data.chunks:
-            data = data.chunk(_scatter_core_dims)
+            data = data.chunk(-1)
 
         codes, unique_groups = pd.factorize(group, sort=True)
         if (codes == -1).any():
@@ -511,30 +516,34 @@ def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset
         position_in_group = pd.Series(codes).groupby(codes).cumcount().to_numpy()
         nterm = data.sizes[TERM_DIM]
 
-        def scatter_terms(values: np.ndarray, fill: Any) -> np.ndarray:
-            """Place each member's term block into its group's padded slot."""
-            leading = values.shape[:-2]
-            blocks = np.full(
-                (*leading, n_groups, nterm, max_size), fill, dtype=values.dtype
-            )
-            blocks[..., codes, :, position_in_group] = np.moveaxis(values, -2, 0)
-            return blocks.reshape((*leading, n_groups, nterm * max_size))
-
-        def group_sum(values: np.ndarray) -> np.ndarray:
-            """Sum values within each group along the last axis, skipping NaNs."""
-            by_element = np.moveaxis(values, -1, 0)
-            summed = np.zeros((n_groups, *by_element.shape[1:]), dtype=values.dtype)
-            np.add.at(summed, codes, np.where(np.isnan(by_element), 0, by_element))
-            return np.moveaxis(summed, 0, -1)
-
         def scatter(da: DataArray, fill: Any) -> DataArray:
-            """Run the term scatter over the core dims, lazily for dask arrays."""
+            """Scatter each member's terms into its group's padded slot, keeping
+            the grouped dimension's position."""
+            da = da.transpose(..., TERM_DIM)
+            dims_in = list(da.dims)
+            group_axis = dims_in.index(group_dim)
+            target = [GROUP_DIM if d == group_dim else d for d in dims_in]
+
+            def scatter_terms(values: np.ndarray) -> np.ndarray:
+                surv_shape = [
+                    n_groups if d == group_dim else s
+                    for d, s in zip(dims_in, values.shape)
+                    if d != TERM_DIM
+                ]
+                out = np.full(
+                    (*surv_shape, nterm, max_size), fill, dtype=values.dtype
+                )
+                grouped = np.moveaxis(out, group_axis, 0)
+                grouped[codes, ..., :, position_in_group] = np.moveaxis(
+                    values, group_axis, 0
+                )
+                return out.reshape((*surv_shape, nterm * max_size))
+
             return xr.apply_ufunc(
                 scatter_terms,
                 da,
-                kwargs={"fill": fill},
-                input_core_dims=[[group_dim, TERM_DIM]],
-                output_core_dims=[[GROUP_DIM, TERM_DIM]],
+                input_core_dims=[dims_in],
+                output_core_dims=[target],
                 exclude_dims={group_dim, TERM_DIM},
                 dask="parallelized",
                 dask_gufunc_kwargs={
@@ -543,21 +552,35 @@ def scatter(da: DataArray, fill: Any) -> DataArray:
                 output_dtypes=[da.dtype],
             )
 
-        const = xr.apply_ufunc(
-            group_sum,
-            data.const,
-            input_core_dims=[[group_dim]],
-            output_core_dims=[[GROUP_DIM]],
-            exclude_dims={group_dim},
-            dask="parallelized",
-            dask_gufunc_kwargs={"output_sizes": {GROUP_DIM: n_groups}},
-            output_dtypes=[data.const.dtype],
-        )
+        def group_sum(da: DataArray) -> DataArray:
+            """Sum the constant term within each group, keeping the grouped
+            dimension's position."""
+            dims_in = list(da.dims)
+            group_axis = dims_in.index(group_dim)
+            target = [GROUP_DIM if d == group_dim else d for d in dims_in]
+
+            def reduce(values: np.ndarray) -> np.ndarray:
+                members = np.moveaxis(values, group_axis, 0)
+                summed = np.zeros((n_groups, *members.shape[1:]), dtype=values.dtype)
+                np.add.at(summed, codes, np.where(np.isnan(members), 0, members))
+                return np.moveaxis(summed, 0, group_axis)
+
+            return xr.apply_ufunc(
+                reduce,
+                da,
+                input_core_dims=[dims_in],
+                output_core_dims=[target],
+                exclude_dims={group_dim},
+                dask="parallelized",
+                dask_gufunc_kwargs={"output_sizes": {GROUP_DIM: n_groups}},
+                output_dtypes=[da.dtype],
+            )
+
         ds = Dataset(
             {
                 "coeffs": scatter(data.coeffs, fill_value["coeffs"]),
                 "vars": scatter(data.vars, fill_value["vars"]),
-                "const": const,
+                "const": group_sum(data.const),
             }
         )
         return ds.assign_coords({GROUP_DIM: unique_groups})

From 6f0c043751875e7ae870b9e7eab9e369ee45adb6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 21 Jul 2026 10:53:57 +0000
Subject: [PATCH 7/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 linopy/expressions.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index 17f600865..539becdc9 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -517,8 +517,10 @@ def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset
         nterm = data.sizes[TERM_DIM]
 
         def scatter(da: DataArray, fill: Any) -> DataArray:
-            """Scatter each member's terms into its group's padded slot, keeping
-            the grouped dimension's position."""
+            """
+            Scatter each member's terms into its group's padded slot, keeping
+            the grouped dimension's position.
+            """
             da = da.transpose(..., TERM_DIM)
             dims_in = list(da.dims)
             group_axis = dims_in.index(group_dim)
@@ -530,9 +532,7 @@ def scatter_terms(values: np.ndarray) -> np.ndarray:
                     for d, s in zip(dims_in, values.shape)
                     if d != TERM_DIM
                 ]
-                out = np.full(
-                    (*surv_shape, nterm, max_size), fill, dtype=values.dtype
-                )
+                out = np.full((*surv_shape, nterm, max_size), fill, dtype=values.dtype)
                 grouped = np.moveaxis(out, group_axis, 0)
                 grouped[codes, ..., :, position_in_group] = np.moveaxis(
                     values, group_axis, 0
@@ -553,8 +553,10 @@ def scatter_terms(values: np.ndarray) -> np.ndarray:
             )
 
         def group_sum(da: DataArray) -> DataArray:
-            """Sum the constant term within each group, keeping the grouped
-            dimension's position."""
+            """
+            Sum the constant term within each group, keeping the grouped
+            dimension's position.
+            """
             dims_in = list(da.dims)
             group_axis = dims_in.index(group_dim)
             target = [GROUP_DIM if d == group_dim else d for d in dims_in]

From c5d08f422c367cef3cd6d9ec53a972b8cfef1d3e Mon Sep 17 00:00:00 2001
From: FBumann <117816358+FBumann@users.noreply.github.com>
Date: Tue, 21 Jul 2026 13:41:52 +0200
Subject: [PATCH 8/9] perf(groupby): sum const before term scatters to restore
 peak-memory parity

Inlining group_sum(data.const) into the Dataset literal made it evaluate
after the coeffs/vars scatters, stacking its nan-masking temporaries on
top of the retained padded arrays (+31% peak in expr_groupby_sum).
Reducing const first frees them before the padded arrays exist, restoring
the exact peak of master (verified with memray: 171 KB vs 275 KB).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 linopy/expressions.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index 539becdc9..ed56763ab 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -494,7 +494,9 @@ def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset
         padded with fill values. The group dimension replaces the grouped one in
         place: all surviving dimensions are passed as core dims, so ``apply_ufunc``
         emits the target axis order and the padded arrays are allocated once,
-        contiguous. Dask arrays are collapsed to a single chunk (no per-dim
+        contiguous. The constant is reduced before the term scatters so its
+        temporaries are freed before the padded arrays exist, keeping peak memory
+        at input + result. Dask arrays are collapsed to a single chunk (no per-dim
         parallelism) since every dimension is now a core dim.
         """
         data = self.data if data is None else data
@@ -578,11 +580,12 @@ def reduce(values: np.ndarray) -> np.ndarray:
                 output_dtypes=[da.dtype],
             )
 
+        const = group_sum(data.const)
         ds = Dataset(
             {
                 "coeffs": scatter(data.coeffs, fill_value["coeffs"]),
                 "vars": scatter(data.vars, fill_value["vars"]),
-                "const": group_sum(data.const),
+                "const": const,
             }
         )
         return ds.assign_coords({GROUP_DIM: unique_groups})

From 8cfa6f8ae28ae4f5ed46dc4e3c1a64dcaedbb786 Mon Sep 17 00:00:00 2001
From: FBumann <117816358+FBumann@users.noreply.github.com>
Date: Tue, 21 Jul 2026 13:42:00 +0200
Subject: [PATCH 9/9] fix(groupby): keep dim position when the grouper reuses
 the grouped dim's name

Restoring the group dim's position inferred the consumed dims as those
absent from the result, so a grouper named like the dim it groups (e.g.
relabelling 'name' to a coarser 'name') shadowed the consumed dim, the
restore early-returned and the dim drifted to the trailing position on
both paths. Pass the grouper's own dims (DataArray/IndexVariable dims,
Series/DataFrame index name) through _grouper_dims instead; unknown
grouper types (grouper mappings, multi-key lists under use_fallback)
keep the inference.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 linopy/expressions.py          | 33 ++++++++++++++++++++++++++++-----
 test/test_linear_expression.py | 10 +++++++---
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/linopy/expressions.py b/linopy/expressions.py
index ed56763ab..99704b33c 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -192,8 +192,22 @@ def _flatten_multidim_group(
     return series, data
 
 
+def _grouper_dims(group: Any) -> list[Hashable] | None:
+    """
+    The dims a grouper consumes when grouping, or None for grouper types whose
+    dims are unknown here (exotic hashables, validated by xarray's groupby).
+    """
+    if isinstance(group, (DataArray, IndexVariable)):
+        return list(group.dims)
+    if isinstance(group, (pd.Series, pd.DataFrame)):
+        return [group.index.name]
+    return None
+
+
 def _restore_group_dim_position(
-    result: Dataset, original_dims: tuple[str, ...]
+    result: Dataset,
+    original_dims: tuple[Hashable, ...],
+    grouped_dims: Sequence[Hashable] | None = None,
 ) -> Dataset:
     """
     Move the new group dimension into the slot the grouped dimension occupied,
@@ -206,15 +220,23 @@ def _restore_group_dim_position(
     This covers the fallback, which appends the group dim, and normalises the
     scatter kernel's target-ordered output.
 
+    ``grouped_dims`` names the dims the grouping consumed. It must be passed
+    whenever the grouper's name can equal a grouped dim (the replacement dim then
+    shadows the consumed one, so it cannot be inferred from the dims alone);
+    when None, the consumed dims are inferred as those absent from the result.
+
     Numpy-backed terms are made C-contiguous so downstream flattening (LP/solver
     export) ravels a view instead of copying; a no-op when already contiguous, as
     the scatter kernel's output is. Dask arrays are left lazy.
     """
-    consumed = [d for d in original_dims if d not in result.dims]
+    if grouped_dims is None:
+        consumed = [d for d in original_dims if d not in result.dims]
+    else:
+        consumed = [d for d in original_dims if d in set(grouped_dims)]
     if not consumed:
         return result
-    new_dims = [str(d) for d in result.dims if d not in original_dims]
-    order: list[str] = []
+    new_dims = [str(d) for d in result.dims if d not in original_dims or d in consumed]
+    order: list[Hashable] = []
     for d in original_dims:
         if d not in consumed:
             order.append(d)
@@ -441,6 +463,7 @@ def sum(
 
         data = self.data
         original_dims = self.data.coeffs.dims
+        grouped_dims = _grouper_dims(group)
         is_multidim_grouper = (
             isinstance(group, DataArray)
             and group.ndim > 1
@@ -482,7 +505,7 @@ def func(ds: Dataset) -> Dataset:
 
             ds = self.groupby.map(func)
 
-        ds = _restore_group_dim_position(ds, original_dims)
+        ds = _restore_group_dim_position(ds, original_dims, grouped_dims)
         return LinearExpression(ds, self.model)
 
     def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset:
diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py
index f16516149..835798fc2 100644
--- a/test/test_linear_expression.py
+++ b/test/test_linear_expression.py
@@ -1695,10 +1695,14 @@ def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None:
     assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
 
 
+@pytest.mark.parametrize("group_name", ["bus", "name"])
 @pytest.mark.parametrize("use_fallback", [False, True])
-def test_linear_expression_groupby_preserves_dim_position(use_fallback: bool) -> None:
+def test_linear_expression_groupby_preserves_dim_position(
+    use_fallback: bool, group_name: str
+) -> None:
     # Grouping must replace the grouped dim in place, like xarray. `name` sits
-    # between `scenario` and `snapshot`, so the group dim must stay in the middle.
+    # between `scenario` and `snapshot`, so the group dim must stay in the
+    # middle — also when the grouper reuses the grouped dim's name.
     m = Model()
     coords = [
         pd.Index(["s1", "s2"], name="scenario"),
@@ -1707,7 +1711,7 @@ def test_linear_expression_groupby_preserves_dim_position(use_fallback: bool) ->
     ]
     v = m.add_variables(coords=coords, name="x")
     groups = xr.DataArray(
-        ["g1", "g1", "g2"], coords={"name": ["a", "b", "c"]}, name="bus"
+        ["g1", "g1", "g2"], coords={"name": ["a", "b", "c"]}, name=group_name
     )
 
     reference = (