From 6b1532f3200297c0734d871f717d41f06053010b Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 13:06:11 +0200 Subject: [PATCH 1/2] perf(csr): avoid CSR copies and full-grid scaling when freezing constraints (#977) Share the lhs matrix when every row stays active, copy only before eliminate_zeros, and pick mask and scaling values at the active rows without expanding broadcast views. Scaling is still validated in full. --- linopy/constraints.py | 55 ++++++++++++++++++++++++++-------- linopy/model.py | 14 +++++---- linopy/scaling.py | 8 +++-- test/test_csr.py | 70 ++++++++++++++++++++++++++++++++++++++----- 4 files changed, 118 insertions(+), 29 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 1f0bbdde..320c863d 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -939,24 +939,53 @@ def _replace(self, **changes: Any) -> CSRConstraint: return new def assign_labels( - self, cindex: int, name: str, scaling: np.ndarray | None = None + self, cindex: int, name: str, scaling: float | DataArray = 1.0 ) -> CSRConstraint: """ Return a copy labelled from ``cindex`` and named ``name``. Rows without terms are dropped, as when freezing a dense constraint; - a zero coefficient counts as a term. ``scaling`` is a row scaling over - the full flat grid. + a zero coefficient counts as a term. ``scaling`` is a scalar or a row + scaling broadcast on the grid; its distinct values are validated + without expanding a broadcast view. """ + values = np.asarray(scaling) + distinct = values[tuple(slice(None) if s else 0 for s in values.strides)] + validate_scaling(distinct, "constraint scaling") kept = self._kept(np.diff(self._csr.indptr) > 0) - kept._csr.eliminate_zeros() - changes: dict[str, Any] = dict(cindex=cindex, name=name) - if scaling is not None: - changes["scaling"] = scaling[kept._active_positions] - return kept._replace(**changes) + csr = kept._csr + if not csr.data.all(): + csr = csr.copy() if csr is self._csr else csr + csr.eliminate_zeros() + if isinstance(scaling, DataArray): + row_scaling = kept._active_values(scaling) + else: + row_scaling = np.full(csr.shape[0], float(scaling)) + return kept._replace(csr=csr, cindex=cindex, name=name, scaling=row_scaling) + + def _active_values(self, values: DataArray) -> np.ndarray: + """ + Values of ``values``, broadcast on the grid, at the active rows. + + A broadcast view is indexed per dimension instead of being expanded + to the full grid, unless the per-dimension indices would be larger. + """ + grid_values = values.transpose(*self._grid.dims).values + positions = self._active_positions + if ( + grid_values.flags.c_contiguous + or positions.size * grid_values.ndim >= grid_values.size + ): + return grid_values.reshape(-1)[positions] + return grid_values[np.unravel_index(positions, grid_values.shape)] def _kept(self, keep: np.ndarray) -> CSRConstraint: - """Copy holding only the active rows where ``keep`` is True.""" + """ + Copy holding only the active rows where ``keep`` is True, sharing the + row arrays when every row is kept. + """ + if keep.all(): + return self._replace() def rows(values: Any) -> Any: is_rows = isinstance(values, np.ndarray) and values.ndim @@ -978,8 +1007,7 @@ def masked(self, mask: DataArray) -> CSRConstraint: Copy with the cells where the boolean ``mask`` is False made inactive, without the dense rectangle. ``mask`` must lie on the constraint grid. """ - flat = mask.transpose(*self._grid.dims).to_numpy().reshape(-1) - return self._kept(flat[self._active_positions]) + return self._kept(self._active_values(mask).astype(bool)) def _assign_coords(self, **coords: Any) -> CSRConstraint: """ @@ -1643,10 +1671,11 @@ def from_csr( sign = maybe_replace_sign(sign) rhs_flat = _rhs_grid_values(expr, rhs) - expr.const active = np.flatnonzero(~np.isnan(rhs_flat)) + all_active = active.size == rhs_flat.size return cls( - expr.csr[active], + expr.csr if all_active else expr.csr[active], active, - rhs_flat[active], + rhs_flat if all_active else rhs_flat[active], sign, grid=expr.grid, model=expr.model, diff --git a/linopy/model.py b/linopy/model.py index 0c6b5c68..1e4f005b 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1371,18 +1371,20 @@ def add_constraints( con = self._constraint_from_lhs(lhs, sign, rhs, coords) if isinstance(con, CSRConstraint) and freeze: if mask is not None: - mask = broadcast_to_coords(mask, con.coords, label="mask") - con = con.masked(mask.astype(bool)) + con = con.masked(broadcast_to_coords(mask, con.coords, label="mask")) _check_infinities(con._sign, con._rhs, name) self.check_force_dim_names(con.coords.to_dataset()) enforce_no_multiindex(con, context=f"constraint {name!r}") - scaling_grid = validate_scaling( - broadcast_to_coords(scaling, con.coords, label="constraint scaling"), - "constraint scaling", + row_scaling = ( + float(scaling) + if isinstance(scaling, int | float | np.number) + else broadcast_to_coords( + scaling, con.coords, label="constraint scaling" + ) ) cindex = self._cCounter self._cCounter += con.full_size - con = con.assign_labels(cindex, name, scaling_grid.values.ravel()) + con = con.assign_labels(cindex, name, row_scaling) return self._soften_added(self.constraints.add(con), penalty) if isinstance(con, CSRConstraint): if chunked: diff --git a/linopy/scaling.py b/linopy/scaling.py index 3011ea07..8cb3aec6 100644 --- a/linopy/scaling.py +++ b/linopy/scaling.py @@ -8,7 +8,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeVar import numpy as np from xarray import DataArray, Dataset @@ -18,8 +18,10 @@ if TYPE_CHECKING: from linopy.model import Model +ScalingT = TypeVar("ScalingT", DataArray, np.ndarray) -def validate_scaling(scaling: DataArray, label: str = "scaling") -> DataArray: + +def validate_scaling(scaling: ScalingT, label: str = "scaling") -> ScalingT: """ Validate and normalize a scaling array. @@ -28,7 +30,7 @@ def validate_scaling(scaling: DataArray, label: str = "scaling") -> DataArray: multiply exported coefficients. They must be finite and strictly positive. """ scaling = scaling.astype(float) - values = scaling.values + values = np.asarray(scaling) if not np.isfinite(values).all() or (values <= 0).any(): raise ValueError(f"{label} must contain only finite positive values.") return scaling diff --git a/test/test_csr.py b/test/test_csr.py index 5189575b..1a371ec5 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -179,8 +179,10 @@ def test_zero_coefficient_rows_stay_active(sparse: bool) -> None: c = base_model(sparse=sparse) lhs = (0.0 * c.gen_p).groupby(c.gbus).sum() lhs = lhs + (0.0 * c.flow).groupby(c.bus0).sum() + nterm = lhs.nterm con = c.m.add_constraints(lhs == c.load, name="bal", freeze=True) assert len(con.active_labels()) == c.load.size + assert lhs.nterm == nterm def test_merge_keeps_absent_cell_absent() -> None: @@ -462,18 +464,72 @@ def test_frozen_invalid_infinite_rhs_raises(sparse: bool) -> None: c.m.add_constraints(c.balance_lhs() <= -np.inf, name="bal", freeze=True) +ROW_SCALINGS: dict[str, Callable[[xr.DataArray], Any]] = { + "scalar": lambda load: 2.0, + "snapshot": lambda load: xr.DataArray( + np.arange(1.0, load.sizes["snapshot"] + 1), coords=[load.indexes["snapshot"]] + ), + "grid": lambda load: load, +} + + +@pytest.mark.parametrize("masked", [False, True], ids=["all", "masked"]) +@pytest.mark.parametrize("scaling", list(ROW_SCALINGS)) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_constraint_applies_row_scaling(sparse: bool) -> None: +def test_frozen_constraint_applies_row_scaling( + sparse: bool, scaling: str, masked: bool +) -> None: require_v1() - c = base_model(sparse=sparse) - snaps = c.load.indexes["snapshot"] - scaling = xr.DataArray(np.arange(1.0, len(snaps) + 1), coords=[snaps]) + c = base_model(n_snap=6, sparse=sparse) + row_scaling = ROW_SCALINGS[scaling](c.load) + mask = c.load.bus == "bus1" if masked else None con = c.m.add_constraints( - c.balance_lhs() == c.load, name="bal", freeze=True, scaling=scaling + c.balance_lhs() == c.load, + name="bal", + freeze=True, + scaling=row_scaling, + mask=mask, ) assert isinstance(con, CSRConstraint) - expected = scaling.broadcast_like(con.scaling).transpose(*con.scaling.dims) - xr.testing.assert_equal(con.scaling, expected) + active = con.labels != -1 + expected = xr.DataArray(row_scaling).broadcast_like(c.load).astype(float) + xr.testing.assert_equal( + con.scaling.where(active), expected.transpose(*con.scaling.dims).where(active) + ) + + +@pytest.mark.parametrize("grid", [False, True], ids=["scalar", "grid"]) +@pytest.mark.parametrize("invalid", [-1.0, np.nan]) +@pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) +def test_frozen_constraint_rejects_invalid_row_scaling_on_masked_rows( + sparse: bool, invalid: float, grid: bool +) -> None: + require_v1() + c = base_model(sparse=sparse) + mask = c.load.bus == "bus1" + scaling = c.load.where(mask, invalid) if grid else invalid + with pytest.raises(ValueError, match="finite positive"): + c.m.add_constraints( + c.balance_lhs() == c.load, freeze=True, scaling=scaling, mask=mask + ) + + +@pytest.mark.parametrize("masked", [False, True], ids=["all", "masked"]) +def test_freeze_does_not_copy_lhs_matrix(masked: bool) -> None: + require_v1() + c = base_model(gens_per_bus=(400,) * 50, n_snap=20, sparse=True) + lhs = c.balance_lhs() + assert lhs._csr is not None + csr = lhs._csr.csr + csr_bytes = csr.data.nbytes + csr.indices.nbytes + csr.indptr.nbytes + mask = c.load.bus == "bus1" if masked else None + tracemalloc.start() + try: + c.m.add_constraints(lhs == c.load, name="bal", freeze=True, mask=mask) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert peak < csr_bytes / 4 def sparse_model_results(c: Case) -> dict[str, Any]: From 22ed32b7d1f22f8d4a3935a3704f12225bc5841c Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 13:23:46 +0200 Subject: [PATCH 2/2] docs(release-notes): note frozen CSR build memory savings (#977) --- doc/release_notes.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 98cbd965..ea9416ad 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -39,6 +39,7 @@ Upcoming Version * New read-only model key ``Model(sparse=True)`` turns on the whole sparse (CSR) path for one model: ``groupby(...).sum()`` and ``@``/``dot`` against a constant return CSR-backed expressions, and ``add_constraints`` freezes every constraint unless ``freeze=False`` is passed. It requires the v1 semantics: ``Model(sparse=True)`` raises under legacy, and so do ``groupby(...).sum()`` and ``@`` on a sparse model, and ``read_netcdf`` of one, once the semantics are switched back to legacy. It rejects ``chunk``. The key is kept by ``Model.copy`` and the netcdf round trip. ``@`` no longer reads ``linopy.options["sparse_groupby"]``. (`#976 `__) * Deprecated in favour of ``Model(sparse=True)``, each with a ``FutureWarning`` and to be removed with the legacy semantics: ``Model(freeze_constraints=...)`` and the ``Model.freeze_constraints`` setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. They keep their current behaviour until then, except that ``@`` ignores ``sparse_groupby``, and netcdf files that store ``freeze_constraints`` still load. (`#976 `__) +* Adding a frozen constraint from a sparse expression no longer copies the lhs matrix when every row stays active, and picks the mask and the row scaling at the active rows without expanding them over the full coordinate grid. This roughly halves the peak memory of ``add_constraints`` on a sparse model. (`#977 `__) *Other*