diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 34c97d10..494cb655 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 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. Version 0.8.0 diff --git a/linopy/expressions.py b/linopy/expressions.py index 85134d47..b1166d08 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -215,19 +215,55 @@ def _unstack_multikey(ds: Dataset, dim: str) -> Dataset: return ds.unstack(dim, fill_value=LinearExpression._fill_value) +def _check_grouper_alignment(group: Any, data: Dataset) -> None: + """ + Ensure an indexed grouper's labels match the data along each shared dim. + + The fast path matches a ``pd.Series``, ``pd.DataFrame`` or ``DataArray`` + grouper to the expression by position, so a grouper whose coordinates are + reordered -- or a different set entirely -- relative to the expression would + silently regroup. linopy does not reindex the grouper: it checks that the + labels match and raises otherwise, leaving the caller to align the grouper + explicitly. A grouper without an index along a dimension has nothing to align + by and keeps the positional match. + """ + shared: list[tuple[Hashable, pd.Index]] + if isinstance(group, (pd.Series, pd.DataFrame)): + shared = [(group.index.name, group.index)] + elif isinstance(group, DataArray): + shared = [ + (dim, group.get_index(dim)) for dim in group.dims if dim in group.indexes + ] + else: + return + for dim, index in shared: + if dim not in data.indexes or index.equals(data.indexes[dim]): + continue + detail = ( + "the same labels in a different order" + if set(index) == set(data.indexes[dim]) + else "a different set of labels" + ) + raise ValueError( + f"the grouper's labels along dimension {dim!r} do not match the " + f"expression's coordinates ({detail}). linopy matches groupers by " + f"position and does not reindex; reorder the grouper to the " + f"expression's {dim!r} coordinates before grouping." + ) + + def _encode_multikey_group( - frame: pd.DataFrame, index: pd.Index + frame: pd.DataFrame, ) -> tuple[pd.Series, tuple[dict, pd.Index]]: """ Encode a multi-key group frame as a single integer-coded Series. ``_grouped_sum`` groups by one key, so each row's tuple of key values is - mapped to an integer. The returned ``(int_map, columns)`` lets - :func:`_restore_multikey_index` rebuild the MultiIndex on the result. + mapped to an integer. The frame is already aligned to the data (see + :func:`_check_grouper_alignment`), so no reindexing is needed. The returned + ``(int_map, columns)`` lets :func:`_restore_multikey_index` rebuild the + MultiIndex on the result. """ - index_name = frame.index.name - frame = frame.reindex(index) - frame.index.name = index_name int_map = get_index_map(*frame.values.T) coded = frame.apply(tuple, axis=1).map(int_map) return coded, (int_map, frame.columns) @@ -352,6 +388,7 @@ def sum( ) group = _resolve_group(self.group, self.data) + _check_grouper_alignment(group, self.data) multikey_frame = ( None if use_fallback else _multikey_value_frame(group, self.data) @@ -381,9 +418,7 @@ def sum( multikey_decode = None if isinstance(group, pd.DataFrame): - group, multikey_decode = _encode_multikey_group( - group, self.data.indexes[group.index.name] - ) + group, multikey_decode = _encode_multikey_group(group) assert isinstance(group, pd.Series) ds = self._grouped_sum(group, data) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 8e0ef7dd..92e2510f 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -1695,6 +1695,75 @@ def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None: assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True)) +class TestGroupbyGrouperAlignment: + """ + A ``pd.Series``/``DataArray`` grouper whose labels are reordered or a + different set relative to the expression must raise, not silently regroup + by position. See https://github.com/PyPSA/linopy/issues/827. + """ + + @staticmethod + def _expr() -> LinearExpression: + m = Model() + v = m.add_variables(coords=[[0, 1, 2, 3]], dims=["i"], name="v") + return 1 * v + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_reordered_series_raises(self, use_fallback: bool) -> None: + expr = self._expr() + s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 3], name="i"), name="g") + with pytest.raises(ValueError, match="different order"): + expr.groupby(s.iloc[::-1]).sum(use_fallback=use_fallback) + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_reordered_dataarray_raises(self, use_fallback: bool) -> None: + expr = self._expr() + da = xr.DataArray([1, 1, 2, 2], coords={"i": [0, 1, 2, 3]}, name="g") + with pytest.raises(ValueError, match="different order"): + expr.groupby(da.isel(i=slice(None, None, -1))).sum( + use_fallback=use_fallback + ) + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_mismatched_label_set_raises(self, use_fallback: bool) -> None: + expr = self._expr() + s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 9], name="i"), name="g") + with pytest.raises(ValueError, match="different set of labels"): + expr.groupby(s).sum(use_fallback=use_fallback) + + def test_reordered_dataframe_raises(self) -> None: + # the DataFrame grouper path aligns by label like Series/DataArray + expr = self._expr() + df = pd.DataFrame({"g": [1, 1, 2, 2]}, index=pd.Index([0, 1, 2, 3], name="i")) + with pytest.raises(ValueError, match="different order"): + expr.groupby(df.iloc[::-1]).sum() + + def test_reordered_ndim_dataarray_raises(self) -> None: + m = Model() + v = m.add_variables(coords=[[0, 1], [0, 1]], dims=["i", "j"], name="v") + groups = xr.DataArray( + [[1, 1], [2, 2]], dims=["i", "j"], coords={"i": [1, 0], "j": [0, 1]} + ) + with pytest.raises(ValueError, match="dimension 'i'"): + (1 * v).groupby(groups).sum() + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_aligned_grouper_unaffected(self, use_fallback: bool) -> None: + # an aligned grouper still groups by label as before + expr = self._expr() + s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 3], name="i"), name="g") + grouped = expr.groupby(s).sum(use_fallback=use_fallback) + assert (grouped.data.g == [1, 2]).all() + assert grouped.nterm == 2 + + def test_unindexed_grouper_matches_positionally(self) -> None: + # a grouper without an index along the dim has nothing to align by + expr = self._expr() + da = xr.DataArray([1, 1, 2, 2], dims=["i"], name="g") # no "i" coordinate + grouped = expr.groupby(da).sum() + assert (grouped.data.g == [1, 2]).all() + + @pytest.mark.parametrize("use_fallback", [True, False]) def test_linear_expression_groupby_with_name(v: Variable, use_fallback: bool) -> None: expr = 1 * v @@ -1844,15 +1913,16 @@ def test_linear_expression_groupby_with_dataarray( def test_linear_expression_groupby_with_dataframe_non_aligned(v: Variable) -> None: + # a DataFrame grouper aligns by label like Series/DataArray: a reordered + # index raises instead of silently regrouping. See issue #827. expr = 1 * v groups = pd.DataFrame( {"a": [1] * 10 + [2] * 10, "b": list(range(4)) * 5}, index=v.indexes["dim_2"] ) - target = expr.groupby(groups).sum() + expr.groupby(groups).sum() # aligned: fine - groups_non_aligned = groups[::-1] - grouped = expr.groupby(groups_non_aligned).sum() - assert_linequal(grouped, target) + with pytest.raises(ValueError, match="different order"): + expr.groupby(groups[::-1]).sum() @pytest.mark.parametrize("use_fallback", [True, False])