diff --git a/doc/release_notes.rst b/doc/release_notes.rst index de87c522..dc3a320a 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -73,6 +73,8 @@ Version 0.9.1 * ``LinearExpression.reindex`` keeps a sparse (CSR-backed) expression sparse for plain label changes — reorder, add, or drop coordinates — instead of expanding to the dense rectangle. New coordinates become absent cells and the result matches the dense reindex, so a ``groupby(sparse=True) → reindex → merge`` chain stays sparse and the build peak stays low (v1 only; other arguments fall back to dense). (`#932 `__) +* ``linopy.merge`` (and ``+`` / ``-`` / ``.add`` / ``.sub`` with an explicit ``join=``) keeps sparse (CSR-backed) expressions sparse when the operands live on different label subsets of the same dimensions, e.g. a nodal balance summing grouped generator, line and load terms. The payloads are aligned row-wise onto the joined grid (``outer`` / ``inner`` / ``left`` / ``right`` / ``override``) instead of falling back to the dense rectangle; the positions the join creates carry the same fill as the dense path (zero, or absent with ``fill_value=linopy.ABSENT``) and the result equals the dense one. Dense operands on a different grid are converted on the fly. (`#749 `__) + **Bug fixes** * An SOS set is now ordered by declaration, not by the values of its coordinates. Labels are names, but they were handed to the solver as SOS weights, so element-for-element identical models could reach different optima — silently changing who is adjacent in a ``sos_type=2`` set — just because their coordinates were named differently. **Behaviour change:** ascending coordinates are unaffected (this covers every piecewise formulation), descending ones now run in reverse with the same adjacency, and only a ``sos_type=2`` set whose numeric coordinates neither ascend nor descend changes meaning — that case now warns, and sorting the index restores the old order. (`#892 `__, `#893 `__) diff --git a/linopy/expressions.py b/linopy/expressions.py index f312bec2..9ab501ac 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -3272,7 +3272,9 @@ def merge( if issubclass(cls, LinearExpression) and not has_quad_expression: from linopy.sparse_expression import try_csr_merge - csr_result = try_csr_merge(exprs, dim=dim, join=join, kwargs=kwargs) + csr_result = try_csr_merge( + exprs, dim=dim, join=join, fill_value=fill_value, kwargs=kwargs + ) if csr_result is not None: return csr_result diff --git a/linopy/sparse_expression.py b/linopy/sparse_expression.py index d1abee6a..8465a3b2 100644 --- a/linopy/sparse_expression.py +++ b/linopy/sparse_expression.py @@ -26,8 +26,10 @@ import pandas as pd import scipy.sparse from xarray import Dataset +from xarray.core.types import JoinOptions from linopy.constants import TERM_DIM +from linopy.semantics import FillValueLike, join_fill if TYPE_CHECKING: from linopy.expressions import LinearExpression @@ -95,20 +97,13 @@ def from_grouper( ) @classmethod - def from_expression( - cls, expr: LinearExpression, template: CSRPayload - ) -> CSRPayload | None: - """Convert a dense expression on the template's grid, else None.""" - if set(expr.coord_dims) != set(template.grid_dims): - return None - for d in expr.coord_dims: - if not expr.data.get_index(d).equals(template.indexes[str(d)]): - return None - first = template.grid_dims[0] - codes = np.arange(len(template.indexes[first])) - return cls._from_scatter( - expr, template.grid_dims, template.indexes, first, first, codes, False - ) + def from_expression(cls, expr: LinearExpression) -> CSRPayload: + """Convert a dense expression to CSR form on its own coordinate grid.""" + grid_dims = tuple(str(d) for d in expr.coord_dims) + indexes = {d: expr.data.get_index(d).rename(d) for d in grid_dims} + first = grid_dims[0] + codes = np.arange(len(indexes[first])) + return cls._from_scatter(expr, grid_dims, indexes, first, first, codes, False) @classmethod def _from_scatter( @@ -163,15 +158,25 @@ def _from_scatter( def scaled(self, factor: float) -> CSRPayload: return replace(self, csr=self.csr * factor, const=self.const * factor) - def reindexed(self, indexes: dict[str, pd.Index]) -> CSRPayload: + def reindexed( + self, + indexes: dict[str, pd.Index], + grid_dims: tuple[str, ...] | None = None, + fill: float = np.nan, + ) -> CSRPayload: """ - Remap rows onto new per-dim indexes without the dense rectangle: - dropped labels vanish, new labels are absent cells (NaN const). + Remap rows onto new per-dim indexes, optionally in a new dim order, + without the dense rectangle: dropped labels vanish, new labels get + ``fill`` as their constant (NaN: absent cells). """ - shape, strides = _grid_layout(self.grid_dims, indexes) + grid_dims = grid_dims or self.grid_dims + shape, strides = _grid_layout(grid_dims, indexes) + stride_of = dict(zip(grid_dims, strides)) positions = [indexes[d].get_indexer(self.indexes[d]) for d in self.grid_dims] valid = _flat_cells([pos == -1 for pos in positions]) == 0 - row_map = _flat_cells([pos * s for pos, s in zip(positions, strides)]) + row_map = _flat_cells( + [pos * stride_of[d] for pos, d in zip(positions, self.grid_dims)] + ) coo = self.csr.tocoo() keep = valid[coo.coords[0]] @@ -181,10 +186,14 @@ def reindexed(self, indexes: dict[str, pd.Index]) -> CSRPayload: coo = scipy.sparse.coo_array( (coo.data[keep], (rows, cols)), shape=(n_cells, self.csr.shape[1]) ) - const = np.full(n_cells, np.nan) + const = np.full(n_cells, fill) const[row_map[valid]] = self.const[valid] return replace( - self, csr=scipy.sparse.csr_array(coo), const=const, indexes=indexes + self, + csr=scipy.sparse.csr_array(coo), + const=const, + grid_dims=grid_dims, + indexes=indexes, ) def filled(self, value: float) -> CSRPayload: @@ -211,17 +220,20 @@ def add(self, other: CSRPayload) -> CSRPayload: Sparse matrix addition == merge along the term dimension. Goes through COO so explicit zero coefficients survive (scipy's ``+`` drops them), keeping a cell with only zero-coefficient terms distinguishable from - an empty cell, as on the dense path. + an empty cell, as on the dense path. A cell absent in either operand + is absent in the sum and carries no terms (v1 dead-term invariant). """ + const = self.const + other.const a, b = self.csr.tocoo(), other.csr.tocoo() shape = (self.n_cells, max(a.shape[1], b.shape[1])) rows = np.concatenate([a.coords[0], b.coords[0]]) cols = np.concatenate([a.coords[1], b.coords[1]]) data = np.concatenate([a.data, b.data]) - coo = scipy.sparse.coo_array((data, (rows, cols)), shape=shape) - return replace( - self, csr=scipy.sparse.csr_array(coo), const=self.const + other.const + present = ~np.isnan(const)[rows] + coo = scipy.sparse.coo_array( + (data[present], (rows[present], cols[present])), shape=shape ) + return replace(self, csr=scipy.sparse.csr_array(coo), const=const) def materialize(self) -> LinearExpression: """ @@ -275,14 +287,55 @@ def _flat_cells(axis_positions: list[np.ndarray]) -> np.ndarray: return cells.reshape(-1) +def _aligned( + payloads: list[CSRPayload], join: JoinOptions | None, fill: float +) -> list[CSRPayload] | None: + """ + Conform the payloads to the grid an explicit join produces, cells the + join creates carrying ``fill`` as constant. None where the dense path + owns the semantics: ``exact`` and the auto-detected join raise there on + differing grids, ``override`` on differing shapes, any join on + non-unique labels. + """ + template = payloads[0] + dims = template.grid_dims + if any(not p.indexes[d].is_unique for p in payloads for d in dims): + return None + if join == "override": + if any(p.grid_dims != dims or p.shape != template.shape for p in payloads): + return None + return [replace(p, indexes=template.indexes) for p in payloads] + if join == "left": + indexes = template.indexes + elif join == "right": + indexes = payloads[-1].indexes + elif join in ("outer", "inner"): + combine = pd.Index.union if join == "outer" else pd.Index.intersection + indexes = {} + for d in dims: + index = template.indexes[d] + for p in payloads[1:]: + index = combine(index, p.indexes[d]) + indexes[d] = pd.Index(index, name=d) + else: + return None + return [p.reindexed(indexes, dims, fill) for p in payloads] + + def try_csr_merge( - exprs: Any, dim: str, join: Any, kwargs: dict + exprs: Any, + dim: str, + join: JoinOptions | None, + fill_value: FillValueLike, + kwargs: dict[str, Any], ) -> LinearExpression | None: """ Sparse branch of :func:`linopy.expressions.merge`: combine plain - LinearExpressions on one shared grid (CSR-backed or dense-convertible), - where any join produces the identical result. Returns None to fall - through to the dense path. + LinearExpressions over one set of grid dimensions (CSR-backed or + dense-convertible) as sparse matrix addition. Grids that differ in + their labels are aligned row-wise onto the joined grid, the cells the + join creates carrying the fill of the dense path (zero, or NaN for + ``fill_value=ABSENT``). Returns None to fall through to the dense path. """ from linopy.expressions import LinearExpression @@ -290,18 +343,23 @@ def try_csr_merge( return None if not all(type(e) is LinearExpression for e in exprs): return None - payloads = [e._payload for e in exprs if e._payload is not None] - if not payloads: + if all(e._payload is None for e in exprs): return None - template = payloads[0] - if not all(template.same_grid(p) for p in payloads[1:]): + dims = set(exprs[0].coord_dims) + if any(set(e.coord_dims) != dims for e in exprs[1:]): + return None + if any(e._payload is None and set(e.data.coords) - dims for e in exprs): return None - combined: CSRPayload | None = None - for e in exprs: - payload = e._payload or CSRPayload.from_expression(e, template) - if payload is None: + payloads = [e._payload or CSRPayload.from_expression(e) for e in exprs] + template = payloads[0] + if not all(template.same_grid(p) for p in payloads[1:]): + aligned = _aligned(payloads, join, join_fill(fill_value, 0.0)) + if aligned is None: return None - combined = payload if combined is None else combined.add(payload) - assert combined is not None + payloads = aligned + + combined = payloads[0] + for payload in payloads[1:]: + combined = combined.add(payload) return LinearExpression._from_payload(combined, exprs[0].model) diff --git a/test/test_sparse_groupby.py b/test/test_sparse_groupby.py index bbd48055..6255e91d 100644 --- a/test/test_sparse_groupby.py +++ b/test/test_sparse_groupby.py @@ -15,6 +15,7 @@ import polars as pl import pytest import xarray as xr +from xarray.core.types import JoinOptions import linopy from linopy import LinearExpression, Model, Variable @@ -35,6 +36,7 @@ class Case: m: Model gen_p: Variable flow: Variable + flow_t: Variable eff: xr.DataArray gbus: pd.Series bus0: pd.Series @@ -63,6 +65,7 @@ def base_model( m = linopy.Model() gen_p = m.add_variables(coords=[gens, snaps], name="gen_p") flow = m.add_variables(coords=[lines, snaps], name="flow") + flow_t = m.add_variables(coords=[snaps, lines], name="flow_t") gbus = pd.Series(buses[gen_bus], index=gens, name="bus") bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") @@ -71,7 +74,7 @@ def base_model( rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" ).sortby("bus") eff = xr.DataArray(rng.uniform(0.5, 1.5, len(gens)), coords=[gens]) - return Case(m, gen_p, flow, eff, gbus, bus0, bus1, load) + return Case(m, gen_p, flow, flow_t, eff, gbus, bus0, bus1, load) def canon(df: pl.DataFrame) -> pl.DataFrame: @@ -444,3 +447,173 @@ def test_rename_stays_csr_and_matches_dense() -> None: dense = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=False).rename(bus="node") assert sparse.coord_dims == ("node", "snapshot") assert_linequal(sparse, dense) + + +def cross_grid_parts( + c: Case, sparse: bool, lines: tuple[str, ...] = ("line1", "line2") +) -> list[LinearExpression]: + """Generation on all buses, flow on a line subset only, snapshot-major on the flow side.""" + gen = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=sparse) + flow_t = 1.0 * c.flow_t.loc[:, list(lines)] + flow = flow_t.groupby(c.bus0.loc[list(lines)]).sum(sparse=sparse) + return [gen, flow] + + +def assert_terms_equal(a: LinearExpression, b: LinearExpression) -> None: + """``assert_linequal`` up to the width of the (non-contractual) term axis.""" + width = max(a.nterm, b.nterm) + padded = [] + for e in (a, b): + pad = {"_term": (0, width - e.nterm)} + fill = {"vars": -1, "coeffs": np.nan} + padded.append(LinearExpression(e.data.pad(pad, constant_values=fill), e.model)) + assert_linequal(*padded) + + +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +@pytest.mark.parametrize("order", ["gen-flow", "flow-gen"]) +def test_cross_grid_merge_stays_csr_and_matches_dense( + join: JoinOptions, order: str +) -> None: + require_v1() + c = base_model() + sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) + if order == "flow-gen": + sparse, dense = sparse[::-1], dense[::-1] + res = linopy.merge(sparse, join=join) + assert res._payload is not None + assert res.coord_dims == dense[0].coord_dims + assert_terms_equal(res, linopy.merge(dense, join=join)) + + +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +def test_three_operand_cross_grid_merge_matches_dense(join: JoinOptions) -> None: + require_v1() + c = base_model() + third_lines = ("line3", "line4") + sparse = cross_grid_parts(c, True) + cross_grid_parts(c, True, third_lines)[1:] + dense = cross_grid_parts(c, False) + cross_grid_parts(c, False, third_lines)[1:] + res = linopy.merge(sparse, join=join) + assert res._payload is not None + assert_terms_equal(res, linopy.merge(dense, join=join)) + + +def test_cross_grid_merge_absent_fill_matches_dense() -> None: + require_v1() + c = base_model() + sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) + res = linopy.merge(sparse, join="outer", fill_value=linopy.ABSENT) + expected = linopy.merge(dense, join="outer", fill_value=linopy.ABSENT) + assert res._payload is not None + filled = res.fillna(0) + assert filled._payload is not None + assert_terms_equal(filled, expected.fillna(0)) + assert_terms_equal(res, expected) + assert res.const.isnull().sum() == 3 * c.load.sizes["snapshot"] + + +def test_cross_grid_merge_keeps_absent_cell_absent() -> None: + require_v1() + c = base_model() + sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) + mask = xr.DataArray([True, False], coords=[dense[1].indexes["bus"]]) + dense[1] = dense[1].where(mask) + res = linopy.merge([sparse[0], dense[1]], join="outer") + assert res._payload is not None + assert_terms_equal(res, linopy.merge(dense, join="outer")) + + +def test_cross_grid_merge_mixed_dense_operand_stays_csr() -> None: + require_v1() + c = base_model() + sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) + res = sparse[0].add(dense[1], join="outer") + assert res._payload is not None + assert_terms_equal(res, dense[0].add(dense[1], join="outer")) + + +@pytest.mark.parametrize( + "kwargs, error", + [ + ({}, ValueError), + ({"join": "exact"}, xr.AlignmentError), + ({"join": "override"}, xr.AlignmentError), + ], + ids=["auto", "exact", "override"], +) +def test_cross_grid_merge_raises_like_dense(kwargs: dict, error: type) -> None: + require_v1() + c = base_model() + with pytest.raises(error): + linopy.merge(cross_grid_parts(c, False), **kwargs) + with pytest.raises(error): + linopy.merge(cross_grid_parts(c, True), **kwargs) + + +def test_merge_with_aux_coord_operand_raises_like_dense() -> None: + require_v1() + c = base_model() + sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) + tag = xr.DataArray(["x", "y"], coords=[dense[1].indexes["bus"]]) + tagged = LinearExpression(dense[1].data.assign_coords(tag=tag), c.m) + with pytest.raises(xr.MergeError, match="conflicting values for variable 'tag'"): + linopy.merge([dense[0], tagged], join="outer") + with pytest.raises(xr.MergeError, match="conflicting values for variable 'tag'"): + linopy.merge([sparse[0], tagged], join="outer") + + +def test_cross_grid_merge_with_duplicate_labels_raises_like_dense() -> None: + require_v1() + c = base_model() + dup = pd.Index(["bus1", "bus1", "bus2"], name="bus") + shed = 1.0 * c.m.add_variables( + coords=[dup, c.load.indexes["snapshot"]], name="shed" + ) + gen = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + dense = (c.eff * c.gen_p).groupby(c.gbus).sum() + with pytest.raises(ValueError, match="cannot reindex or align"): + linopy.merge([dense, shed], join="left") + with pytest.raises(ValueError, match="cannot reindex or align"): + linopy.merge([gen, shed], join="left") + + +def test_override_merge_same_shape_stays_csr() -> None: + require_v1() + c = base_model() + gen = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + flow = (1.0 * c.flow).groupby(c.bus1.str.upper()).sum(sparse=True) + res = linopy.merge([gen, flow], join="override") + assert res._payload is not None + dense = [ + (c.eff * c.gen_p).groupby(c.gbus).sum(), + (1.0 * c.flow).groupby(c.bus1.str.upper()).sum(), + ] + assert_terms_equal(res, linopy.merge(dense, join="override")) + + +def test_cross_grid_balance_freezes_csr() -> None: + require_v1() + c1, c2 = base_model(), base_model() + lhs1 = linopy.merge(cross_grid_parts(c1, False), join="outer") + con1 = c1.m.add_constraints(lhs1 == c1.load, name="bal") + lhs2 = linopy.merge(cross_grid_parts(c2, True), join="outer") + assert lhs2._payload is not None + con2 = c2.m.add_constraints(lhs2 == c2.load, name="bal", freeze=True) + assert isinstance(con2, CSRConstraint) + assert_frozen_equal(con1, con2) + + +def test_cross_grid_merge_peak_memory() -> None: + require_v1() + sizes = (200,) + (1,) * 299 + n_snap = 50 + dense_rectangle_bytes = len(sizes) * n_snap * max(sizes) * 16 + c = base_model(gens_per_bus=sizes, n_snap=n_snap) + tracemalloc.start() + try: + lhs = linopy.merge(cross_grid_parts(c, True), join="outer") + c.m.add_constraints(lhs == c.load, name="bal", freeze=True) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert peak < dense_rectangle_bytes / 4