From e610c04770b58d17b33fae1e503361ae4c6bff6b Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:45:39 +0200 Subject: [PATCH 1/4] fix: don't leak aux coords of summed dims onto the term dimension LinearExpression._sum dropped a summed dim's index coords via reset_index but let auxiliary coords defined on that dim ride the stack onto _term, breaking later arithmetic with CoordinateValidationError (#295). Drop them before stacking; aux coords on remaining dims still propagate. Co-Authored-By: Claude Fable 5 --- doc/release_notes.rst | 1 + linopy/expressions.py | 18 +++++++++++++++--- test/test_linear_expression.py | 24 ++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index bffb2daa3..cf23d7cc5 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``. * ``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) Version 0.8.0 ------------- diff --git a/linopy/expressions.py b/linopy/expressions.py index dd8c5c6c1..3df28dc69 100644 --- a/linopy/expressions.py +++ b/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 @@ -135,6 +135,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. @@ -1651,8 +1664,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) ) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 0294fe015..b4ec6f47a 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -487,6 +487,30 @@ def test_linear_expression_sum_drop_zeros(z: Variable) -> None: assert res.nterm == 2 +def test_linear_expression_sum_drops_aux_coord_on_summed_dim() -> None: + # an auxiliary coordinate on a summed dimension must not leak onto the + # helper term dimension, see https://github.com/PyPSA/linopy/issues/295 + m = Model() + v = m.add_variables(coords=[[1, 2, 3]], dims=["A"], name="v") + v = v.assign_coords({"B": ("A", [311, 311, 322])}) + + summed = (1 * v).sum("A") + assert "B" not in summed.coords + con = summed >= xr.DataArray(10) # previously CoordinateValidationError + assert "B" not in con.coords + + assert "B" not in (1 * v).sum().coords + + +def test_linear_expression_sum_keeps_aux_coord_on_other_dim() -> None: + m = Model() + v = m.add_variables(coords=[[1, 2], [1, 2]], dims=["A", "k"], name="v") + v = v.assign_coords({"B": ("k", ["x", "y"])}) + summed = (1 * v).sum("A") + assert "B" in summed.coords + assert summed.coords["B"].dims == ("k",) + + def test_linear_expression_sum_warn_using_dims(z: Variable) -> None: with pytest.warns(DeprecationWarning): (1 * z).sum(dims="dim_0") From 7fb776a3ff04646b7bfc7be7cf6f0d9cd640e24c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:59:50 +0200 Subject: [PATCH 2/4] test: merge the two aux-coord sum tests into one Co-Authored-By: Claude Fable 5 --- test/test_linear_expression.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index b4ec6f47a..c72d4d38c 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -487,30 +487,22 @@ def test_linear_expression_sum_drop_zeros(z: Variable) -> None: assert res.nterm == 2 -def test_linear_expression_sum_drops_aux_coord_on_summed_dim() -> None: +def test_linear_expression_sum_aux_coords() -> None: # an auxiliary coordinate on a summed dimension must not leak onto the # helper term dimension, see https://github.com/PyPSA/linopy/issues/295 m = Model() - v = m.add_variables(coords=[[1, 2, 3]], dims=["A"], name="v") - v = v.assign_coords({"B": ("A", [311, 311, 322])}) + v = m.add_variables(coords=[[1, 2], [1, 2]], dims=["A", "k"], name="v") + v = v.assign_coords({"B": ("A", [311, 322]), "C": ("k", ["x", "y"])}) summed = (1 * v).sum("A") assert "B" not in summed.coords + assert summed.coords["C"].dims == ("k",) con = summed >= xr.DataArray(10) # previously CoordinateValidationError assert "B" not in con.coords assert "B" not in (1 * v).sum().coords -def test_linear_expression_sum_keeps_aux_coord_on_other_dim() -> None: - m = Model() - v = m.add_variables(coords=[[1, 2], [1, 2]], dims=["A", "k"], name="v") - v = v.assign_coords({"B": ("k", ["x", "y"])}) - summed = (1 * v).sum("A") - assert "B" in summed.coords - assert summed.coords["B"].dims == ("k",) - - def test_linear_expression_sum_warn_using_dims(z: Variable) -> None: with pytest.warns(DeprecationWarning): (1 * z).sum(dims="dim_0") From 9cc525d6381b1b2ad541e15e144ad56ef84e4f35 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 23 Jul 2026 14:25:44 +0200 Subject: [PATCH 3/4] test: cover aux-coord drop for matmul and groupby collapse (#295) --- test/test_linear_expression.py | 58 +++++++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index c72d4d38c..71da93b90 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -487,20 +487,54 @@ def test_linear_expression_sum_drop_zeros(z: Variable) -> None: assert res.nterm == 2 -def test_linear_expression_sum_aux_coords() -> None: - # an auxiliary coordinate on a summed dimension must not leak onto the - # helper term dimension, see https://github.com/PyPSA/linopy/issues/295 - m = Model() - v = m.add_variables(coords=[[1, 2], [1, 2]], dims=["A", "k"], name="v") - v = v.assign_coords({"B": ("A", [311, 322]), "C": ("k", ["x", "y"])}) +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 - summed = (1 * v).sum("A") - assert "B" not in summed.coords - assert summed.coords["C"].dims == ("k",) - con = summed >= xr.DataArray(10) # previously CoordinateValidationError - assert "B" not in con.coords + @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 - assert "B" not in (1 * v).sum().coords + @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: From 8f7f907c276131ba23d28208062039b904ef2b37 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:26:39 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/test_linear_expression.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 02eb088ae..54b76dd42 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -504,8 +504,10 @@ def expr(self) -> LinearExpression: [ 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]}), + lambda e: ( + e + @ xr.DataArray([1.0, 1.0, 1.0], dims=["A"], coords={"A": [1, 2, 3]}) + ), id="matmul", ), ],