fix(copy): keep frozen constraints frozen in Model.copy

doc/release_notes.rst

@@ -81,6 +81,7 @@ Upcoming Version
 * ``CSRConstraint.loc``, ``update`` and ``from_rule``, and assigning to ``coeffs``, ``vars`` or ``sign``, now raise an ``AttributeError`` saying the operation is not supported on a frozen constraint and naming ``.mutable()`` as the way out (for ``from_rule``: build with ``Constraint.from_rule`` and ``.freeze()`` the result), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 <https://github.com/PyPSA/linopy/issues/963>`__)
 * ``linopy.merge`` with ``join="outer"``, ``"left"`` or ``"right"`` no longer raises an xarray ``MergeError`` when only one operand carries an auxiliary coordinate on a joined dimension. The constant was reindexed with a fill value of ``0`` that was also written into the auxiliary coordinate, so it no longer matched the coefficients' copy. The coordinate now follows its rows onto the joined grid and is NaN where no operand provides it, under both semantics. The same holds for ``.add`` / ``.sub`` / ``.mul`` / ``.div`` with a constant and a reindexing ``join=``.
 * 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 <https://github.com/PyPSA/linopy/issues/933>`__)
+* ``Model.copy`` (and the ``copy.copy``/``copy.deepcopy`` protocols) keeps frozen constraints as ``CSRConstraint`` instead of rebuilding them as dense ``Constraint``; ``deep=True`` copies their arrays, and duals are copied only with ``include_solution=True``. A CSR-backed objective is now deep-copied with ``deep=True`` as well. The auxiliary coordinates of a CSR grid are stored read-only, so a copy can share them safely. (`#981 <https://github.com/PyPSA/linopy/issues/981>`__)
 * ``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/csr.py

@@ -77,6 +77,11 @@ class Grid:
     indexes: dict[str, pd.Index]
     aux: AuxCoords = field(default_factory=dict)
 
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self, "aux", {n: (d, _readonly(v)) for n, (d, v) in self.aux.items()}
+        )
+
     @classmethod
     def from_coords(cls, coords: Iterable[pd.Index]) -> Grid:
         """Build from one index per dimension, each named after its dim."""
@@ -749,6 +754,15 @@ def _densify_notice(reason: str) -> None:
         )
 
 
+def _readonly(values: np.ndarray) -> np.ndarray:
+    """``values`` as a read-only array, copied first if it is writable."""
+    if not values.flags.writeable:
+        return values
+    values = values.copy()
+    values.setflags(write=False)
+    return values
+
+
 def _aux_coords(ds: Dataset | DataArray, dims: set[str]) -> AuxCoords:
     """Scalar and one-dimensional auxiliary coordinates of ``ds`` lying on ``dims``."""
     aux: AuxCoords = {}

linopy/io.py

@@ -17,11 +17,12 @@
 from io import BufferedWriter
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import numpy as np
 import pandas as pd
 import polars as pl
+import scipy.sparse
 import xarray as xr
 from tqdm import tqdm
 
@@ -34,6 +35,8 @@
 from linopy.objective import Objective, linear_part
 from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup
 
+Buffer = TypeVar("Buffer", np.ndarray, scipy.sparse.sparray)
+
 if TYPE_CHECKING:
     from cuopt.linear_programming import DataModel as cuoptDataModel
     from cupdlpx import Model as cupdlpxModel
