feat: support mixed per-element signs in frozen Constraint, add rhs/lhs setters

- Store _sign as str (uniform, fast) or np.ndarray (mixed, per-row)
- Add rhs/lhs setters on Constraint via _refreeze_after pattern
- Invalidate _dual on mutation; update netcdf serialization for array signs
- Add tests for setters, mixed-sign freeze/roundtrip/sanitize/repr/netcdf

Author: FabianHofmann

linopy/constraints.py

@@ -511,7 +511,7 @@ def __init__(
         csr: scipy.sparse.csr_array,
         con_labels: np.ndarray,
         rhs: np.ndarray,
-        sign: str,
+        sign: str | np.ndarray,
         coords: list[pd.Index],
         model: Model,
         name: str = "",
@@ -596,7 +596,7 @@ def coord_names(self) -> list[str]:
         return [c.name for c in self._coords]
 
     def _active_to_dataarray(
-        self, active_values: np.ndarray, fill: float | int = -1
+        self, active_values: np.ndarray, fill: float | int | str = -1
     ) -> DataArray:
         full = np.full(self.full_size, fill, dtype=active_values.dtype)
         full[self.active_positions] = active_values
@@ -633,14 +633,40 @@ def vars(self) -> DataArray:
 
     @property
     def sign(self) -> DataArray:
-        """Get sign DataArray (scalar, same sign for all entries)."""
-        return DataArray(np.full(self.shape, self._sign), coords=self._coords)
+        """Get sign DataArray."""
+        if isinstance(self._sign, str):
+            return DataArray(np.full(self.shape, self._sign), coords=self._coords)
+        return self._active_to_dataarray(self._sign, fill="")
 
     @property
     def rhs(self) -> DataArray:
         """Get RHS DataArray, shape (*coord_dims)."""
         return self._active_to_dataarray(self._rhs, fill=np.nan)
 
+    @rhs.setter
+    def rhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None:
+        self._refreeze_after(lambda mc: setattr(mc, "rhs", value))
+
+    @property
+    def lhs(self) -> expressions.LinearExpression:
+        """Get LHS as LinearExpression (triggers Dataset reconstruction)."""
+        return self.mutable().lhs
+
+    @lhs.setter
+    def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None:
+        self._refreeze_after(lambda mc: setattr(mc, "lhs", value))
+
+    def _refreeze_after(self, mutate: Callable[[MutableConstraint], None]) -> None:
+        mc = self.mutable()
+        mutate(mc)
+        refrozen = Constraint.from_mutable(mc, self._cindex)
+        self._csr = refrozen._csr
+        self._con_labels = refrozen._con_labels
+        self._rhs = refrozen._rhs
+        self._sign = refrozen._sign
+        self._coords = refrozen._coords
+        self._dual = None
+
     @property
     @has_optimized_model
     def dual(self) -> DataArray:
@@ -745,7 +771,8 @@ def row_expr(row: int) -> str:
             coeffs_row = np.zeros(nterm, dtype=csr.dtype)
             vars_row[: end - start] = vlabels[csr.indices[start:end]]
             coeffs_row[: end - start] = csr.data[start:end]
-            return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[self._sign]} {self._rhs[row]}"
+            sign = self._sign if isinstance(self._sign, str) else self._sign[row]
+            return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}"
 
         if size > 1:
             for indices in generate_indices_for_printout(shape, max_lines):
@@ -787,21 +814,22 @@ def to_netcdf_ds(self) -> Dataset:
             "rhs": DataArray(self._rhs, dims=["_flat"]),
             "_con_labels": DataArray(self._con_labels, dims=["_flat"]),
         }
+        if isinstance(self._sign, np.ndarray):
+            data_vars["_sign"] = DataArray(self._sign, dims=["_flat"])
         data_vars.update(coords_to_dataset_vars(self._coords))
         if self._dual is not None:
             data_vars["dual"] = DataArray(self._dual, dims=["_flat"])
         dim_names = [c.name for c in self._coords]
