diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 49ac04c8..5bccdc84 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 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 793c6dba..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,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] diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index 7442b543..187a4bb0 100644 --- a/linopy/persistent/snapshot.py +++ b/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. """ diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index f86a5037..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,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: