diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 80f9076e..99128ed9 100644 --- a/doc/release_notes.rst +++ b/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. +* Frozen (CSR-backed) constraints now stay consistent when the model's variables change after freezing. The cached CSR records the active-variable ordering it was built against and is lazily reconciled to the current ordering on access: adding variables (e.g. under ``Model(freeze_constraints=True)`` while building incrementally) previously failed with ``ValueError: incompatible dimensions for axis 1`` when assembling the constraint matrix, and ``Model.remove_variables`` left the cached column positions of surviving frozen constraints stale, producing a wrong-width matrix or an ``IndexError`` at solve time. All consumers of the CSR (matrix assembly, ``.data`` reconstruction, ``repr``, ``to_polars``, netcdf export) go through the reconciliation, and a variable removed behind the model's back while still referenced is reported with a clear error. +* ``repr`` of a frozen constraint on variables with masked entries no longer fails (or shows the wrong variables): CSR column positions are now decoded to variable labels before printing. * ``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 diff --git a/linopy/constraints.py b/linopy/constraints.py index 0b9dbb0a..c5ed9d12 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -501,10 +501,12 @@ 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_rows, n_active_vars). Rows are the active + (non-masked) flat positions of the constraint grid; column indices + are dense positions in the model's active-variable ordering + (``label_index.vlabels``) as of ``vlabels``. rhs : np.ndarray - Shape (n_flat,). Right-hand-side values. + Shape (n_active_rows,). Right-hand-side values. sign : str or np.ndarray Constraint sign. Either a single str ('=', '<=', '>=') for uniform signs, or a per-row np.ndarray of sign strings for mixed signs. @@ -517,11 +519,18 @@ class CSRConstraint(ConstraintBase): cindex : int or None Starting label assigned by the model. None if not yet assigned. dual : np.ndarray or None - Shape (n_flat,). Dual values after solving, or None. + Shape (n_active_rows,). Dual values after solving, or None. + vlabels : np.ndarray or None + The active-variable basis (``label_index.vlabels``) the CSR's column + indices refer to. None (default) means the model's current basis. + Whenever the model's basis changes afterwards (variables added or + removed), the CSR is lazily reconciled on access, see + ``_reconciled_csr``. """ __slots__ = ( "_csr", + "_vlabels", "_con_labels", "_rhs", "_sign", @@ -547,8 +556,12 @@ def __init__( dual: np.ndarray | None = None, binvar_labels: np.ndarray | None = None, binval: int | np.ndarray | None = None, + vlabels: np.ndarray | None = None, ) -> None: self._csr = csr + self._vlabels = ( + model.variables.label_index.vlabels if vlabels is None else vlabels + ) self._con_labels = con_labels self._rhs = rhs self._sign = sign @@ -740,7 +753,7 @@ def _to_dataset(self, nterm: int) -> Dataset: ------- Dataset with variables ``labels``, ``coeffs``, ``vars``. """ - csr = self._csr + csr = self._reconciled_csr() counts = np.diff(csr.indptr) shape = self.shape full_size = self.full_size @@ -806,7 +819,8 @@ def __repr__(self) -> str: dim_names = self.coord_names size = self.full_size # Map active rows (CSR is active-only) back to full flat positions - csr = self._csr + csr = self._reconciled_csr() + vlabels = self._model.variables.label_index.vlabels nterm = self.nterm active_positions = self.active_positions masked_entries = size - len(active_positions) @@ -823,7 +837,8 @@ def row_expr(row: int) -> str: start, end = int(csr.indptr[row]), int(csr.indptr[row + 1]) vars_row = np.full(nterm, -1, dtype=np.int64) coeffs_row = np.zeros(nterm, dtype=csr.dtype) - vars_row[: end - start] = csr.indices[start:end] + # csr.indices are column positions into vlabels; decode to labels + vars_row[: end - start] = vlabels[csr.indices[start:end]] coeffs_row[: end - start] = csr.data[start:end] sign = self._sign if isinstance(self._sign, str) else self._sign[row] return f"{format_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}" @@ -852,15 +867,90 @@ def row_expr(row: int) -> str: return "\n".join(lines) + def _reconciled_csr(self) -> scipy.sparse.csr_array: + """ + Return the stored CSR reconciled to the current active-variable basis. + + The CSR's column indices are dense positions in the model's + active-variable ordering (``label_index.vlabels``) as of the last + reconciliation (initially: construction/freeze time). Any change to + the active-variable set afterwards changes that ordering, so every + consumer that interprets column indices must go through this method + first. Three cases: + + - The basis object is unchanged (the common case; ``label_index`` + caches ``vlabels`` until variables are added or removed): return + the stored CSR as is, O(1). + - The stored basis is a prefix of the current one (variables were + only added): positions are stable, only the width grows. Metadata + change only, the index arrays are reused. + - Otherwise (variables were removed, possibly combined with + additions): remap each stored position through its variable label + to the new position. Surviving labels keep their relative order in + the new basis, so per-row indices remain sorted. + + The reconciled CSR is stored back together with the current basis, so + subsequent accesses take the fast path. ``_csr`` is rebound, never + mutated in place (copy-on-write, see ``sanitize_zeros``); holding a + reference to the (shared) ``vlabels`` array is cheap because all + constraints reconciled against the same basis share one object. + """ + label_index = self._model.variables.label_index + new_vlabels = label_index.vlabels + old_vlabels = self._vlabels + if new_vlabels is old_vlabels: + return self._csr + + csr = self._csr + n_old = len(old_vlabels) + if len(new_vlabels) >= n_old and np.array_equal( + new_vlabels[:n_old], old_vlabels + ): + if csr.shape[1] != len(new_vlabels): + csr = scipy.sparse.csr_array( + (csr.data, csr.indices, csr.indptr), + shape=(csr.shape[0], len(new_vlabels)), + ) + else: + referenced = old_vlabels[csr.indices] + new_positions = label_index.label_to_pos[referenced] + if (new_positions < 0).any(): + missing = np.unique(referenced[new_positions < 0]) + raise ValueError( + f"Frozen constraint '{self._name}' references variables that " + f"are no longer part of the model (labels {missing.tolist()}). " + "Use `Model.remove_variables` to remove variables; it also " + "removes the constraints referencing them." + ) + csr = scipy.sparse.csr_array( + (csr.data, new_positions, csr.indptr), + shape=(csr.shape[0], len(new_vlabels)), + ) + self._csr = csr + self._vlabels = new_vlabels + return csr + 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 the stored CSR matrix and con_labels. + + The ``label_index`` parameter exists for interface parity with + ``Constraint.to_matrix``; the basis is always the model's own + ``label_index`` (the same object callers pass in). + """ + return self._reconciled_csr(), self._con_labels def to_netcdf_ds(self) -> Dataset: - """Return a Dataset with raw CSR components for netcdf serialization.""" - csr = self._csr + """ + Return a Dataset with raw CSR components for netcdf serialization. + + The CSR is reconciled first, so the serialized column indices always + refer to the basis of the variables serialized alongside; loading + (``from_netcdf_ds``) then correctly adopts the load-time basis. + """ + csr = self._reconciled_csr() data_vars: dict[str, DataArray] = { "indptr": DataArray(csr.indptr, dims=["_indptr"]), "indices": DataArray(csr.indices, dims=["_nnz"]), @@ -937,20 +1027,19 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: ) def has_variable(self, variable: variables.Variable) -> bool: + csr = self._reconciled_csr() vlabels = self._model.variables.label_index.vlabels - return bool( - np.isin(vlabels[self._csr.indices], variable.labels.values.ravel()).any() - ) + return bool(np.isin(vlabels[csr.indices], variable.labels.values.ravel()).any()) 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, reconciled to the current basis.""" 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._reconciled_csr(), self._con_labels, self._rhs, sense def active_labels(self) -> np.ndarray: return self._con_labels @@ -1008,7 +1097,7 @@ def mutable(self) -> Constraint: def to_polars(self) -> pl.DataFrame: """Convert frozen constraint to polars DataFrame directly from CSR.""" - csr = self._csr + csr = self._reconciled_csr() sign_dtype = pl.Enum(["=", "<=", ">="]) if csr.nnz == 0: return pl.DataFrame( @@ -1052,21 +1141,23 @@ def iterate_slices( would be misleading. Do not call ``.data``, ``.mutable()``, or any coord-dependent property on batch slices. """ - nnz = self._csr.nnz + csr = self._reconciled_csr() + nnz = csr.nnz if slice_size is None or nnz <= slice_size: yield self return - for rows in _equal_nnz_slices(self._csr.indptr, slice_size): + for rows in _equal_nnz_slices(csr.indptr, slice_size): sign = self._sign if isinstance(self._sign, str) else self._sign[rows] yield CSRConstraint( - csr=self._csr[rows], + csr=csr[rows], con_labels=self._con_labels[rows], rhs=self._rhs[rows], sign=sign, coords=[], model=self._model, name=self._name, + vlabels=self._vlabels, ) @classmethod @@ -2234,6 +2325,9 @@ def reset_dual(self) -> None: c._name, cindex=c._cindex, dual=None, + binvar_labels=c._binvar_labels, + binval=c._binval, + vlabels=c._vlabels, ) elif isinstance(c, Constraint): if "dual" in c.data: diff --git a/test/test_frozen_constraint_active_vars.py b/test/test_frozen_constraint_active_vars.py new file mode 100644 index 00000000..2fd54922 --- /dev/null +++ b/test/test_frozen_constraint_active_vars.py @@ -0,0 +1,490 @@ +""" +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. + +The mechanism under test is ``CSRConstraint._reconciled_csr``: the frozen CSR +records the active-variable basis its column indices refer to and is lazily +reconciled to the model's current basis on access (widened after additions, +remapped through variable labels after removals). Scenarios J-O cover the +consumers beyond matrix assembly (``.data``, ``to_polars``, ``repr``, +``has_variable``, ``iterate_slices``, netcdf round-trip) and the guard against +dangling references. +""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +import linopy +from linopy import Model, read_netcdf +from linopy.constraints import CSRConstraint + + +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)) + + +# --------------------------------------------------------------------------- +# Scenario I: shrink -- remove the LAST-added variable array, so the new basis +# is a proper prefix of the stored one. A pure "widen if too narrow" check +# misses this case, and a naive elementwise basis comparison would compare +# arrays of unequal length. +# --------------------------------------------------------------------------- +def test_I_remove_last_var_shrinks_matrix() -> 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 >= 1, name="cx") # frozen with width covering x and y + m.add_objective(x.sum() + y.sum()) + return m + + m_mut = build(False) + m_frz = build(True) + m_mut.remove_variables("y") # cx survives; basis shrinks + m_frz.remove_variables("y") + A_mut = _dense_A(m_mut) + A_frz = _dense_A(m_frz) + assert A_frz.shape == (2, 2) + np.testing.assert_array_equal(A_mut, A_frz) + + +# --------------------------------------------------------------------------- +# Scenario J: non-matrix consumers after basis changes. The CSR's column +# indices are interpreted against the current variable basis not only when +# assembling the constraint matrix, but also by ``.data`` (Dataset +# reconstruction), ``to_polars`` (LP export), ``iterate_slices``, +# ``has_variable``, and ``repr``. All must decode to the same variable labels +# as the mutable path after variables were added and removed. +# --------------------------------------------------------------------------- +def _assemble_straddle_then_mutate(freeze: bool) -> Model: + m = _assemble_straddle(freeze) + m.remove_variables("y") # remap: z's positions shift down + m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="w") # widen + return m + + +def test_J_data_reconstruction_matches_mutable() -> None: + m_mut = _assemble_straddle_then_mutate(freeze=False) + m_frz = _assemble_straddle_then_mutate(freeze=True) + con_mut = m_mut.constraints["cxz"] + con_frz = m_frz.constraints["cxz"] + assert isinstance(con_frz, CSRConstraint) + ds_mut = con_mut.data + ds_frz = con_frz.data + np.testing.assert_array_equal(ds_mut.vars.values, ds_frz.vars.values) + np.testing.assert_array_equal(ds_mut.coeffs.values, ds_frz.coeffs.values) + np.testing.assert_array_equal(ds_mut.labels.values, ds_frz.labels.values) + + +def test_J_to_polars_decodes_current_labels() -> None: + m = _assemble_straddle_then_mutate(freeze=True) + con = m.constraints["cxz"] + assert isinstance(con, CSRConstraint) + df = con.to_polars() + expected_vars = np.stack( + [ + m.variables["x"].labels.values, + m.variables["z"].labels.values, + ] + ).T.ravel() # per row: x[i], z[i] + np.testing.assert_array_equal(df["vars"].to_numpy(), expected_vars) + + +def test_J_iterate_slices_decodes_current_labels() -> None: + m = _assemble_straddle_then_mutate(freeze=True) + con = m.constraints["cxz"] + assert isinstance(con, CSRConstraint) + dfs = [s.to_polars() for s in con.iterate_slices(slice_size=2)] + assert len(dfs) > 1 # actually sliced + vars_concat = np.concatenate([df["vars"].to_numpy() for df in dfs]) + expected_vars = np.stack( + [ + m.variables["x"].labels.values, + m.variables["z"].labels.values, + ] + ).T.ravel() + np.testing.assert_array_equal(vars_concat, expected_vars) + + +def test_J_has_variable_after_mutation() -> None: + m = _assemble_straddle_then_mutate(freeze=True) + con = m.constraints["cxz"] + assert isinstance(con, CSRConstraint) + assert con.has_variable(m.variables["x"]) + assert con.has_variable(m.variables["z"]) + assert not con.has_variable(m.variables["w"]) + + +def test_J_repr_after_mutation_matches_mutable_body() -> None: + m_mut = _assemble_straddle_then_mutate(freeze=False) + m_frz = _assemble_straddle_then_mutate(freeze=True) + r_mut = repr(m_mut.constraints["cxz"]) + r_frz = repr(m_frz.constraints["cxz"]) + for snippet in ("x[0]", "z[0]", "x[1]", "z[1]"): + assert snippet in r_frz + # the expression bodies (everything after the header) agree + assert r_mut.split("\n")[2:] == r_frz.split("\n")[2:] + + +# --------------------------------------------------------------------------- +# Scenario K: netcdf round-trip after basis changes. Serialization reconciles +# first, so the stored column indices always match the variables stored +# alongside, and the loaded model adopts the load-time basis. +# --------------------------------------------------------------------------- +def test_K_netcdf_roundtrip_after_mutation(tmp_path: Path) -> None: + m = _assemble_straddle_then_mutate(freeze=True) + fn = tmp_path / "frozen_mutated.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + assert isinstance(p.constraints["cxz"], CSRConstraint) + np.testing.assert_array_equal(_dense_A(m), _dense_A(p)) + + +# --------------------------------------------------------------------------- +# Scenario L: a dangling reference (variable removed behind the model's back, +# bypassing Model.remove_variables) is detected and reported instead of +# producing a silently wrong matrix. +# --------------------------------------------------------------------------- +def test_L_dangling_reference_raises() -> None: + m = Model(freeze_constraints=True) + x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x") + m.add_constraints(x >= 1, name="cx") + m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y") + m.variables.remove("x") # bypasses constraint cleanup on purpose + con = m.constraints["cx"] + assert isinstance(con, CSRConstraint) + with pytest.raises(ValueError, match="no longer part of the model"): + con.to_matrix() + + +# --------------------------------------------------------------------------- +# Scenario M: masked variables. Masked entries hold label -1, so dense +# positions and labels diverge; decoding CSR columns as labels without going +# through the basis is wrong for any masked model. +# --------------------------------------------------------------------------- +def _assemble_masked(freeze: bool) -> Model: + m = Model(freeze_constraints=freeze) + mask = xr.DataArray([False, True, True], coords=[pd.RangeIndex(3, name="i")]) + x = m.add_variables( + lower=0, coords=[pd.RangeIndex(3, name="i")], name="x", mask=mask + ) + 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 test_M_masked_matrix_matches_mutable() -> None: + np.testing.assert_array_equal( + _dense_A(_assemble_masked(False)), _dense_A(_assemble_masked(True)) + ) + + +def test_M_masked_repr_decodes_labels() -> None: + m = _assemble_masked(freeze=True) + r = repr(m.constraints["cx"]) + assert "x[1]" in r + assert "x[2]" in r + + +# --------------------------------------------------------------------------- +# Scenario N: an empty CSR (all rows masked away) survives additions and +# removals of other variables. +# --------------------------------------------------------------------------- +def test_N_empty_csr_reconciles() -> 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") + mask = xr.DataArray([False, False], coords=[pd.RangeIndex(2, name="i")]) + m.add_constraints(x >= 1, name="c_empty", mask=mask) + y = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y") + m.add_constraints(y >= 2, name="cy") + m.add_objective(x.sum() + y.sum()) + return m + + m_mut = build(False) + m_frz = build(True) + m_mut.remove_variables("x") + m_frz.remove_variables("x") + np.testing.assert_array_equal(_dense_A(m_mut), _dense_A(m_frz)) + + +# --------------------------------------------------------------------------- +# Scenario O: whitebox -- reconcile caching and copy-on-write semantics. +# --------------------------------------------------------------------------- +def test_O_reconcile_caches_and_rebinds() -> None: + m = Model(freeze_constraints=True) + x = m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="x") + m.add_constraints(x >= 1, name="cx") + con = m.constraints["cx"] + assert isinstance(con, CSRConstraint) + + csr1 = con._reconciled_csr() + assert con._reconciled_csr() is csr1 # fast path: same basis, same object + + m.add_variables(lower=0, coords=[pd.RangeIndex(2, name="i")], name="y") + csr2 = con._reconciled_csr() + assert csr2.shape[1] == m.variables.label_index.n_active_vars # widened + # copy-on-write: the previously returned csr object is not mutated + assert csr1.shape[1] == 2 + np.testing.assert_array_equal(csr1.todense(), csr2.todense()[:, :2]) + # basis reference updated -> subsequent accesses take the fast path again + assert con._vlabels is m.variables.label_index.vlabels + assert con._reconciled_csr() is csr2