fix: don't leak aux coords of summed dims onto the term dimension

doc/release_notes.rst

@@ -42,6 +42,7 @@ Upcoming Version
 * ``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.
+* Summing an expression over a dimension that carries an auxiliary (non-dimension) coordinate no longer leaks that coordinate onto the internal term dimension, where it broke later arithmetic with a ``CoordinateValidationError``. Auxiliary coordinates on the remaining dimensions still propagate. (https://github.com/PyPSA/linopy/issues/295)
 * ``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)
 
 Version 0.8.0

linopy/expressions.py

@@ -11,7 +11,7 @@
 import logging
 import operator
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, field
 from itertools import product, zip_longest
 from typing import TYPE_CHECKING, Any, Self, TypeAlias, TypeVar, cast, overload
@@ -136,6 +136,19 @@ def _expr_unwrap(
 logger = logging.getLogger(__name__)
 
 
+def _drop_coords_on_dims(ds: Dataset, dims: Iterable[Hashable]) -> Dataset:
+    """
+    Drop every coordinate touching the given dimensions.
+
+    A coordinate on a dimension that is being reduced away must not survive
+    the reduction: an auxiliary coordinate left in place would ride a
+    subsequent stack onto the helper term dimension and break later
+    alignment (https://github.com/PyPSA/linopy/issues/295).
+    """
+    dims = set(dims)
+    return ds.drop_vars(k for k, c in ds.coords.items() if dims & set(c.dims))
+
+
 def _resolve_group(group: Any, data: Dataset) -> Any:
     """
     Normalize a groupby key.
@@ -1811,8 +1824,7 @@ def _sum(
         else:
             dim = [d for d in dim if d != TERM_DIM]
             ds = (
-                data[["coeffs", "vars"]]
-                .reset_index(dim, drop=True)
+                _drop_coords_on_dims(data[["coeffs", "vars"]], dim)
                 .rename({TERM_DIM: STACKED_TERM_DIM})
                 .stack({TERM_DIM: [STACKED_TERM_DIM] + dim}, create_index=False)
             )

test/test_linear_expression.py

@@ -487,6 +487,58 @@ def test_linear_expression_sum_drop_zeros(z: Variable) -> None:
     assert res.nterm == 2
 
 
+class TestCollapseAuxCoords:
+    # collapsing a dimension carrying an auxiliary coordinate must drop it, not
+    # leak it onto the term dimension where it breaks later arithmetic with a
+    # CoordinateValidationError, see https://github.com/PyPSA/linopy/issues/295
+
+    @pytest.fixture
+    def expr(self) -> LinearExpression:
+        m = Model()
+        v = m.add_variables(coords=[[1, 2, 3], [1, 2]], dims=["A", "k"], name="v")
+        v = v.assign_coords({"B": ("A", [311, 311, 322]), "C": ("k", ["p", "q"])})
+        return 1 * v
+
+    @pytest.mark.parametrize(
+        "collapse",
+        [
+            pytest.param(lambda e: e.sum("A"), id="sum"),
+            pytest.param(
+                lambda e: (
+                    e
+                    @ xr.DataArray([1.0, 1.0, 1.0], dims=["A"], coords={"A": [1, 2, 3]})
+                ),
+                id="matmul",
+            ),
+        ],
+    )
+    def test_contracting_a_dim_drops_its_aux_coord(
+        self, expr: LinearExpression, collapse: Any
+    ) -> None:
+        res = collapse(expr)
+        assert "B" not in res.coords
+        assert res.coords["C"].dims == ("k",)
+        con = res >= xr.DataArray(10)
+        assert "B" not in con.coords
+
+    @pytest.mark.parametrize("use_fallback", [False, True])
+    def test_groupby_sum_keeps_aux_coord_on_kept_dim(
+        self, expr: LinearExpression, use_fallback: bool
+    ) -> None:
+        grouped = expr.groupby("B").sum(use_fallback=use_fallback)
+        assert "A" not in grouped.coords
+        assert grouped.coords["C"].dims == ("k",)
+        con = grouped >= xr.DataArray(10)
+        assert set(con.coords) == {"B", "k", "C"}
+
+    def test_summing_all_dims_drops_all_aux_coords(
+        self, expr: LinearExpression
+    ) -> None:
+        summed = expr.sum()
+        assert "B" not in summed.coords
+        assert "C" not in summed.coords
+
+
 def test_linear_expression_sum_warn_using_dims(z: Variable) -> None:
     with pytest.warns(DeprecationWarning):
         (1 * z).sum(dims="dim_0")