perf(matrices): trim the Model.matrices build and drop redundant builds

benchmark/benchmark_sparse_export.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""
+Benchmark the build and export of a PyPSA-like dispatch model.
+
+Run as ``python benchmark/benchmark_sparse_export.py {sparse,dense,frozen}``.
+``sparse`` uses ``Model(sparse=True)``, ``dense`` keeps mutable constraints and
+``frozen`` freezes each constraint of a dense model. Every phase reports its
+wall time, its peak traced memory and the operations that densified. The
+``matrices``, ``to_highspy`` and ``to_file(lp)`` phases each build the matrices.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import tempfile
+import time
+import tracemalloc
+import warnings
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import linopy
+from linopy.constants import PerformanceWarning
+from linopy.io import to_highspy
+
+
+class Phases:
+    def __init__(self) -> None:
+        self.rows: list[tuple[str, float, float, list[str]]] = []
+
+    @contextmanager
+    def __call__(self, name: str) -> Iterator[None]:
+        gc.collect()
+        tracemalloc.start()
+        start = time.perf_counter()
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            yield
+        elapsed = time.perf_counter() - start
+        peak = tracemalloc.get_traced_memory()[1]
+        tracemalloc.stop()
+        densified = [
+            str(w.message)[:90]
+            for w in caught
+            if issubclass(w.category, PerformanceWarning)
+        ]
+        self.rows.append((name, elapsed, peak / 1e6, densified))
+
+    def report(self) -> None:
+        for name, elapsed, peak, densified in self.rows:
+            note = "DENSIFY: " + "; ".join(densified) if densified else ""
+            print(f"{name:28s} {elapsed * 1e3:9.0f} ms   peak {peak:8.1f} MB   {note}")
+        total = sum(row[1] for row in self.rows)
+        print(f"{'TOTAL':28s} {total * 1e3:9.0f} ms")
+
+
+def run(mode: str, n_bus: int, n_t: int, n_gen: int, n_line: int, n_sto: int) -> None:
+    linopy.options["semantics"] = "v1"
+    linopy.options["warn_on_densify"] = True
+    rng = np.random.default_rng(0)
+
+    bus = pd.RangeIndex(n_bus, name="bus")
+    t = pd.RangeIndex(n_t, name="t")
+    gen = pd.RangeIndex(n_gen, name="gen")
+    line = pd.RangeIndex(n_line, name="line")
+    sto = pd.RangeIndex(n_sto, name="sto")
+
+    gen_bus = rng.integers(0, n_bus, n_gen)
+    gen_bus[rng.random(n_gen) < 0.4] = 0
+    gen_bus = pd.Series(gen_bus, index=gen, name="bus")
+    sto_bus = pd.Series(np.arange(n_sto) % n_bus, index=sto, name="bus")
+    bus0 = np.arange(n_line) % n_bus
+    bus1 = (bus0 + 1 + rng.integers(0, 3, n_line)) % n_bus
+    incidence = np.zeros((n_line, n_bus))
+    incidence[np.arange(n_line), bus0] = -1
+    incidence[np.arange(n_line), bus1] = 1
+    incidence = xr.DataArray(incidence, coords=[line, bus])
+    demand = xr.DataArray(rng.uniform(10, 100, (n_bus, n_t)), coords=[bus, t])
+    pmax = xr.DataArray(rng.uniform(0, 1, (n_gen, n_t)), coords=[gen, t])
+
+    freeze = {"sparse": None, "dense": False, "frozen": True}[mode]
+    m = linopy.Model(sparse=mode == "sparse")
+    phase = Phases()
+
+    with phase("vars"):
+        p = m.add_variables(0, coords=[gen, t], name="p")
+        flow = m.add_variables(-100, 100, coords=[line, t], name="flow")
+        soc = m.add_variables(0, 100, coords=[sto, t], name="soc")
+        ch = m.add_variables(0, 50, coords=[sto, t], name="ch")
+        dis = m.add_variables(0, 50, coords=[sto, t], name="dis")
+    with phase("expr: groupby nodal supply"):
+        supply = p.groupby(gen_bus).sum()
+    with phase("expr: sto groupby"):
+        sto_net = (dis - ch).groupby(sto_bus).sum()
+    with phase("expr: flow @ incidence"):
+        net_flow = flow @ incidence
+    with phase("expr: balance sum"):
+        balance = supply + sto_net + net_flow
+    with phase("cons: balance"):
+        m.add_constraints(balance == demand, name="balance", freeze=freeze)
+    with phase("expr: soc shift"):
+        soc_expr = soc - 0.99 * soc.shift(t=1) - 0.95 * ch + dis / 0.95
+    with phase("cons: soc"):
+        m.add_constraints(soc_expr == 0, name="soc", freeze=freeze)
+    with phase("expr+cons: p <= pmax"):
+        m.add_constraints(p <= pmax * 100, name="p_max", freeze=freeze)
+    with phase("objective"):
+        m.add_objective((p * 2).sum() + (ch + dis).sum())
+    with phase("matrices"):
+        m.matrices.c
+    with phase("to_highspy"):
+        to_highspy(m)
+    with phase("to_file(lp)"), tempfile.TemporaryDirectory() as tmp:
+        m.to_file(Path(tmp) / "model.lp", progress=False)
+
+    print(
+        f"=== {mode}  buses={n_bus} snapshots={n_t} generators={n_gen} "
+        f"lines={n_line} storage={n_sto}  nvars={m.nvars} ncons={m.ncons}"
+    )
+    phase.report()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("mode", choices=["sparse", "dense", "frozen"])
+    parser.add_argument("--buses", type=int, default=200)
+    parser.add_argument("--snapshots", type=int, default=720)
+    parser.add_argument("--generators", type=int, default=2000)
+    parser.add_argument("--lines", type=int, default=300)
+    parser.add_argument("--storage", type=int, default=200)
+    args = parser.parse_args()
+    run(
+        args.mode, args.buses, args.snapshots, args.generators, args.lines, args.storage
+    )
+
+
+if __name__ == "__main__":
+    main()

