fix: support adding/removing variables after adding frozen constraint

doc/release_notes.rst

@@ -31,6 +31,8 @@ Upcoming Version
 
 * LP file export now honors bounds tightened below ``[0, 1]`` on a binary variable via the ``.lower``/``.upper`` setters after creation (e.g. ``upper = 0``). Previously such bounds were written only by ``io_api="direct"`` and dropped by ``io_api="lp"``. (https://github.com/PyPSA/linopy/issues/776)
 * Freezing an empty constraint group (e.g. an empty ``isel`` slice) no longer raises ``ValueError: cannot reshape array of size 0``. ``Model(freeze_constraints=True)`` and ``Constraint.freeze()`` now round-trip zero-row constraints losslessly.
+* Adding variables after freezing a constraint (e.g. under ``Model(freeze_constraints=True)`` while building incrementally) no longer fails with ``ValueError: incompatible dimensions for axis 1`` when assembling the constraint matrix. Frozen (CSR-backed) constraints now widen their cached CSR to the current active-variable count, matching the ``shape (n_active_cons, n_active_vars)`` invariant of the mutable path.
+* Removing a variable after freezing constraints now remaps the cached CSR columns of the surviving frozen constraints. Previously ``Model.remove_variables`` left their absolute variable positions stale after the removal renumbered them, producing a wrong-width constraint matrix or an ``IndexError`` at solve time.
 * ``Variable.where`` no longer raises ``ValueError: exact match required for all data variable names`` once a solution is attached (after ``Model.solve``) or the variable is fixed. The fill value now covers auxiliary data variables (``solution``, stashed bounds) instead of only ``labels``/``lower``/``upper``.
 
 Version 0.8.0

linopy/constraints.py

@@ -852,11 +852,73 @@ def row_expr(row: int) -> str:
 
         return "\n".join(lines)
 
+    def _csr_for_index(
+        self, label_index: VariableLabelIndex | None
+    ) -> scipy.sparse.csr_array:
+        """
+        Return the cached CSR widened to the current active-variable count.
+
+        The CSR is cached at freeze time with as many columns as there were
+        active variables then. If variables were added afterwards (e.g. an
+        auxiliary variable introduced by a constraint added after this one was
+        frozen), the active-variable count has grown and the cached CSR is too
+        narrow to be stacked with constraints frozen later, so
+        ``Constraints.to_matrix`` fails with an axis-1 dimension mismatch.
+
+        Variable positions are assigned in encounter order and are stable under
+        additions, so widening the column dimension only appends empty trailing
+        columns and keeps every per-constraint block consistent with
+        ``label_index.n_active_vars``.
+        """
+        csr = self._csr
+        if label_index is None:
+            return csr
+        n_active_vars = label_index.n_active_vars
+        if csr.shape[1] < n_active_vars:
+            csr = scipy.sparse.csr_array(
+                (csr.data, csr.indices, csr.indptr),
+                shape=(csr.shape[0], n_active_vars),
+            )
+        return csr
+
+    def _remap_columns(
+        self, old_vlabels: np.ndarray, new_label_index: VariableLabelIndex
+    ) -> None:
+        """
+        Rewrite cached CSR column positions after variables were removed.
+
+        The cached CSR stores absolute dense variable positions from freeze
+        time. Removing a variable renumbers the positions of every later
+        variable, so those cached positions go stale (wrong columns and wrong
+        width). ``old_vlabels`` gives the variable label at each column position
+        currently used by the CSR; ``new_label_index`` provides the post-removal
+        label -> position mapping. A frozen constraint that referenced a removed
+        variable has already been dropped, so every remaining column maps to a
+        valid new position.
+        """
+        csr = self._csr
+        n_active_vars = new_label_index.n_active_vars
+        if csr.nnz == 0:
+            self._csr = scipy.sparse.csr_array(
+                (csr.shape[0], n_active_vars), dtype=csr.dtype
+            )
+            return
+        new_positions = new_label_index.label_to_pos[old_vlabels[csr.indices]]
+        if (new_positions < 0).any():
+            raise ValueError(
+                "Frozen constraint references a removed variable while remapping "
+                "columns; this should not happen."
+            )
+        self._csr = scipy.sparse.csr_array(
+            (csr.data, new_positions, csr.indptr),
+            shape=(csr.shape[0], n_active_vars),
+        )
+
     def to_matrix(
         self, label_index: VariableLabelIndex | None = None
     ) -> tuple[scipy.sparse.csr_array, np.ndarray]:
         """Return the stored CSR matrix and con_labels."""
-        return self._csr, self._con_labels
+        return self._csr_for_index(label_index), self._con_labels
 
     def to_netcdf_ds(self) -> Dataset:
         """Return a Dataset with raw CSR components for netcdf serialization."""
@@ -945,12 +1007,17 @@ def has_variable(self, variable: variables.Variable) -> bool:
     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) — pre-stored, widened if needed.
+
+        See ``_csr_for_index`` for why the cached CSR may need widening to the
+        current active-variable count before it can be stacked.
+        """
         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._csr_for_index(label_index), self._con_labels, self._rhs, sense
 
     def active_labels(self) -> np.ndarray:
         return self._con_labels

linopy/model.py

@@ -1281,8 +1281,23 @@ def remove_variables(self, name: str) -> None:
             for k in to_remove:
                 self.constraints.remove(k)
 
+        # Frozen (CSR-backed) constraints cache absolute variable positions.
+        # Removing a variable renumbers the positions of every later variable,
+        # so capture the pre-removal ordering to remap the surviving frozen
+        # constraints once the variable is gone.
+        frozen_cons = [
+            c for _, c in self.constraints.items() if isinstance(c, CSRConstraint)
+        ]
+        old_vlabels = self.variables.label_index.vlabels.copy() if frozen_cons else None
+
         self.variables.remove(name)
 
+        if frozen_cons:
+            assert old_vlabels is not None  # captured whenever frozen_cons is non-empty
+            new_label_index = self.variables.label_index
+            for con in frozen_cons:
+                con._remap_columns(old_vlabels, new_label_index)
+
         self.objective = self.objective.sel(
             {TERM_DIM: ~self.objective.vars.isin(variable.labels)}
         )

test/test_frozen_constraint_active_vars.py

@@ -0,0 +1,267 @@
+"""
+Differential tests for frozen (CSR-backed) constraints: changing the variables.
+
+The invariant under test (stated in PR #630): ``to_matrix`` on both the mutable
+``Constraint`` and the frozen ``CSRConstraint`` "always returns a csr matrix of
+shape (n_active_constraints, n_active_variables)". A frozen constraint caches
+its CSR at freeze time, so any change to the active-variable set *after* freezing
+(adding or removing variables) must still yield a matrix consistent with the
+mutable path.
+
+Strategy: build the same model twice under an identical construction order --
+once with ``freeze_constraints=False`` (mutable, treated as ground truth) and
+once with ``freeze_constraints=True`` -- and assert the assembled constraint
+matrix, shape, and solved objective agree. Any divergence is a frozen-path bug.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import linopy
+from linopy import Model
+
+
+def _assemble(freeze: bool) -> Model:
+    """
+    Return a model built in a fixed order that freezes before adding a variable.
+
+    Order: add x, add a constraint on x, THEN add y, add a constraint on y.
+    With freeze_constraints=True the first constraint is frozen while only x
+    exists.
+    """
+    m = Model(freeze_constraints=freeze)
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x")
+    m.add_constraints(x >= 1, name="cx")
+    y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="j")], name="y")
+    m.add_constraints(y >= 2, name="cy")
+    m.add_objective(x.sum() + y.sum())
+    return m
+
+
+def _dense_A(m: Model) -> np.ndarray:
+    A = m.matrices.A
+    assert A is not None
+    return np.asarray(A.todense())
+
+
+# ---------------------------------------------------------------------------
+# Scenario A: baseline -- all variables added before any frozen constraint.
+# This is what upstream fixtures do; should pass with or without the fix.
+# ---------------------------------------------------------------------------
+def test_A_all_vars_first() -> None:
+    m_mut = Model()
+    x = m_mut.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x")
+    y = m_mut.add_variables(lower=0, coords=[pd.RangeIndex(2, name="j")], name="y")
+    m_mut.add_constraints(x >= 1, name="cx", freeze=False)
+    m_mut.add_constraints(y >= 2, name="cy", freeze=False)
+
+    m_frz = Model()
+    x = m_frz.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x")
+    y = m_frz.add_variables(lower=0, coords=[pd.RangeIndex(2, name="j")], name="y")
+    m_frz.add_constraints(x >= 1, name="cx", freeze=True)
+    m_frz.add_constraints(y >= 2, name="cy", freeze=True)
+
+    np.testing.assert_array_equal(_dense_A(m_mut), _dense_A(m_frz))
+
+
+# ---------------------------------------------------------------------------
+# Scenario B: interleaved -- variable added AFTER a frozen constraint.
+# This is the widening case. Expected to FAIL on unpatched master.
+# ---------------------------------------------------------------------------
+def test_B_var_added_after_freeze_matrix_shape() -> None:
+    m = _assemble(freeze=True)
+    n_vars = m.variables.label_index.n_active_vars
+    n_cons = m.constraints.label_index.n_active_cons
+    A = m.matrices.A
+    assert A is not None
+    assert A.shape == (n_cons, n_vars)
+
+
+def test_B_interleaved_matrix_matches_mutable() -> None:
+    A_mut = _dense_A(_assemble(freeze=False))
+    A_frz = _dense_A(_assemble(freeze=True))
+    assert A_mut.shape == A_frz.shape
+    np.testing.assert_array_equal(A_mut, A_frz)
+
+
+def test_B_interleaved_solves_equal() -> None:
+    if "highs" not in linopy.available_solvers:
+        pytest.skip("highs unavailable")
+    m_mut = _assemble(freeze=False)
+    m_frz = _assemble(freeze=True)
+    m_mut.solve(solver_name="highs")
+    m_frz.solve(solver_name="highs")
+    obj_mut = m_mut.objective.value
+    obj_frz = m_frz.objective.value
+    assert obj_mut is not None and obj_frz is not None
+    assert obj_frz == pytest.approx(obj_mut)
+
+
+# ---------------------------------------------------------------------------
+# Scenario C: a frozen constraint that links a variable added later, plus an
+# extra late variable that widens the index further.
+# ---------------------------------------------------------------------------
+def _assemble_linking(freeze: bool) -> Model:
+    m = Model(freeze_constraints=freeze)
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x")
+    m.add_constraints(x >= 1, name="cx")
+    y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y")
+    # constraint linking x and y, frozen while both exist
+    m.add_constraints(x + y >= 5, name="cxy")
+    z = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="z")
+    m.add_constraints(z >= 3, name="cz")
+    m.add_objective(x.sum() + y.sum() + z.sum())
+    return m
+
+
+def test_C_linking_matrix_matches_mutable() -> None:
+    A_mut = _dense_A(_assemble_linking(freeze=False))
+    A_frz = _dense_A(_assemble_linking(freeze=True))
+    assert A_mut.shape == A_frz.shape
+    np.testing.assert_array_equal(A_mut, A_frz)
+
+
+# ---------------------------------------------------------------------------
+# Scenario D: variable removal AFTER freezing. Removing a variable renumbers
+# dense positions for all later variables. A frozen constraint caches absolute
+# positions, so this is the silent-corruption hypothesis.
+#
+# We remove the FIRST variable (lowest positions) so that a later frozen
+# constraint referencing a higher-positioned variable has its cached column
+# indices invalidated by the renumbering.
+# ---------------------------------------------------------------------------
+def _assemble_for_removal(freeze: bool) -> Model:
+    m = Model(freeze_constraints=freeze)
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x")
+    m.add_constraints(x >= 1, name="cx")
+    y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="j")], name="y")
+    m.add_constraints(y >= 7, name="cy")
+    m.add_objective(x.sum() + y.sum())
+    return m
+
+
+def test_D_remove_first_var_matrix_matches_mutable() -> None:
+    m_mut = _assemble_for_removal(freeze=False)
+    m_frz = _assemble_for_removal(freeze=True)
+    # Removing x also removes cx (references x); cy (on y) must survive intact.
+    m_mut.remove_variables("x")
+    m_frz.remove_variables("x")
+    A_mut = _dense_A(m_mut)
+    A_frz = _dense_A(m_frz)
+    assert A_mut.shape == A_frz.shape
+    np.testing.assert_array_equal(A_mut, A_frz)
+
+
+def test_D_remove_first_var_solves_equal() -> None:
+    if "highs" not in linopy.available_solvers:
+        pytest.skip("highs unavailable")
+    m_mut = _assemble_for_removal(freeze=False)
+    m_frz = _assemble_for_removal(freeze=True)
+    m_mut.remove_variables("x")
+    m_frz.remove_variables("x")
+    m_mut.solve(solver_name="highs")
+    m_frz.solve(solver_name="highs")
+    obj_mut = m_mut.objective.value
+    obj_frz = m_frz.objective.value
+    assert obj_mut is not None and obj_frz is not None
+    assert obj_frz == pytest.approx(obj_mut)
+
+
+# ---------------------------------------------------------------------------
+# Scenario E: a surviving frozen constraint whose terms STRADDLE the removed
+# variable (references a variable before it and one after it). Exercises the
+# column remap within a single CSR row and the sorted-indices assumption.
+# ---------------------------------------------------------------------------
+def _assemble_straddle(freeze: bool) -> Model:
+    m = Model(freeze_constraints=freeze)
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x")
+    y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y")
+    z = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="z")
+    # constraint links x and z (straddles y), frozen while all three exist
+    m.add_constraints(x + z >= 4, name="cxz")
+    m.add_objective(x.sum() + y.sum() + z.sum())
+    return m
+
+
+def test_E_straddle_remove_middle_matches_mutable() -> None:
+    m_mut = _assemble_straddle(freeze=False)
+    m_frz = _assemble_straddle(freeze=True)
+    m_mut.remove_variables("y")  # y does not appear in cxz -> cxz survives
+    m_frz.remove_variables("y")
+    A_mut = _dense_A(m_mut)
+    A_frz = _dense_A(m_frz)
+    assert A_mut.shape == A_frz.shape
+    np.testing.assert_array_equal(A_mut, A_frz)
+    # sanity: CSR indices remain sorted within rows after remap
+    A_frz_sp = m_frz.matrices.A
+    assert A_frz_sp is not None
+    dense_before = np.asarray(A_frz_sp.todense())
+    A_frz_sp.sort_indices()
+    np.testing.assert_array_equal(np.asarray(A_frz_sp.todense()), dense_before)
+
+
+# ---------------------------------------------------------------------------
+# Scenario F: multiple sequential removals, each renumbering again.
+# ---------------------------------------------------------------------------
+def _assemble_three(freeze: bool) -> Model:
+    m = Model(freeze_constraints=freeze)
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x")
+    m.add_constraints(x >= 1, name="cx")
+    y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y")
+    m.add_constraints(y >= 2, name="cy")
+    z = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="z")
+    m.add_constraints(z >= 3, name="cz")
+    m.add_objective(x.sum() + y.sum() + z.sum())
+    return m
+
+
+def test_F_multiple_sequential_removals_match_mutable() -> None:
+    m_mut = _assemble_three(freeze=False)
+    m_frz = _assemble_three(freeze=True)
+    for name in ("x", "y"):  # remove two, leaving only z + cz
+        m_mut.remove_variables(name)
+        m_frz.remove_variables(name)
+        np.testing.assert_array_equal(_dense_A(m_mut), _dense_A(m_frz))
+
+
+# ---------------------------------------------------------------------------
+# Scenario G: interleave removal and addition (remove, then add a new var and
+# a new frozen constraint that widens again).
+# ---------------------------------------------------------------------------
+def test_G_remove_then_add_matches_mutable() -> None:
+    def build(freeze: bool) -> Model:
+        m = Model(freeze_constraints=freeze)
+        x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x")
+        m.add_constraints(x >= 1, name="cx")
+        y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y")
+        m.add_constraints(y >= 2, name="cy")
+        m.remove_variables("x")  # renumber: cy remapped
+        w = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="w")
+        m.add_constraints(w >= 5, name="cw")  # widen again
+        m.add_objective(y.sum() + w.sum())
+        return m
+
+    np.testing.assert_array_equal(_dense_A(build(False)), _dense_A(build(True)))
+
+
+# ---------------------------------------------------------------------------
+# Scenario H: removing a variable also drops a frozen constraint referencing it
+# even when that constraint references surviving variables too (parity check).
+# ---------------------------------------------------------------------------
+def test_H_removal_drops_referencing_frozen_constraint() -> None:
+    def build(freeze: bool) -> Model:
+        m = Model(freeze_constraints=freeze)
+        x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x")
+        y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y")
+        m.add_constraints(x + y >= 3, name="cxy")
+        m.add_constraints(y >= 1, name="cy")
+        m.add_objective(x.sum() + y.sum())
+        return m
+
+    m_mut = build(False)
+    m_frz = build(True)
+    m_mut.remove_variables("x")  # drops cxy in both; cy survives
+    m_frz.remove_variables("x")
+    assert set(m_mut.constraints) == set(m_frz.constraints)
+    np.testing.assert_array_equal(_dense_A(m_mut), _dense_A(m_frz))