diff --git a/doc/release_notes.rst b/doc/release_notes.rst index bffb2daa..34c97d10 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -38,6 +38,7 @@ Upcoming Version * LP file export now honors bounds tightened below ``[0, 1]`` on a binary variable via the ``.lower``/``.upper`` setters after creation (e.g. ``upper = 0``). Previously such bounds were written only by ``io_api="direct"`` and dropped by ``io_api="lp"``. (https://github.com/PyPSA/linopy/issues/776) * 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. * ``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/constants.py b/linopy/constants.py index 8723a0bb..7936ef1c 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -76,6 +76,7 @@ class PerformanceWarning(UserWarning): LP_PIECE_DIM = f"{BREAKPOINT_DIM}_piece" PWL_LINK_DIM = "_pwl_var" GROUP_DIM = "_group" +GROUP_STACK_DIM = "_group_stack" FACTOR_DIM = "_factor" CONCAT_DIM = "_concat" CV_DIM = "_cv" diff --git a/linopy/expressions.py b/linopy/expressions.py index dd8c5c6c..85134d47 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -74,6 +74,7 @@ FACTOR_DIM, GREATER_EQUAL, GROUP_DIM, + GROUP_STACK_DIM, HELPER_DIMS, LESS_EQUAL, STACKED_TERM_DIM, @@ -173,6 +174,24 @@ def _multikey_value_frame(group: Any, data: Dataset) -> pd.DataFrame | None: return data[names].to_dataframe()[names] +def _flatten_multidim_group( + group: DataArray, data: Dataset +) -> tuple[pd.Series, Dataset]: + """ + Flatten an N-D DataArray grouper into a 1-D key over a stacked dimension. + + A multi-dimensional grouper defines its groups jointly over all its dims, so + the data is stacked over those dims and the grouper's values are raveled into + a Series aligned with the stacked axis. This lets the fast path reduce over + every grouper dim at once instead of leaking the extra dims into the result. + """ + group_dims = list(group.dims) + data = data.stack({GROUP_STACK_DIM: group_dims}, create_index=False) + series = pd.Series(group.transpose(*group_dims).values.reshape(-1), name=group.name) + series.index.name = GROUP_STACK_DIM + return series, data + + def _unstack_multikey(ds: Dataset, dim: str) -> Dataset: """ Unstack a stacked multi-key group dimension into one dimension per key. @@ -340,6 +359,15 @@ def sum( if multikey_frame is not None: group = multikey_frame + data = self.data + is_multidim_grouper = ( + isinstance(group, DataArray) + and group.ndim > 1 + and set(group.dims) <= set(data.dims) + ) + if is_multidim_grouper and not use_fallback: + group, data = _flatten_multidim_group(group, data) + fast_path_types = (pd.Series, pd.DataFrame, xr.DataArray) if isinstance(group, fast_path_types) and not use_fallback: final_group_name = ( @@ -358,7 +386,7 @@ def sum( ) assert isinstance(group, pd.Series) - ds = self._grouped_sum(group) + ds = self._grouped_sum(group, data) if multikey_decode is not None: ds = _restore_multikey_index(ds, multikey_decode) @@ -375,7 +403,7 @@ def func(ds: Dataset) -> Dataset: return self.map(func) - def _grouped_sum(self, group: pd.Series) -> Dataset: + def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset: """ Sum groups by scattering all terms directly into the final padded arrays. @@ -384,7 +412,7 @@ def _grouped_sum(self, group: pd.Series) -> Dataset: padded with fill values. Only the result arrays are allocated, keeping peak memory at input + result. """ - data = self.data + data = self.data if data is None else data group_dim = group.index.name fill_value = LinearExpression._fill_value diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 0294fe01..8e0ef7dd 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -1668,9 +1668,8 @@ def test_multiindex_level( assert grouped.vars.transpose(level, TERM_DIM).values.tolist() == vars_ -@pytest.mark.parametrize("use_fallback", [True]) +@pytest.mark.parametrize("use_fallback", [True, False]) def test_linear_expression_groupby_ndim(z: Variable, use_fallback: bool) -> None: - # TODO: implement fallback for n-dim groupby, see https://github.com/PyPSA/linopy/issues/299 expr = 1 * z groups = xr.DataArray([[1, 1, 2], [1, 3, 3]], coords=z.coords) grouped = expr.groupby(groups).sum(use_fallback=use_fallback) @@ -1678,6 +1677,22 @@ def test_linear_expression_groupby_ndim(z: Variable, use_fallback: bool) -> None # there are three groups, 1, 2 and 3, the largest group has 3 elements assert (grouped.data.group == [1, 2, 3]).all() assert grouped.nterm == 3 + assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True)) + + +def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None: + # a 2-D grouper reduces over both its dims jointly, leaving unrelated dims + # untouched, see https://github.com/PyPSA/linopy/issues/823 + m = Model() + v = m.add_variables(coords=[[0, 1], [0, 1], [0, 1]], dims=["i", "j", "k"]) + groups = xr.DataArray( + [[1, 1], [2, 2]], dims=["i", "j"], coords={"i": [0, 1], "j": [0, 1]}, name="g" + ) + expr = 1 * v + grouped = expr.groupby(groups).sum() + assert set(grouped.dims) == {"g", "k", "_term"} + assert (grouped.data.g == [1, 2]).all() + assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True)) @pytest.mark.parametrize("use_fallback", [True, False])