fix(constraints): CSRConstraint stores label columns; frozen constraints survive later variable changes

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 <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>`__)
 * ``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

@@ -531,15 +531,57 @@ 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.
 
     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
@@ -665,7 +707,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:
@@ -807,9 +849,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
@@ -912,10 +952,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."""
@@ -936,6 +990,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,
@@ -955,11 +1010,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
@@ -999,18 +1069,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
@@ -1086,12 +1160,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 = (
@@ -1149,8 +1222,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)
@@ -1724,9 +1801,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)
@@ -1741,14 +1818,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
@@ -1766,9 +1843,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
@@ -1792,9 +1869,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()
 
@@ -2410,9 +2487,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
         -------

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(

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")

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