doc/release_notes.rst

@@ -41,6 +41,7 @@ Upcoming Version
 * ``@``/``dot`` against a sparse constant no longer scales with the total number of variables in the model: each chunk of the sparse product now runs on only the variables it uses, which removes several seconds of allocation overhead on models with millions of variables. (`#990 <https://github.com/PyPSA/linopy/pull/990>`__)
 * 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>`__)
+* Building ``Model.matrices`` no longer allocates label-sized scaling lookups and skips ``eliminate_zeros`` for frozen constraints. On a model with 2M variables this makes each build of a sparse or frozen model about 20 to 35 ms faster. (`#1008 <https://github.com/PyPSA/linopy/issues/1008>`__)
 
 *Other*
 

linopy/dualization.py

@@ -479,12 +479,13 @@ def _add_dual_feasibility_constraints(
     dual_vars : dict
         ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``.
     """
-    A = m.matrices.A
+    M = m.matrices
+    A = M.A
     if A is None:
         raise ValueError("Constraint matrix is None, model has no constraints.")
 
-    vlabels = np.asarray(m.matrices.vlabels, dtype=np.int64)
-    clabels = np.asarray(m.matrices.clabels, dtype=np.int64)
+    vlabels = np.asarray(M.vlabels, dtype=np.int64)
+    clabels = np.asarray(M.clabels, dtype=np.int64)
 
     flat_con_to_dual = _build_flat_con_to_dual_label_lookup(m, dual_vars)
     if not len(flat_con_to_dual):
@@ -496,7 +497,7 @@ def _add_dual_feasibility_constraints(
     flat_v, flat_d, nnz_data = _extract_dual_feas_entries(
         A, vlabels, clabels, flat_con_to_dual
     )
-    c_lookup = _build_obj_coeff_lookup(vlabels, m.matrices.c)
+    c_lookup = _build_obj_coeff_lookup(vlabels, M.c)
 
     logger.debug("Building dual feasibility constraints for each primal variable.")
     for var_name, var in m.variables.items():

linopy/expressions.py

@@ -1743,11 +1743,15 @@ def _map_solution(self) -> DataArray:
         Replace variable labels by solution values.
         """
         m = self.model
-        M = m.matrices
-        sol = pd.Series(M.sol, M.vlabels)
+        labels = self.vars.values
+        known = np.append(m.variables.label_index.label_to_pos != -1, True)
+        if not known[labels].all():
+            raise KeyError("Expression references variables missing from the model.")
+        sol = np.full(m._xCounter + 1, np.nan)
+        for var in m.variables.data.values():
+            sol[var.labels.values] = var.solution.values
         sol[-1] = np.nan
-        idx = np.ravel(self.vars)
-        values = np.asarray(sol[idx]).reshape(self.vars.shape)
+        values = sol[labels]
         return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords)
 
     @property

linopy/matrices.py

@@ -15,35 +15,31 @@
 from numpy import ndarray
 
 from linopy import expressions