@@ -1335,7 +1338,12 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model:
     Model
         A deep or shallow copy of the model.
     """
-    from linopy.constraints import Constraint, ConstraintBase, Constraints
+    from linopy.constraints import (
+        Constraint,
+        ConstraintBase,
+        Constraints,
+        CSRConstraint,
+    )
     from linopy.expressions import Expressions, LinearExpression, QuadraticExpression
     from linopy.model import Model, Objective
     from linopy.variables import Variable, Variables
@@ -1378,17 +1386,28 @@ def _copy_expr(
         new_model,
     )
 
-    def _copy_con_data(con: ConstraintBase) -> xr.Dataset:
-        d = con.mutable().data
-        if include_solution:
-            return d.copy(deep=deep)
-        return d[con.data_attrs].copy(deep=deep)
+    def _buffer(value: Buffer) -> Buffer:
+        return value.copy() if deep else value
+
+    def _copy_con(name: str, con: ConstraintBase) -> ConstraintBase:
+        if isinstance(con, CSRConstraint):
+            buffer_types = (np.ndarray, scipy.sparse.sparray)
+            changes: dict[str, Any] = {"model": new_model}
+            if deep:
+                kwargs = con._init_kwargs().items()
+                changes |= {
+                    k: v.copy() for k, v in kwargs if isinstance(v, buffer_types)
+                }
+            if not include_solution:
+                changes["dual"] = None
+            return con._replace(**changes)
+        d = con.data
+        if not include_solution:
+            d = d[con.data_attrs]
+        return Constraint(d.copy(deep=deep), new_model, name)
 
     new_model._constraints = Constraints(
-        {
-            name: Constraint(_copy_con_data(con), new_model, name)
-            for name, con in m.constraints.items()
-        },
+        {name: _copy_con(name, con) for name, con in m.constraints.items()},
         new_model,
     )
 
@@ -1397,7 +1416,12 @@ def _copy_con_data(con: ConstraintBase) -> xr.Dataset:
     obj_expr = (
         type(expr)(expr.data.copy(deep=deep), new_model)
         if csr is None
-        else LinearExpression._from_csr(replace(csr, model=new_model), new_model)
+        else LinearExpression._from_csr(
+            replace(
+                csr, csr=_buffer(csr.csr), const=_buffer(csr.const), model=new_model
+            ),
+            new_model,
+        )
     )
     new_model._objective = Objective(
         obj_expr, new_model, m.objective.sense, m.objective.scaling

test/test_csr.py

@@ -2203,3 +2203,83 @@ def test_soften_rejects(error: str, freeze: bool) -> None:
     call, exc, match = SOFTEN_ERRORS[error]
     with pytest.raises(exc, match=match):
         call(con)
+
+
+def frozen_model(soften: bool) -> tuple[Model, CSRConstraint]:
+    c = base_model()
+    for var in (c.gen_p, c.flow):
+        var.update(lower=0, upper=1)
+    with no_densify():
+        c.m.add_objective((1.0 * c.gen_p).groupby(c.gbus).sum(sparse=True).sum())
+        con = c.m.add_constraints(
+            c.balance_lhs(sparse=True), ">=", c.load, name="c", freeze=True
+        )
+        if soften:
+            con.soften(penalty=2.0, max_violation=20.0)
+    assert isinstance(con, CSRConstraint)
+    return c.m, con
+
+
+@pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs")
+@pytest.mark.parametrize("include_solution", [True, False])
+@pytest.mark.parametrize("deep", [True, False])
+def test_copy_keeps_frozen_constraints(deep: bool, include_solution: bool) -> None:
+    require_v1()
+    m, con = frozen_model(soften=True)
+    m.solve("highs")
+    with no_densify():
+        c = m.copy(deep=deep, include_solution=include_solution)
+    copied = c.constraints["c"]
+    assert copied.model is c and copied.name == "c"
+    assert c.objective.expression.is_sparse
+    assert con.slack is not None and copied.slack is not None
+    assert copied.slack.positive.labels.equals(con.slack.positive.labels)
+    assert copied.slack.positive.model is c
+    assert ("dual" in copied.mutable().data) == include_solution
+
+
+@pytest.mark.parametrize("deep", [True, False])
+def test_softening_frozen_copy_leaves_original(deep: bool) -> None:
+    require_v1()
+    m, con = frozen_model(soften=False)
+    with no_densify():
+        copied = m.copy(deep=deep)
+        copied.constraints["c"].soften(penalty=2.0)
+    assert con.slack is None and "c_slack_pos" not in m.variables
+    assert m.objective.expression.nterm < copied.objective.expression.nterm
+    assert con.nterm < copied.constraints["c"].nterm
+
+
+@pytest.mark.parametrize("deep", [True, False])
+def test_copy_frozen_matrices(deep: bool) -> None:
+    m = Model(freeze_constraints=True)
+    i = pd.RangeIndex(4, name="i")
+    x = m.add_variables(lower=0, coords=[i], name="x")
+    b = m.add_variables(coords=[i], binary=True, name="b")
+    lhs = (2 * x).assign_coords(aux=("i", [1.0, 2, 3, 4]))
+    mask = xr.DataArray([True, False, True, True], coords=[i])
+    scaling = xr.DataArray([1.0, 2, 3, 4], coords=[i])
+    m.add_constraints(lhs >= 1, name="c", mask=mask, scaling=scaling)
+    m.add_indicator_constraints(b, 1, x <= 5, name="ind")
+    c = m.copy(deep=deep)
+    for name, con in m.constraints.items():
+        copied = c.constraints[name]
+        assert isinstance(con, CSRConstraint) and isinstance(copied, CSRConstraint)
+        assert np.shares_memory(con._csr.data, copied._csr.data) != deep
+    got, want = c.matrices, m.matrices
+    assert got.A is not None and want.A is not None
+    assert got.indicator_A is not None and want.indicator_A is not None
+    pairs = [
+        (got.A.toarray(), want.A.toarray()),
+        (got.b, want.b),
+        (got.sense, want.sense),
+        (got.clabels, want.clabels),
+        (got.indicator_A.toarray(), want.indicator_A.toarray()),
+        (got.indicator_b, want.indicator_b),
+        (got.indicator_binvar, want.indicator_binvar),
+    ]
+    for g, w in pairs:
+        assert np.array_equal(g, w)
+    with pytest.raises(ValueError, match="read-only"):
+        c.constraints["c"].coords["aux"].values[0] = -1
+    assert c.constraints["c"].coords["aux"].equals(m.constraints["c"].coords["aux"])