perf(constraints): cache positional CSR on CSRConstraint so the snapshot diff keeps its identity fast path

doc/release_notes.rst

@@ -45,6 +45,8 @@ Upcoming Version
 * ``Model.copy`` (and the ``copy.copy``/``copy.deepcopy`` protocols) no longer downgrades a quadratic objective to a linear one. The objective was rebuilt as a ``LinearExpression`` regardless of its type, so the copy silently solved a different problem. ``linopy.testing.assert_model_equal`` now compares the objective by expression type as well, which it could not do before. (`#903 <https://github.com/PyPSA/linopy/issues/903>`__)
 * ``Model.to_netcdf``/``linopy.read_netcdf`` downgraded a quadratic objective the same way; the expression type is now stored alongside the objective and restored on read. Files written by earlier versions are read as before. (`#903 <https://github.com/PyPSA/linopy/issues/903>`__)
 * Adding or removing variables after a constraint was added with ``freeze=True`` no longer breaks ``model.matrices``. The frozen constraint stored dense variable positions as its matrix columns, so blocks frozen at different times disagreed on their width and stacking them raised ``ValueError: inconsistent shapes``. Raw variable labels are stored instead and mapped to positions when the matrix is assembled. Frozen constraints in netCDF files written by earlier versions are read as before. (`#926 <https://github.com/PyPSA/linopy/issues/926>`__)
+
+* A frozen constraint caches the label-to-position mapping of its matrix columns by weak reference and only rebuilds it when the constraint or the set of variables changes. While a persistent snapshot holds the arrays, repeated matrix assembly on an unchanged model returns the same objects, so the snapshot diff can again skip the comparison of untouched frozen constraints by object identity; one-off exports such as ``to_file`` retain no extra memory. (`#933 <https://github.com/PyPSA/linopy/issues/933>`__)
 * ``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 <https://github.com/PyPSA/linopy/pull/899>`__)
 
 **Breaking Changes**

linopy/constraints.py

@@ -8,6 +8,7 @@
 
 import functools
 import warnings
+import weakref
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Generator, Hashable, ItemsView, Iterator, Sequence
 from dataclasses import dataclass
@@ -559,18 +560,26 @@ def _csr_from_label_columns(
     index_dtype = (
         indptr.dtype if n_cols <= np.iinfo(indptr.dtype).max else np.dtype(np.int64)
     )
+    return _positional_csr(
+        data, cols.astype(index_dtype, copy=False), indptr, n_rows, n_cols
+    )
+
+
+def _positional_csr(
+    data: np.ndarray, cols: np.ndarray, indptr: np.ndarray, n_rows: int, n_cols: int
+) -> scipy.sparse.csr_array:
+    """Wrap positional CSR components without copying ``data`` or ``cols``."""
     csr = scipy.sparse.csr_array(
-        (
-            data,
-            cols.astype(index_dtype, copy=False),
-            indptr.astype(index_dtype, copy=False),
-        ),
-        shape=(n_rows, n_cols),
+        (data, cols, indptr.astype(cols.dtype, copy=False)), shape=(n_rows, n_cols)
     )
     csr.data = data
+    csr.indices = cols
     return csr
 
 
+_PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, "weakref.ref[np.ndarray]"]
+
+
 class CSRConstraint(ConstraintBase):
     """
     Frozen constraint backed by a CSR sparse matrix.
@@ -615,6 +624,7 @@ class CSRConstraint(ConstraintBase):
         "_dual",
         "_binvar_labels",
         "_binval",
+        "_positional_cache",
     )
 
     def __init__(
@@ -648,6 +658,7 @@ def __init__(
         self._dual = dual
         self._binvar_labels = binvar_labels
         self._binval = binval
+        self._positional_cache: _PositionalCache | None = None
 
     @property
     def model(self) -> Model:
@@ -718,9 +729,14 @@ def nterm(self) -> int:
     def coord_names(self) -> list[str]:
         return [str(c.name) for c in self._coords]
 
-    def _replace(self, **changes: Any) -> CSRConstraint:
-        """Copy with the given constructor arguments replaced."""
-        kwargs: dict[str, Any] = dict(
+    def __getstate__(self) -> dict[str, Any]:
+        return self._init_kwargs()
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__init__(**state)  # type: ignore[misc]
+
+    def _init_kwargs(self) -> dict[str, Any]:
+        return dict(
             csr=self._csr,
             active_positions=self._active_positions,
             rhs=self._rhs,
@@ -734,8 +750,15 @@ def _replace(self, **changes: Any) -> CSRConstraint:
             binval=self._binval,
             scaling=self._scaling,
         )
+
+    def _replace(self, **changes: Any) -> CSRConstraint:
+        """Copy with the given constructor arguments replaced."""
+        kwargs = self._init_kwargs()
         kwargs.update(changes)
-        return CSRConstraint(**kwargs)
+        new = CSRConstraint(**kwargs)
+        if kwargs["csr"] is self._csr:
+            new._positional_cache = self._positional_cache
+        return new
 
     def assign_labels(
         self, cindex: int, name: str, scaling: np.ndarray | None = None
@@ -1009,13 +1032,31 @@ def _to_positional_csr(
         """
         Return the stored CSR with label columns replaced by dense positions.
 