-from linopy.constraints import CSRConstraint
-from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup
+from linopy.constraints import ConstraintBase, CSRConstraint
 
 if TYPE_CHECKING:
     from linopy.model import Model
 
 
-def _stack(csrs: list) -> scipy.sparse.csr_array | None:
-    """
-    Vertically stack CSR blocks, or None when there are none.
-
-    Explicit zeros are dropped: expressions that broadcast against a dense
-    coordinate store one coefficient per pair, most of them zero, and a zero
-    coefficient never changes a constraint. Keeping them only inflates the
-    stored nnz handed to the solvers/writers (e.g. ``highspy.addRows`` scales
-    with stored nnz), so we prune them once, centrally, for every backend.
-    """
+def _stack(csrs: list[scipy.sparse.csr_array]) -> scipy.sparse.csr_array | None:
+    """Vertically stack CSR blocks, or None when there are none."""
     if not csrs:
         return None
-    stacked = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr"))
-    stacked.eliminate_zeros()
-    return stacked
+    return cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr"))
 
 
 def _concat(arrays: list, dtype: type | None = None) -> ndarray:
     """Concatenate arrays, or an empty array when there are none."""
     return np.concatenate(arrays) if arrays else np.array([], dtype=dtype)
 
 
+def _row_scaling(c: ConstraintBase) -> ndarray:
+    """Scaling of the active rows of a constraint, in row order."""
+    if isinstance(c, CSRConstraint):
+        return c._scaling
+    return c.scaling.values.ravel()[c.active_row_mask()]
+
+
 def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray:
     """Broadcast an indicator triggering value to one entry per active row."""
     if np.ndim(binval) == 0:
@@ -72,18 +68,12 @@ def __init__(self, model: Model) -> None:
 
     def _build_vars(self) -> None:
         m = self._parent
-        label_index = m.variables.label_index
-        self.vlabels: ndarray = label_index.vlabels
-        var_scaling_by_label = variable_scaling_lookup(m)
-        self.var_scaling: ndarray = (
-            var_scaling_by_label[self.vlabels]
-            if len(self.vlabels)
-            else np.array([], dtype=float)
-        )
+        self.vlabels: ndarray = m.variables.label_index.vlabels
 
         lb_list = []
         ub_list = []
         vtypes_list = []
+        scaling_list = []
 
         for name, var in m.variables.items():
             labels = var.labels.values.ravel()
@@ -101,7 +91,9 @@ def _build_vars(self) -> None:
             lb_list.append(var.lower.values.ravel()[mask])
             ub_list.append(var.upper.values.ravel()[mask])
             vtypes_list.append(np.full(mask.sum(), vtype))
+            scaling_list.append(var.solver_scaling.values.ravel()[mask])
 
+        self.var_scaling: ndarray = _concat(scaling_list, dtype=float)
         if lb_list:
             self.lb: ndarray = np.concatenate(lb_list) * self.var_scaling
             self.ub: ndarray = np.concatenate(ub_list) * self.var_scaling
@@ -115,15 +107,13 @@ def _build_cons(self) -> None:
         m = self._parent
         label_index = m.variables.label_index
         label_to_pos = label_index.label_to_pos
-        con_scaling_by_label = constraint_scaling_lookup(m)
         unit_cols = bool((self.var_scaling == 1).all())
 
         def scale_rows_and_cols(
-            csr: scipy.sparse.csr_array, con_labels: np.ndarray, b: np.ndarray
+            csr: scipy.sparse.csr_array, row_scaling: np.ndarray, b: np.ndarray
         ) -> tuple[scipy.sparse.csr_array, np.ndarray]:
             if csr.shape[0] == 0:
                 return csr, b
-            row_scaling = con_scaling_by_label[con_labels]
             unit_rows = bool((row_scaling == 1).all())
             if unit_rows and unit_cols:
                 return csr, b
