From 715830bca26118565c516fa4ea61d37bc6753134 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 08:08:23 +0200 Subject: [PATCH 1/4] perf(constraints): cache positional CSR on CSRConstraint keyed on label index identity (#933) --- linopy/constraints.py | 17 ++++++++++++++--- linopy/persistent/snapshot.py | 10 +++++----- test/test_persistent_snapshot_buffers.py | 15 ++++++++++++++- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 793c6dba..564362f9 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -648,6 +648,9 @@ def __init__( self._dual = dual self._binvar_labels = binvar_labels self._binval = binval + self._positional_cache: ( + tuple[scipy.sparse.csr_array, np.ndarray, scipy.sparse.csr_array] | None + ) = None @property def model(self) -> Model: @@ -1009,13 +1012,21 @@ 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 result is cached + on the identity of the stored CSR and of ``label_index.label_to_pos``, + which is rebuilt whenever variables are added or removed, so repeated + calls on an unchanged model also return the same ``indices`` array. """ csr = self._csr - return _csr_from_label_columns( + label_to_pos = label_index.label_to_pos + cache = self._positional_cache + if cache is not None and cache[0] is csr and cache[1] is label_to_pos: + return cache[2] + 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, positional) + return positional def to_netcdf_ds(self) -> Dataset: """Return a Dataset with raw CSR components for netcdf serialization.""" diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index 7442b543..c87c648f 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -75,11 +75,11 @@ 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 - mutation path rebinds whole arrays (copy-on-write), and the diff uses - object identity to skip comparisons on untouched containers. + ``CSRConstraint`` returns its stored ``indptr``/``data`` and a cached + positional ``indices`` array that is only rebuilt when the constraint or + the variable label index changes. 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. """ csr, con_labels, b, sense = con.to_matrix_with_rhs(var_label_index) return ContainerConBuffers( diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index f86a5037..3af1ac52 100644 --- a/test/test_persistent_snapshot_buffers.py +++ b/test/test_persistent_snapshot_buffers.py @@ -141,7 +141,20 @@ 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( + frozen_model: Model, +) -> None: + con = frozen_model.constraints["c2"] + label_index = frozen_model.variables.label_index + before = _extract_con_buffers(con, label_index).indices + frozen_model.add_variables(0, 1, coords=[range(2)], name="z") + after = _extract_con_buffers(con, label_index).indices + assert after is not before + np.testing.assert_array_equal(after, before) + assert _extract_con_buffers(con, label_index).indices is after def test_untouched_frozen_constraint_needs_no_rebuild(frozen_model: Model) -> None: From d723f0229a3f7b0881edb5653b20d98af602e4a9 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 08:17:02 +0200 Subject: [PATCH 2/4] perf(constraints): keep positional CSR cache across _replace, clear it on csr rebind; add slot and tests (#933) --- linopy/constraints.py | 22 ++++++++---- test/test_persistent_snapshot_buffers.py | 43 +++++++++++++++++++----- 2 files changed, 50 insertions(+), 15 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 564362f9..3da58320 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -571,6 +571,9 @@ def _csr_from_label_columns( return csr +_PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, scipy.sparse.csr_array] + + class CSRConstraint(ConstraintBase): """ Frozen constraint backed by a CSR sparse matrix. @@ -615,6 +618,7 @@ class CSRConstraint(ConstraintBase): "_dual", "_binvar_labels", "_binval", + "_positional_cache", ) def __init__( @@ -648,9 +652,7 @@ def __init__( self._dual = dual self._binvar_labels = binvar_labels self._binval = binval - self._positional_cache: ( - tuple[scipy.sparse.csr_array, np.ndarray, scipy.sparse.csr_array] | None - ) = None + self._positional_cache: _PositionalCache | None = None @property def model(self) -> Model: @@ -738,7 +740,10 @@ def _replace(self, **changes: Any) -> CSRConstraint: scaling=self._scaling, ) 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 @@ -1019,9 +1024,10 @@ def _to_positional_csr( """ csr = self._csr label_to_pos = label_index.label_to_pos - cache = self._positional_cache - if cache is not None and cache[0] is csr and cache[1] is label_to_pos: - return cache[2] + if self._positional_cache is not None: + cached_csr, cached_label_to_pos, positional = self._positional_cache + if cached_csr is csr and cached_label_to_pos is label_to_pos: + return positional positional = _csr_from_label_columns( csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self._name ) @@ -1165,6 +1171,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: @@ -1188,6 +1195,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] diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index 3af1ac52..5bda83b8 100644 --- a/test/test_persistent_snapshot_buffers.py +++ b/test/test_persistent_snapshot_buffers.py @@ -123,6 +123,17 @@ def test_csr_capture_deterministic(baseline_model: Model) -> None: np.testing.assert_array_equal(b1.data, b2.data) +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() @@ -144,19 +155,35 @@ def test_frozen_con_buffers_keep_data_identity(frozen_model: Model) -> None: assert b1.indices is b2.indices, name -def test_frozen_con_positional_csr_follows_variable_changes( - frozen_model: Model, -) -> None: - con = frozen_model.constraints["c2"] - label_index = frozen_model.variables.label_index +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 - frozen_model.add_variables(0, 1, coords=[range(2)], name="z") + np.testing.assert_array_equal(before, [3, 4]) + m.remove_variables("x") after = _extract_con_buffers(con, label_index).indices - assert after is not before - np.testing.assert_array_equal(after, before) + 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: snap = ModelSnapshot.capture(frozen_model) diff = ModelDiff.from_snapshot(snap, frozen_model) From 6921fac93c105083552bed61580bfff41ae54ef8 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 08:17:54 +0200 Subject: [PATCH 3/4] doc: release note for the positional CSR cache (#933) --- doc/release_notes.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 49ac04c8..d1f065f7 100644 --- a/doc/release_notes.rst +++ b/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 `__) * ``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 `__) * 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 `__) + +* A frozen constraint caches the label-to-position mapping of its matrix columns and only rebuilds it when the constraint or the set of variables changes. Repeated matrix assembly on an unchanged model returns the same arrays, so the persistent snapshot diff can again skip the comparison of untouched frozen constraints by object identity. (`#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 `__) **Breaking Changes** From 8f38f7e848b27b25484b14786194afed8f4e7ef3 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 08:40:51 +0200 Subject: [PATCH 4/4] perf(constraints): hold the positional CSR cache by weak reference; keep CSRConstraint picklable (#933) --- doc/release_notes.rst | 2 +- linopy/constraints.py | 60 +++++++++++++++++------- linopy/persistent/snapshot.py | 11 +++-- test/test_persistent_snapshot_buffers.py | 21 +++++++++ 4 files changed, 70 insertions(+), 24 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d1f065f7..5bccdc84 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -46,7 +46,7 @@ Upcoming Version * ``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 `__) * 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 `__) -* A frozen constraint caches the label-to-position mapping of its matrix columns and only rebuilds it when the constraint or the set of variables changes. Repeated matrix assembly on an unchanged model returns the same arrays, so the persistent snapshot diff can again skip the comparison of untouched frozen constraints by object identity. (`#933 `__) +* 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 `__) * ``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 `__) **Breaking Changes** diff --git a/linopy/constraints.py b/linopy/constraints.py index 3da58320..bcc974f6 100644 --- a/linopy/constraints.py +++ b/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,19 +560,24 @@ 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, scipy.sparse.csr_array] +_PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, "weakref.ref[np.ndarray]"] class CSRConstraint(ConstraintBase): @@ -723,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, @@ -739,6 +750,10 @@ 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) new = CSRConstraint(**kwargs) if kwargs["csr"] is self._csr: @@ -1017,21 +1032,30 @@ def _to_positional_csr( """ Return the stored CSR with label columns replaced by dense positions. - ``indptr`` and ``data`` stay the stored arrays. The result is cached - on the identity of the stored CSR and of ``label_index.label_to_pos``, - which is rebuilt whenever variables are added or removed, so repeated - calls on an unchanged model also return the same ``indices`` array. + ``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 label_to_pos = label_index.label_to_pos if self._positional_cache is not None: - cached_csr, cached_label_to_pos, positional = self._positional_cache - if cached_csr is csr and cached_label_to_pos is label_to_pos: - return positional + 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, positional) + self._positional_cache = (csr, label_to_pos, weakref.ref(positional.indices)) return positional def to_netcdf_ds(self) -> Dataset: diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index c87c648f..187a4bb0 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -75,11 +75,12 @@ def _extract_con_buffers( Mutable ``Constraint`` objects build fresh arrays in ``to_matrix_with_rhs``, so the buffers are exclusively owned. - ``CSRConstraint`` returns its stored ``indptr``/``data`` and a cached - positional ``indices`` array that is only rebuilt when the constraint or - the variable label index changes. 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. + ``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. """ csr, con_labels, b, sense = con.to_matrix_with_rhs(var_label_index) return ContainerConBuffers( diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index 5bda83b8..fcacb1b8 100644 --- a/test/test_persistent_snapshot_buffers.py +++ b/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,24 @@ 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()