-        Only ``indices`` is freshly allocated; ``indptr`` and ``data`` stay the
-        stored arrays, so identity comparisons against a snapshot still hold.
+        ``indptr`` and ``data`` stay the stored arrays. The positional
+        ``indices`` are cached by weak reference, keyed on the identity of the
+        stored CSR and of ``label_index.label_to_pos``, which is rebuilt
+        whenever variables are added or removed. While a caller such as a
+        ``ModelSnapshot`` holds the array, repeated calls on an unchanged model
+        return the same object; otherwise it is freed and rebuilt on demand.
         """
         csr = self._csr
-        return _csr_from_label_columns(
+        label_to_pos = label_index.label_to_pos
+        if self._positional_cache is not None:
+            cached_csr, cached_label_to_pos, ref = self._positional_cache
+            cols = ref()
+            if (
+                cached_csr is csr
+                and cached_label_to_pos is label_to_pos
+                and cols is not None
+            ):
+                return _positional_csr(
+                    csr.data, cols, csr.indptr, csr.shape[0], label_index.n_active_vars
+                )
+        positional = _csr_from_label_columns(
             csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self._name
         )
+        self._positional_cache = (csr, label_to_pos, weakref.ref(positional.indices))
+        return positional
 
     def to_netcdf_ds(self) -> Dataset:
         """Return a Dataset with raw CSR components for netcdf serialization."""
@@ -1154,6 +1195,7 @@ def sanitize_zeros(self) -> CSRConstraint:
             csr.data[zeros] = 0
             csr.eliminate_zeros()
             self._csr = csr
+            self._positional_cache = None
         return self
 
     def sanitize_missings(self) -> CSRConstraint:
@@ -1177,6 +1219,7 @@ def sanitize_infinities(self) -> CSRConstraint:
             return self
         keep = ~invalid
         self._csr = self._csr[keep]
+        self._positional_cache = None
         self._active_positions = self._active_positions[keep]
         self._rhs = self._rhs[keep]
         self._scaling = self._scaling[keep]

linopy/persistent/snapshot.py

@@ -75,9 +75,10 @@ def _extract_con_buffers(
 
     Mutable ``Constraint`` objects build fresh arrays in
     ``to_matrix_with_rhs``, so the buffers are exclusively owned.
-    ``CSRConstraint`` returns a freshly gathered ``indices`` array, since its
-    label columns are mapped to dense positions on every call, and its stored
-    ``indptr``/``data`` — the latter share memory with the constraint, every
+    ``CSRConstraint`` returns its stored ``indptr``/``data`` and a weakly
+    cached positional ``indices`` array: as long as a snapshot holds it and
+    neither the constraint nor the variable label index changed, the same
+    array is returned. The buffers share memory with the constraint, every
     mutation path rebinds whole arrays (copy-on-write), and the diff uses
     object identity to skip comparisons on untouched containers.
     """

test/test_persistent_snapshot_buffers.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import pickle
+import weakref
+
 import numpy as np
 import pytest
 
@@ -123,6 +126,35 @@ def test_csr_capture_deterministic(baseline_model: Model) -> None:
         np.testing.assert_array_equal(b1.data, b2.data)
 
 
+def test_frozen_positional_indices_freed_without_holder(frozen_model: Model) -> None:
+    con = frozen_model.constraints["c2"]
+    label_index = frozen_model.variables.label_index
+    csr, _ = con.to_matrix(label_index)
+    ref = weakref.ref(csr.indices)
+    del csr
+    assert ref() is None
+    csr, _ = con.to_matrix(label_index)
+    np.testing.assert_array_equal(csr.indices, [0, 1, 2, 3, 4])
+
+
+def test_frozen_model_pickles_after_matrix_assembly(frozen_model: Model) -> None:
+    frozen_model.constraints.to_matrix()
+    restored = pickle.loads(pickle.dumps(frozen_model))
+    csr, _ = restored.constraints["c2"].to_matrix(restored.variables.label_index)
+    np.testing.assert_array_equal(csr.indices, [0, 1, 2, 3, 4])
+
+
+def test_frozen_capture_keeps_buffer_identity(frozen_model: Model) -> None:
+    s1 = ModelSnapshot.capture(frozen_model)
+    frozen_model.reset_solution()
+    s2 = ModelSnapshot.capture(frozen_model)
+    for name in s1.con_buffers:
+        b1, b2 = s1.con_buffers[name], s2.con_buffers[name]
+        assert b1.indptr is b2.indptr, name
+        assert b1.indices is b2.indices, name
+        assert b1.data is b2.data, name
+
+
 @pytest.fixture
 def frozen_model() -> Model:
     m = Model()
@@ -141,7 +173,36 @@ def test_frozen_con_buffers_keep_data_identity(frozen_model: Model) -> None:
         b2 = _extract_con_buffers(con, label_index)
         assert b1.data is b2.data, name
         assert b1.indptr is b2.indptr, name
-        np.testing.assert_array_equal(b1.indices, b2.indices)
+        assert b1.indices is b2.indices, name
+
+
+def test_frozen_con_positional_csr_follows_variable_changes() -> None:
+    m = Model()
+    m.add_variables(0, 10, coords=[range(3)], name="x")
+    y = m.add_variables(0, 5, coords=[range(2)], name="y")
+    m.add_constraints(y >= 1, name="c", freeze=True)
+    con = m.constraints["c"]
+    label_index = m.variables.label_index
+    before = _extract_con_buffers(con, label_index).indices
+    np.testing.assert_array_equal(before, [3, 4])
+    m.remove_variables("x")
+    after = _extract_con_buffers(con, label_index).indices
+    np.testing.assert_array_equal(after, [0, 1])
+    assert _extract_con_buffers(con, label_index).indices is after
+
+
+def test_frozen_con_positional_csr_follows_csr_rebind() -> None:
+    m = Model()
+    x = m.add_variables(0, 10, coords=[range(3)], name="x")
+    z = m.add_variables(0, 1, name="z")
+    m.add_constraints(x + 1e-12 * z >= 1, name="c", freeze=True)
+    con = m.constraints["c"]
+    label_index = m.variables.label_index
+    before = _extract_con_buffers(con, label_index)
+    con.sanitize_zeros()
+    after = _extract_con_buffers(con, label_index)
+    np.testing.assert_array_equal(before.indices, [0, 3, 1, 3, 2, 3])
+    np.testing.assert_array_equal(after.indices, [0, 1, 2])
 
 
 def test_untouched_frozen_constraint_needs_no_rebuild(frozen_model: Model) -> None: