diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 29d75df35..a80cd8744 100644 --- a/doc/release_notes.rst +++ b/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 diff --git a/linopy/expressions.py b/linopy/expressions.py index 99704b33c..21a4160e9 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 @@ -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) ) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 835798fc2..54b76dd42 100644 --- a/test/test_linear_expression.py +++ b/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")