perf(csr): avoid CSR copies and full-grid scaling when freezing constraints

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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/issues/977>`__)
 
 *Other*
 

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,

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:

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

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]: