feat: mixed per-element signs in frozen Constraint, rhs/lhs setters, DRY helpers

Author: FabianHofmann

linopy/constraints.py

@@ -497,9 +497,9 @@ class Constraint(ConstraintBase):
         constraint grid (including masked/empty rows).
     rhs : np.ndarray
         Shape (n_flat,). Right-hand-side values.
-    sign : str
-        Constraint sign: one of '=', '<=', '>='.
-        Note: per-element signs are not supported (documented regression vs MutableConstraint).
+    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.
     coords : list of pd.Index
         One index per coordinate dimension defining the constraint grid.
     model : Model
@@ -781,13 +781,11 @@ def __repr__(self) -> str:
         header_string = f"{self.type} `{self._name}`" if self._name else f"{self.type}"
         lines = []
 
-        vlabels = self._model.variables.label_index.vlabels
-
         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] = vlabels[csr.indices[start:end]]
+            vars_row[: end - start] = 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"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}"
@@ -859,9 +857,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint:
             shape=shape,
         )
         rhs = ds["rhs"].values
-        sign: str | np.ndarray = (
-            ds["_sign"].values if "_sign" in ds else attrs["sign"]
-        )
+        sign: str | np.ndarray = ds["_sign"].values if "_sign" in ds else attrs["sign"]
         _cindex_raw = int(attrs["cindex"])
         cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None
         coord_dims = attrs["coord_dims"]
@@ -1657,12 +1653,7 @@ def set_blocks(self, block_map: np.ndarray) -> None:
 
             res = res.where(not_missing.any(constraint.term_dim), -1)
             res = res.where(not_zero.any(constraint.term_dim), 0)
-            if isinstance(constraint, MutableConstraint):
-                constraint._data = assign_multiindex_safe(constraint.data, blocks=res)
-            else:
-                mc = constraint.mutable()
-                mc._data = assign_multiindex_safe(mc.data, blocks=res)
-                self.data[name] = Constraint.from_mutable(mc, constraint._cindex)
+            constraint._data = assign_multiindex_safe(constraint.data, blocks=res)
 
     @property
     def flat(self) -> pd.DataFrame:
@@ -1719,7 +1710,18 @@ def reset_dual(self) -> None:
         """
         for k, c in self.items():
             if isinstance(c, Constraint):
-                c._dual = None
+                if c._dual is not None:
+                    self.data[k] = Constraint(
+                        c._csr,
+                        c._con_labels,
+                        c._rhs,
+                        c._sign,
+                        c._coords,
+                        c._model,
+                        c._name,
+                        cindex=c._cindex,
+                        dual=None,
+                    )
             elif isinstance(c, MutableConstraint):
                 if "dual" in c.data:
                     c._data = c.data.drop_vars("dual")

test/test_constraint.py

@@ -709,10 +709,10 @@ def test_constraints_equalities(m: Model) -> None:
 
 def test_freeze_mutable_roundtrip(m: Model) -> None:
     frozen = m.constraints["c"]
-    assert isinstance(frozen, Constraint)
+    assert isinstance(frozen, linopy.constraints.Constraint)
     mc = frozen.mutable()
     assert isinstance(mc, MutableConstraint)
-    refrozen = Constraint.from_mutable(mc, frozen._cindex)
+    refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex)
     assert_equal(frozen.labels, refrozen.labels)
     assert_equal(frozen.rhs, refrozen.rhs)
     assert_equal(frozen.sign, refrozen.sign)
@@ -727,7 +727,7 @@ def test_freeze_mutable_roundtrip_with_masking() -> None:
     m.add_constraints(x.where(mask) >= 0, name="c")
     frozen = m.constraints["c"]
     mc = frozen.mutable()
-    refrozen = Constraint.from_mutable(mc, frozen._cindex)
+    refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex)
     assert_equal(frozen.labels, refrozen.labels)
     assert_equal(frozen.rhs, refrozen.rhs)
     assert frozen.ncons == refrozen.ncons == 3
@@ -740,7 +740,7 @@ def test_from_mutable_mixed_signs() -> None:
     mc = m.constraints["mixed"]
     assert isinstance(mc, MutableConstraint)
     mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"])
-    frozen = Constraint.from_mutable(mc)
+    frozen = linopy.constraints.Constraint.from_mutable(mc)
     assert isinstance(frozen._sign, np.ndarray)
     assert list(frozen._sign) == ["<=", ">=", "<="]
     assert_equal(frozen.sign, mc.sign)
@@ -801,7 +801,7 @@ def bound(m, i):
         return x.at[i] == 0.0
 
     con = m.add_constraints(bound, coords=coords, name="mixed_rule")
-    assert isinstance(con, Constraint)
+    assert isinstance(con, linopy.constraints.Constraint)
     assert isinstance(con._sign, np.ndarray)
     assert con.ncons == 4
     expected_signs = ["=", ">=", "=", ">="]
@@ -814,7 +814,7 @@ def test_frozen_rhs_setter() -> None:
     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)
+    assert isinstance(con, linopy.constraints.Constraint)
     con.rhs = 10
     np.testing.assert_array_equal(con._rhs, np.full(5, 10.0))
     factor = pd.Series(range(5), index=time)
@@ -828,7 +828,7 @@ def test_frozen_lhs_setter() -> None:
     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)
+    assert isinstance(con, linopy.constraints.Constraint)
     con.lhs = 3 * x + 2 * y
     lhs = con.mutable().lhs
     assert lhs.nterm == 2