From 14e1abe939373278faf8ec6b893a17a72dc7cbed Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 13:43:58 +0200 Subject: [PATCH 1/2] fix(constraints): CSRConstraint.indexes builds a valid xarray Indexes --- linopy/constraints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 0c887db9..94fed577 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -661,7 +661,7 @@ def sizes(self) -> Frozen[Hashable, int]: @property def indexes(self) -> Indexes: - return Indexes({c.name: c for c in self._coords}) + return Dataset(coords={c.name: c for c in self._coords}).indexes @property def nterm(self) -> int: From 9d671250fcdf6b9984cd2f16051d16a7392d0ba5 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 13:43:59 +0200 Subject: [PATCH 2/2] fix(constraints): CSRConstraint stores label columns; map to positions at matrix assembly (#926) --- doc/release_notes.rst | 1 + linopy/constraints.py | 139 ++++++++++++++++++----- linopy/persistent/snapshot.py | 9 +- test/test_constraint.py | 48 ++++++++ test/test_io.py | 30 +++++ test/test_persistent_snapshot_buffers.py | 28 +++++ 6 files changed, 220 insertions(+), 35 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index e3580e05..49ac04c8 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -44,6 +44,7 @@ Upcoming Version * The ``linopy.options`` context manager now restores the values that were active on entry instead of resetting all options to their defaults. * ``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 `__) * ``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 94fed577..e275be1f 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -527,6 +527,46 @@ def _equal_nnz_slices( start = end +def _csr_from_label_columns( + data: np.ndarray, + vlabel_cols: np.ndarray, + indptr: np.ndarray, + n_rows: int, + label_index: VariableLabelIndex, + name: str, +) -> scipy.sparse.csr_array: + """ + Build a CSR matrix from raw variable labels as column indices. + + Labels are mapped to dense positions in the active variable array, so the + result is ``(n_rows, label_index.n_active_vars)`` wide regardless of when + the label columns were assembled. ``data`` and, where its dtype is wide + enough, ``indptr`` are the very arrays passed in, so callers can compare + buffers by identity. + """ + n_cols = label_index.n_active_vars + cols = label_index.label_to_pos[vlabel_cols] + if cols.size and cols.min() < 0: + missing = vlabel_cols[cols < 0] + raise ValueError( + f"Constraint '{name}' references variable labels that are no longer " + f"part of the model, e.g. {missing[:5].tolist()}." + ) + index_dtype = ( + indptr.dtype if n_cols <= np.iinfo(indptr.dtype).max else np.dtype(np.int64) + ) + csr = scipy.sparse.csr_array( + ( + data, + cols.astype(index_dtype, copy=False), + indptr.astype(index_dtype, copy=False), + ), + shape=(n_rows, n_cols), + ) + csr.data = data + return csr + + class CSRConstraint(ConstraintBase): """ Frozen constraint backed by a CSR sparse matrix. @@ -534,8 +574,10 @@ class CSRConstraint(ConstraintBase): Parameters ---------- csr : scipy.sparse.csr_array - Shape (n_flat, model._xCounter). Each row is a flat position in the - constraint grid (including masked/empty rows). + Shape (n_active_cons, model._xCounter). Rows are the active rows of + the constraint grid, column indices are raw variable labels. They are + mapped to dense positions in ``to_matrix``/``to_matrix_with_rhs``, so + adding or removing variables after freezing stays consistent. rhs : np.ndarray Shape (n_flat,). Right-hand-side values. sign : str or np.ndarray @@ -803,9 +845,7 @@ def _to_dataset(self, nterm: int) -> Dataset: if csr.nnz > 0: row_indices = np.repeat(active_positions, counts) term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) - # csr.indices are column positions into vlabels; map back to variable labels - vlabels = self._model.variables.label_index.vlabels - vars_2d[row_indices, term_cols] = vlabels[csr.indices] + vars_2d[row_indices, term_cols] = csr.indices coeffs_2d[row_indices, term_cols] = csr.data dim_names = self.coord_names @@ -908,10 +948,24 @@ def row_expr(row: int) -> str: return "\n".join(lines) def to_matrix( - self, label_index: VariableLabelIndex | None = None + self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - """Return the stored CSR matrix and con_labels.""" - return self._csr, self._con_labels + """Return the CSR matrix with dense variable positions and con_labels.""" + return self._to_positional_csr(label_index), self._con_labels + + def _to_positional_csr( + self, label_index: VariableLabelIndex + ) -> scipy.sparse.csr_array: + """ + 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. + """ + csr = self._csr + return _csr_from_label_columns( + csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self._name + ) def to_netcdf_ds(self) -> Dataset: """Return a Dataset with raw CSR components for netcdf serialization.""" @@ -932,6 +986,7 @@ def to_netcdf_ds(self) -> Dataset: dim_names = [c.name for c in self._coords] attrs: dict[str, Any] = { "_linopy_format": "csr", + "_csr_columns": "labels", "cindex": self._cindex if self._cindex is not None else -1, "shape": list(csr.shape), "coord_dims": dim_names, @@ -951,11 +1006,26 @@ def to_netcdf_ds(self) -> Dataset: @classmethod def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: - """Reconstruct a Constraint from a netcdf Dataset (CSR format).""" + """ + Reconstruct a Constraint from a netcdf Dataset (CSR format). + + Files without the ``_csr_columns`` attribute were written before #926 + and hold dense variable positions instead of labels. + """ attrs = ds.attrs shape = tuple(attrs["shape"]) + indices = ds["indices"].values + columns = attrs.get("_csr_columns", "positions") + if columns == "positions": + vlabels = model.variables.label_index.vlabels + indices = vlabels[indices] + shape = (shape[0], int(vlabels[-1]) + 1 if len(vlabels) else 1) + elif columns != "labels": + raise ValueError( + f"Unknown CSR column format '{columns}' for constraint '{name}'." + ) csr = scipy.sparse.csr_array( - (ds["data"].values, ds["indices"].values, ds["indptr"].values), + (ds["data"].values, indices, ds["indptr"].values), shape=shape, ) rhs = ds["rhs"].values @@ -995,18 +1065,22 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: ) def has_labels(self, labels: np.ndarray) -> bool: - label_to_pos = self._model.variables.label_index.label_to_pos - return contains_labels(self._csr.indices, label_to_pos[labels]) + return contains_labels(self._csr.indices, labels) def to_matrix_with_rhs( self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: - """Return (csr, con_labels, b, sense) — all pre-stored, no recomputation.""" + """Return (csr, con_labels, b, sense); only the columns are recomputed.""" if isinstance(self._sign, str): sense = np.full(len(self._rhs), self._sign[0]) else: sense = np.array([s[0] for s in self._sign]) - return self._csr, self._con_labels, self._rhs, sense + return ( + self._to_positional_csr(label_index), + self._con_labels, + self._rhs, + sense, + ) def active_labels(self) -> np.ndarray: return self._con_labels @@ -1082,12 +1156,11 @@ def to_polars(self) -> pl.DataFrame: ) rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) - vlabels = self._model.variables.label_index.vlabels data: dict[str, Any] = { "labels": self._con_labels[rows], "coeffs": csr.data, - "vars": vlabels[csr.indices], + "vars": csr.indices, "rhs": self._rhs[rows], } sign_expr: pl.Expr | pl.Series = ( @@ -1145,8 +1218,12 @@ def from_mutable( cindex : int or None Starting label index, if assigned. """ - label_index = con.model.variables.label_index - csr, con_labels = con.to_matrix(label_index) + con_labels, _, vlabel_cols, data, indptr = con._matrix_export_data() + csr = scipy.sparse.csr_array( + (data, vlabel_cols, indptr), + shape=(len(con_labels), con.model._xCounter), + ) + csr.sum_duplicates() csr.eliminate_zeros() coords = [con.indexes[d] for d in con.coord_dims] # Build active_mask aligned with con_labels (rows in csr) @@ -1595,9 +1672,9 @@ def has_labels(self, labels: np.ndarray) -> bool: return contains_labels(self.data["vars"].values.ravel(), labels) def _matrix_export_data( - self, label_index: VariableLabelIndex + self, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - label_to_pos = label_index.label_to_pos + """Return (con_labels, row_mask, vlabel_cols, data, indptr) with raw labels.""" labels_flat = self.labels.values.ravel() vars_vals = self.vars.values n_rows = len(labels_flat) @@ -1612,14 +1689,14 @@ def _matrix_export_data( vars_final = vars_2d[row_mask] coeffs_final = self.coeffs.values.ravel().reshape(vars_2d.shape)[row_mask] valid_final = (vars_final != -1) & (coeffs_final != 0) - cols = label_to_pos[vars_final[valid_final]] + vlabel_cols = vars_final[valid_final] data = coeffs_final[valid_final] counts = valid_final.sum(axis=1) indptr = np.empty(len(con_labels) + 1, dtype=np.int32) indptr[0] = 0 np.cumsum(counts, out=indptr[1:]) - return con_labels, row_mask, cols, data, indptr + return con_labels, row_mask, vlabel_cols, data, indptr def to_matrix( self, label_index: VariableLabelIndex @@ -1637,9 +1714,9 @@ def to_matrix( con_labels : np.ndarray Active constraint labels in row order. """ - con_labels, _, cols, data, indptr = self._matrix_export_data(label_index) - csr = scipy.sparse.csr_array( - (data, cols, indptr), shape=(len(con_labels), label_index.n_active_vars) + con_labels, _, vlabel_cols, data, indptr = self._matrix_export_data() + csr = _csr_from_label_columns( + data, vlabel_cols, indptr, len(con_labels), label_index, self.name ) csr.sum_duplicates() return csr, con_labels @@ -1663,9 +1740,9 @@ def to_matrix_with_rhs( self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: """Return (csr, con_labels, b, sense) in one pass.""" - con_labels, row_mask, cols, data, indptr = self._matrix_export_data(label_index) - csr = scipy.sparse.csr_array( - (data, cols, indptr), shape=(len(con_labels), label_index.n_active_vars) + con_labels, row_mask, vlabel_cols, data, indptr = self._matrix_export_data() + csr = _csr_from_label_columns( + data, vlabel_cols, indptr, len(con_labels), label_index, self.name ) csr.sum_duplicates() @@ -2281,9 +2358,9 @@ def to_matrix(self) -> tuple[scipy.sparse.csr_array, np.ndarray]: """ Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. - Each per-constraint CSR is already dense: rows are active constraints - only, column indices are dense variable positions (not raw labels). - Shape is ``(n_active_cons, n_active_vars)``. + Each per-constraint CSR has active constraint rows only and dense + variable positions as column indices. Shape is + ``(n_active_cons, n_active_vars)``. Returns ------- diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index fd758ea3..7442b543 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -75,10 +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 its stored arrays — 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 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. """ csr, con_labels, b, sense = con.to_matrix_with_rhs(var_label_index) return ContainerConBuffers( diff --git a/test/test_constraint.py b/test/test_constraint.py index 7fe50a57..aac1cedc 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -902,6 +902,54 @@ def test_freeze_mutable_roundtrip(m: Model) -> None: np.testing.assert_array_equal(frozen._con_labels, refrozen._con_labels) +def test_frozen_csr_stores_variable_labels(m: Model, x: linopy.Variable) -> None: + frozen = m.constraints["c"] + assert isinstance(frozen, linopy.constraints.CSRConstraint) + np.testing.assert_array_equal( + np.unique(frozen._csr.indices), np.unique(x.labels.values) + ) + assert frozen._csr.shape[1] == m._xCounter + csr, _ = frozen.to_matrix(m.variables.label_index) + assert csr.shape[1] == m.variables.label_index.n_active_vars + + +def _model_with_frozen_and_mutation(freeze: bool, mutation: str) -> Model: + m = Model() + i = pd.RangeIndex(3, name="i") + m.add_variables(coords=[i], name="a") + b = m.add_variables(coords=[i], name="b") + m.add_constraints(2 * b >= 1, name="c1", freeze=freeze) + d = m.add_variables(coords=[pd.RangeIndex(2, name="i")], name="d") + m.add_constraints(d <= 5, name="c2", freeze=freeze) + if mutation == "remove": + m.remove_variables("a") + return m + + +@pytest.mark.parametrize("mutation", ["add", "remove"]) +def test_frozen_matrices_after_variable_mutation(mutation: str) -> None: + frozen = _model_with_frozen_and_mutation(True, mutation) + mutable = _model_with_frozen_and_mutation(False, mutation) + np.testing.assert_array_equal( + frozen.matrices.A.toarray(), mutable.matrices.A.toarray() + ) + np.testing.assert_array_equal(frozen.matrices.b, mutable.matrices.b) + np.testing.assert_array_equal(frozen.matrices.clabels, mutable.matrices.clabels) + np.testing.assert_array_equal(frozen.matrices.vlabels, mutable.matrices.vlabels) + + +@pytest.mark.parametrize("freeze", [True, False]) +def test_constraint_removed_with_referenced_variable(freeze: bool) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + a = m.add_variables(coords=[i], name="a") + b = m.add_variables(coords=[i], name="b") + m.add_constraints(a + b >= 1, name="c", freeze=freeze) + with pytest.warns(UserWarning, match="also removes constraints"): + m.remove_variables("a") + assert "c" not in m.constraints + + def test_freeze_mutable_roundtrip_with_masking() -> None: m = Model() x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") diff --git a/test/test_io.py b/test/test_io.py index 44169eaf..d7e8904c 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -137,6 +137,36 @@ def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: assert_model_equal(m, p) +def test_model_from_netcdf_frozen_constraint_legacy_positions(tmp_path: Path) -> None: + """Files written before #926 stored dense positions as CSR columns.""" + from linopy.constraints import CSRConstraint + + m = Model() + i = pd.RangeIndex(3, name="i") + mask = xr.DataArray([True, False, False], dims=["i"]) + m.add_variables(coords=[i], name="a", mask=mask) + b = m.add_variables(coords=[i], name="b") + m.add_constraints(2 * b >= 1, name="c", freeze=True) + + fn = tmp_path / "test_frozen_legacy.nc" + m.to_netcdf(fn) + + ds = xr.load_dataset(fn) + label_to_pos = m.variables.label_index.label_to_pos + labels = ds["constraints-c-indices"].values + positions = label_to_pos[labels] + assert not np.array_equal(labels, positions) + ds["constraints-c-indices"] = xr.DataArray(positions, dims=["constraints-c-_nnz"]) + del ds.attrs["constraints-c-_csr_columns"] + legacy_fn = tmp_path / "legacy.nc" + ds.to_netcdf(legacy_fn) + + p = read_netcdf(legacy_fn) + assert isinstance(p.constraints["c"], CSRConstraint) + np.testing.assert_array_equal(p.matrices.A.toarray(), m.matrices.A.toarray()) + assert_model_equal(m, p) + + def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None: from linopy.constraints import CSRConstraint diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index 76bfa7e6..f86a5037 100644 --- a/test/test_persistent_snapshot_buffers.py +++ b/test/test_persistent_snapshot_buffers.py @@ -123,6 +123,34 @@ def test_csr_capture_deterministic(baseline_model: Model) -> None: np.testing.assert_array_equal(b1.data, b2.data) +@pytest.fixture +def frozen_model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 5, coords=[range(2)], name="y") + m.add_constraints(2 * x >= 4, name="c1", freeze=True) + m.add_constraints(x.sum() + y.sum() <= 20, name="c2", freeze=True) + m.add_objective(x.sum()) + return m + + +def test_frozen_con_buffers_keep_data_identity(frozen_model: Model) -> None: + label_index = frozen_model.variables.label_index + for name, con in frozen_model.constraints.items(): + b1 = _extract_con_buffers(con, label_index) + 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) + + +def test_untouched_frozen_constraint_needs_no_rebuild(frozen_model: Model) -> None: + snap = ModelSnapshot.capture(frozen_model) + diff = ModelDiff.from_snapshot(snap, frozen_model) + assert isinstance(diff, ModelDiff) + assert diff.is_empty + + def test_duplicate_variable_terms_summed() -> None: m1 = Model() x1 = m1.add_variables(0, 10, coords=[range(3)], name="x")