diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d79eb96c..8af76caa 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -53,6 +53,7 @@ Upcoming Version **Bug fixes** +* ``sum()`` over a dimension no longer raises when another dimension of the expression has size 0; it returns an expression without terms over the kept coordinates, as summing over the empty dimension itself already did. (https://github.com/PyPSA/linopy/issues/906) * A multi-key ``groupby`` now returns its groups sorted by key tuple, like the single-key path. The key combinations were numbered by iterating a ``set``, so the group order was arbitrary and changed between processes with ``PYTHONHASHSEED``. * ``Solver.close()`` no longer leaves dangling native handles behind. The solver model is now dropped before the environment that owns it, instead of after. And the COPT and MindOpt file interfaces no longer hand back a model they already disposed: after a file-based COPT or MindOpt solve, ``model.solver_model`` is ``None`` rather than a handle into freed memory. (`#899 `__) diff --git a/linopy/expressions.py b/linopy/expressions.py index b7e05205..f4b93823 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -181,6 +181,24 @@ def _expr_unwrap( logger = logging.getLogger(__name__) +def _stack_into_term_dim(ds: Dataset, dims: list[Hashable]) -> Dataset: + """ + Stack ``dims`` into the term dimension. + + ``Dataset.stack`` reshapes with an inferred ``-1`` that numpy cannot + resolve on zero-element data (https://github.com/PyPSA/linopy/issues/906), + so an empty dataset is stacked onto a term dimension of size 0 instead. + """ + if ds.vars.size: + return ds.stack({TERM_DIM: dims}, create_index=False) + + for name, da in ds.data_vars.items(): + kept = [d for d in da.dims if d not in dims] + shape = [da.sizes[d] for d in kept] + [0] + ds[name] = DataArray(da.values.reshape(shape), dims=kept + [TERM_DIM]) + return ds + + def _drop_coords_on_dims(ds: Dataset, dims: Iterable[Hashable]) -> Dataset: """ Drop every coordinate touching the given dimensions. @@ -2135,11 +2153,10 @@ def _sum( ds = xr.Dataset({"vars": vars, "coeffs": coeffs, "const": const}) else: dim = [d for d in dim if d != TERM_DIM] - ds = ( - _drop_coords_on_dims(data[["coeffs", "vars"]], dim) - .rename({TERM_DIM: STACKED_TERM_DIM}) - .stack({TERM_DIM: [STACKED_TERM_DIM] + dim}, create_index=False) + ds = _drop_coords_on_dims(data[["coeffs", "vars"]], dim).rename( + {TERM_DIM: STACKED_TERM_DIM} ) + ds = _stack_into_term_dim(ds, [STACKED_TERM_DIM] + dim) ds = assign_multiindex_safe(ds, const=data.const.sum(dim)) return ds diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 454245aa..7550bb6b 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -30,7 +30,7 @@ merge, options, ) -from linopy.constants import HELPER_DIMS, TERM_DIM +from linopy.constants import FACTOR_DIM, HELPER_DIMS, TERM_DIM from linopy.expressions import ScalarLinearExpression from linopy.testing import assert_linequal, assert_quadequal from linopy.variables import ScalarVariable @@ -455,6 +455,27 @@ def test_linear_expression_sum( assert len(expr.coords["dim_2"]) == 10 +@pytest.mark.parametrize("dim", ["line", ["line", "cycle"], None]) +def test_linear_expression_sum_with_empty_sibling_dim( + m: Model, dim: str | list[str] | None +) -> None: + coords = [ + pd.RangeIndex(4, name="t"), + pd.Index([], name="cycle"), + pd.Index(["a", "b"], name="line"), + ] + x = m.add_variables(coords=coords, name="xe") + expr = 2 * x + + summed = expr.sum(dim) + reduced = expr.const.sum(dim) + assert summed.sizes == {**reduced.sizes, TERM_DIM: summed.nterm} + assert_linequal(summed, LinearExpression.from_constant(m, reduced)) + quad = (x * x).sum(dim) + assert quad.sizes == {**reduced.sizes, FACTOR_DIM: 2, TERM_DIM: quad.nterm} + assert quad.nterm == summed.nterm + + @pytest.mark.legacy def test_linear_expression_sum_with_const( x: Variable, y: Variable, z: Variable, v: Variable