-        return Dataset(
-            data_vars,
-            attrs={
-                "_linopy_format": "csr",
-                "sign": self._sign,
-                "cindex": self._cindex if self._cindex is not None else -1,
-                "shape": list(csr.shape),
-                "coord_dims": dim_names,
-                "name": self._name,
-            },
-        )
+        attrs: dict[str, Any] = {
+            "_linopy_format": "csr",
+            "cindex": self._cindex if self._cindex is not None else -1,
+            "shape": list(csr.shape),
+            "coord_dims": dim_names,
+            "name": self._name,
+        }
+        if isinstance(self._sign, str):
+            attrs["sign"] = self._sign
+        return Dataset(data_vars, attrs=attrs)
 
     @classmethod
     def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint:
@@ -813,7 +841,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint:
             shape=shape,
         )
         rhs = ds["rhs"].values
-        sign = attrs["sign"]
+        sign: str | np.ndarray = ds["_sign"].values if "_sign" in ds else attrs["sign"]
         cindex = int(attrs["cindex"])
         cindex = cindex if cindex >= 0 else None
         coord_dims = attrs["coord_dims"]
@@ -841,7 +869,10 @@ 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."""
-        sense = np.full(len(self._rhs), self._sign[0])
+        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
 
     def sanitize_zeros(self) -> Constraint:
@@ -856,18 +887,25 @@ def sanitize_missings(self) -> Constraint:
 
     def sanitize_infinities(self) -> Constraint:
         """Mask out rows with invalid infinite RHS values (mutates in-place)."""
-        if self._sign == LESS_EQUAL:
-            invalid = self._rhs == np.inf
-        elif self._sign == GREATER_EQUAL:
-            invalid = self._rhs == -np.inf
+        if isinstance(self._sign, str):
+            if self._sign == LESS_EQUAL:
+                invalid = self._rhs == np.inf
+            elif self._sign == GREATER_EQUAL:
+                invalid = self._rhs == -np.inf
+            else:
+                return self
         else:
-            return self
+            invalid = ((self._sign == LESS_EQUAL) & (self._rhs == np.inf)) | (
+                (self._sign == GREATER_EQUAL) & (self._rhs == -np.inf)
+            )
         if not invalid.any():
             return self
         keep = ~invalid
         self._csr = self._csr[keep]
         self._con_labels = self._con_labels[keep]
         self._rhs = self._rhs[keep]
+        if not isinstance(self._sign, str):
+            self._sign = self._sign[keep]
         return self
 
     def freeze(self) -> Constraint:
@@ -907,13 +945,14 @@ def from_mutable(
         active_mask = (labels_flat != -1) & (vars_flat != -1).any(axis=1)
         rhs = con.rhs.values.ravel()[active_mask]
         sign_vals = con.sign.values.ravel()
-        unique_signs = np.unique(sign_vals[active_mask])
-        if len(unique_signs) > 1:
-            raise ValueError(
-                "Constraint has per-element signs; cannot freeze to immutable Constraint. "
-                "This is a known limitation — use MutableConstraint instead."
-            )
-        sign = str(unique_signs[0]) if len(unique_signs) == 1 else "="
+        active_signs = sign_vals[active_mask]
+        unique_signs = np.unique(active_signs)
+        if len(unique_signs) == 0:
+            sign: str | np.ndarray = "="
+        elif len(unique_signs) == 1:
+            sign = str(unique_signs[0])
+        else:
+            sign = active_signs
         dual = (
             con.data["dual"].values.ravel()[active_mask] if "dual" in con.data else None
         )

test/test_constraint.py

@@ -734,15 +734,17 @@ def test_freeze_mutable_roundtrip_with_masking() -> None:
     assert frozen.ncons == refrozen.ncons == 3
 
 
-def test_from_mutable_mixed_signs_raises() -> None:
+def test_from_mutable_mixed_signs() -> None:
     m = Model()
     x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x")
     m.add_constraints(x >= 0, name="mixed", freeze=False)
     mc = m.constraints["mixed"]
     assert isinstance(mc, MutableConstraint)
     mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"])
-    with pytest.raises(ValueError, match="per-element signs"):
-        Constraint.from_mutable(mc)
+    frozen = Constraint.from_mutable(mc)
+    assert isinstance(frozen._sign, np.ndarray)
+    assert list(frozen._sign) == ["<=", ">=", "<="]
+    assert_equal(frozen.sign, mc.sign)
 
 
 def test_variable_label_index(m: Model) -> None:
@@ -787,3 +789,102 @@ def test_constraint_repr_shows_variable_names(m: Model) -> None:
     c = m.constraints["c"]
     r = repr(c)
     assert "x" in r
+
+
+def test_freeze_mixed_signs_from_rule() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+    coords = [pd.RangeIndex(4, name="i")]
+
+    def bound(m, i):
+        if i % 2:
+            return x.at[i] >= i
+        return x.at[i] == 0.0
+
+    con = m.add_constraints(bound, coords=coords, name="mixed_rule")
+    assert isinstance(con, Constraint)
+    assert isinstance(con._sign, np.ndarray)
+    assert con.ncons == 4
+    expected_signs = ["=", ">=", "=", ">="]
+    assert list(con._sign) == expected_signs
+    np.testing.assert_array_equal(con.sign.values, expected_signs)
+
+
+def test_frozen_rhs_setter() -> None:
+    m = Model()
+    time = pd.RangeIndex(5, name="t")
+    x = m.add_variables(lower=0, coords=[time], name="x")
+    con = m.add_constraints(x >= 1, name="c")
+    assert isinstance(con, Constraint)
+    con.rhs = 10
+    np.testing.assert_array_equal(con._rhs, np.full(5, 10.0))
+    factor = pd.Series(range(5), index=time)
+    con.rhs = 2 * factor
+    np.testing.assert_array_equal(con._rhs, 2 * np.arange(5, dtype=float))
+
+
+def test_frozen_lhs_setter() -> None:
+    m = Model()
+    time = pd.RangeIndex(5, name="t")
+    x = m.add_variables(lower=0, coords=[time], name="x")
+    y = m.add_variables(lower=0, coords=[time], name="y")
+    con = m.add_constraints(x >= 0, name="c")
+    assert isinstance(con, Constraint)
+    con.lhs = 3 * x + 2 * y
+    lhs = con.mutable().lhs
+    assert lhs.nterm == 2
+
+
+def test_frozen_setter_invalidates_dual() -> None:
+    m = Model()
+    x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x")
+    con = m.add_constraints(x >= 0, name="c")
+    con._dual = np.array([1.0, 2.0, 3.0])
+    con.rhs = 10
+    assert con._dual is None
+
+
+def test_mixed_sign_to_matrix_with_rhs() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+    coords = [pd.RangeIndex(4, name="i")]
+
+    def bound(m, i):
+        if i % 2:
+            return x.at[i] >= i
+        return x.at[i] == 0.0
+
+    con = m.add_constraints(bound, coords=coords, name="c")
+    li = m.variables.label_index
+    csr, con_labels, b, sense = con.to_matrix_with_rhs(li)
+    assert len(sense) == 4
+    assert list(sense) == ["=", ">", "=", ">"]
+
+
+def test_mixed_sign_sanitize_infinities() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+    m.add_constraints(x >= 0, name="c", freeze=False)
+    mc = m.constraints["c"]
+    mc._data["sign"] = xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"])
+    mc._data["rhs"] = xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"])
+    frozen = mc.freeze()
+    frozen.sanitize_infinities()
+    assert frozen.ncons == 2
+    np.testing.assert_array_equal(frozen._rhs, [1.0, 2.0])
+
+
+def test_mixed_sign_repr() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+    coords = [pd.RangeIndex(4, name="i")]
+
+    def bound(m, i):
+        if i % 2:
+            return x.at[i] >= i
+        return x.at[i] == 0.0
+
+    con = m.add_constraints(bound, coords=coords, name="c")
+    r = repr(con)
+    assert "≥" in r
+    assert "=" in r

test/test_io.py

@@ -95,6 +95,31 @@ def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None:
     assert_model_equal(m, p)
 
 
+def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None:
+    from linopy.constraints import Constraint
+
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+
+    def bound(m, i):
+        if i % 2:
+            return x.at[i] >= i
+        return x.at[i] == 0.0
+
+    m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c")
+    assert isinstance(m.constraints["c"], Constraint)
+
+    fn = tmp_path / "test_mixed_sign.nc"
+    m.to_netcdf(fn)
+    p = read_netcdf(fn)
+
+    assert isinstance(p.constraints["c"], Constraint)
+    import numpy as np
+
+    np.testing.assert_array_equal(m.constraints["c"]._sign, p.constraints["c"]._sign)
+    assert_model_equal(m, p)
+
+
 def test_model_to_netcdf_with_sense(model: Model, tmp_path: Path) -> None:
     m = model
     m.objective.sense = "max"