@@ -144,17 +134,19 @@ def scale_rows_and_cols(
         for c in m.constraints.data.values():
             if c.is_indicator:
                 cc = c if isinstance(c, CSRConstraint) else c.freeze()
-                csr, con_labels, b, sense = cc.to_matrix_with_rhs(label_index)
-                csr, b = scale_rows_and_cols(csr, con_labels, b)
+                csr, _, b, sense = cc.to_matrix_with_rhs(label_index)
+                csr, b = scale_rows_and_cols(csr, cc._scaling, b)
                 ind_csrs.append(csr)
                 ind_b.append(b)
                 ind_sense.append(sense)
                 ind_binvar.append(label_to_pos[cc._binvar_labels])
                 binval = cast("int | np.ndarray", cc._binval)
                 ind_binval.append(_binval_per_row(binval, len(b)))
             else:
-                csr, con_labels, b, sense = c.to_matrix_with_rhs(label_index)
-                csr, b = scale_rows_and_cols(csr, con_labels, b)
+                csr, _, b, sense = c.to_matrix_with_rhs(label_index)
+                if not isinstance(c, CSRConstraint):
+                    csr.eliminate_zeros()
+                csr, b = scale_rows_and_cols(csr, _row_scaling(c), b)
                 reg_csrs.append(csr)
                 reg_b.append(b)
                 reg_sense.append(sense)

linopy/solvers.py

@@ -2799,10 +2799,8 @@ def _build_solver_model(
             rowind = np.empty(0, dtype=np.int64)
             rowcoef = np.empty(0, dtype=float)
 
-        lb = np.asarray(M.lb, dtype=float)
-        ub = np.asarray(M.ub, dtype=float)
-        np.place(lb, np.isneginf(lb), -xpress.infinity)
-        np.place(ub, np.isposinf(ub), xpress.infinity)
+        lb = np.where(np.isneginf(M.lb), -xpress.infinity, M.lb)
+        ub = np.where(np.isposinf(M.ub), xpress.infinity, M.ub)
 
         rowtype: np.ndarray
         rhs: np.ndarray

test/test_expressions.py

@@ -141,3 +141,19 @@ def test_expressions_solution() -> None:
     assert isinstance(sol, xr.Dataset)
     assert "double_x" in sol
     assert (sol["double_x"] == 4).all()
+
+
+def test_expression_solution_maps_variable_values() -> None:
+    m = Model()
+    y = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="y")
+    z = m.add_variables(coords=[pd.RangeIndex(2, name="j")], name="z")
+    m._mock_solve()
+    y.solution = xr.DataArray([10.0, 20.0, 30.0], coords=y.coords)
+    z.solution = xr.DataArray([1.0, 2.0], coords=z.coords)
+    expected = 2 * y.solution + z.solution
+    xr.testing.assert_equal((2 * y + z).solution, expected.rename("solution"))
+    expected = y.solution * z.solution
+    xr.testing.assert_equal((y * z).solution, expected.rename("solution"))
+    m.remove_variables("z")
+    with pytest.raises(KeyError):
+        (2 * y + z).solution

test/test_matrices.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 import xarray as xr
 
 from linopy import EQUAL, GREATER_EQUAL, Model
@@ -68,7 +69,8 @@ def test_matrices_duplicated_variables() -> None:
     assert np.isin(np.unique(np.array(A)), [0.0, 2.0]).all()
 
 
-def test_matrices_drops_explicit_zeros() -> None:
+@pytest.mark.parametrize("freeze", [False, True])
+def test_matrices_drops_explicit_zeros(freeze: bool) -> None:
     # https://github.com/PyPSA/linopy/issues/814
     # Expressions that broadcast against a dense coordinate store one coefficient
     # per pair, most of them structurally zero. Those must not reach A, whose
@@ -79,7 +81,9 @@ def test_matrices_drops_explicit_zeros() -> None:
     coeff = xr.DataArray(
         np.eye(4), dims=["j", "i"], coords={"j": range(4), "i": range(4)}
     )
-    m.add_constraints((coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c")
+    m.add_constraints(
+        (coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c", freeze=freeze
+    )
 
     A = m.matrices.A
     assert A is not None
@@ -98,3 +102,19 @@ def test_matrices_float_c() -> None:
 
     c = m.matrices.c
     assert np.all(c == np.array([1.5, 1.5]))
+
+
+def test_matrices_sol_aligned_with_vlabels() -> None:
+    m = Model()
+    i = pd.RangeIndex(3, name="i")
+    x = m.add_variables(coords=[i], name="x", mask=pd.Series([True, False, True], i))
+    y = m.add_variables(coords=[i], name="y")
+    m.add_constraints(x + y >= 0, name="c")
+    with pytest.raises(ValueError, match="not optimized"):
+        m.matrices.sol
+    m._mock_solve()
+    x.solution = xr.DataArray([1.0, np.nan, 3.0], coords=[i])
+    y.solution = xr.DataArray([4.0, 5.0, 6.0], coords=[i])
+    M = m.matrices
+    np.testing.assert_array_equal(M.vlabels, [0, 2, 3, 4, 5])
+    np.testing.assert_array_equal(M.sol, [1.0, 3.0, 4.0, 5.0, 6.0])