fix: review fixes for PR #630 (matrix accessor rewrite)

- Fix __repr__ passing CSR positions instead of variable labels
- Fix set_blocks failing on frozen Constraint
- Extract _active_to_dataarray helper to reduce DRY violations
- Simplify reset_dual to direct mutation instead of reconstruction
- Add tests for freeze/mutable roundtrip, VariableLabelIndex,
  to_matrix_with_rhs, from_mutable mixed signs, repr correctness

Author: FabianHofmann

linopy/constraints.py

@@ -595,16 +595,19 @@ def nterm(self) -> int:
     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
+    ) -> DataArray:
+        full = np.full(self.full_size, fill, dtype=active_values.dtype)
+        full[self.active_positions] = active_values
+        return DataArray(full.reshape(self.shape), coords=self._coords)
+
     @property
     def labels(self) -> DataArray:
         """Get labels DataArray, shape (*coord_dims)."""
         if self._cindex is None:
             return DataArray([])
-        shape = self.shape
-        full_size = self.full_size
-        labels_flat = np.full(full_size, -1, dtype=np.int64)
-        labels_flat[self.active_positions] = self._con_labels
-        return DataArray(labels_flat.reshape(shape), coords=self._coords)
+        return self._active_to_dataarray(self._con_labels, fill=-1)
 
     @property
     def coeffs(self) -> DataArray:
@@ -636,10 +639,7 @@ def sign(self) -> DataArray:
     @property
     def rhs(self) -> DataArray:
         """Get RHS DataArray, shape (*coord_dims)."""
-        shape = self.shape
-        rhs_full = np.full(self.full_size, np.nan)
-        rhs_full[self.active_positions] = self._rhs
-        return DataArray(rhs_full.reshape(shape), coords=self._coords)
+        return self._active_to_dataarray(self._rhs, fill=np.nan)
 
     @property
     @has_optimized_model
@@ -649,9 +649,7 @@ def dual(self) -> DataArray:
             raise AttributeError(
                 "Underlying is optimized but does not have dual values stored."
             )
-        dual_full = np.full(self.full_size, np.nan)
-        dual_full[self.active_positions] = self._dual
-        return DataArray(dual_full.reshape(self.shape), coords=self._coords)
+        return self._active_to_dataarray(self._dual, fill=np.nan)
 
     @dual.setter
     def dual(self, value: DataArray) -> None:
@@ -713,24 +711,10 @@ def _to_dataset(self, nterm: int) -> Dataset:
     def data(self) -> Dataset:
         """Reconstruct the xarray Dataset from the CSR representation."""
         ds = self._to_dataset(self.nterm)
-        shape = self.shape
-        active_pos = self.active_positions
-        rhs_full = np.full(self.full_size, np.nan)
-        rhs_full[active_pos] = self._rhs
-        ds = ds.assign(
-            sign=DataArray(np.full(shape, self._sign), coords=self._coords),
-            rhs=DataArray(rhs_full.reshape(shape), coords=self._coords),
-        )
+        ds = ds.assign(sign=self.sign, rhs=self.rhs)
         if self._dual is not None:
-            dual_full = np.full(self.full_size, np.nan)
-            dual_full[active_pos] = self._dual
-            ds = ds.assign(
-                dual=DataArray(dual_full.reshape(shape), coords=self._coords)
-            )
-        attrs: dict[str, Any] = {"name": self._name}
-        if self._cindex is not None:
-            attrs["label_range"] = (self._cindex, self._cindex + self.full_size)
-        return ds.assign_attrs(attrs)
+            ds = ds.assign(dual=self._active_to_dataarray(self._dual, fill=np.nan))
+        return ds.assign_attrs(self.attrs)
 
     def __repr__(self) -> str:
         """Print the constraint without reconstructing the full Dataset."""
@@ -753,11 +737,13 @@ 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] = csr.indices[start:end]
+            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]}"
 
@@ -1610,7 +1596,12 @@ 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)
-            constraint._data = assign_multiindex_safe(constraint.data, blocks=res)
+            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)
 
     @property
     def flat(self) -> pd.DataFrame:
@@ -1667,18 +1658,7 @@ def reset_dual(self) -> None:
         """
         for k, c in self.items():
             if isinstance(c, Constraint):
-                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,
-                    )
+                c._dual = None
             else:
                 if "dual" in c.data:
                     c._data = c.data.drop_vars("dual")

test/test_constraint.py

@@ -706,3 +706,84 @@ def test_constraints_inequalities(m: Model) -> None:
 
 def test_constraints_equalities(m: Model) -> None:
     assert isinstance(m.constraints.equalities, Constraints)
+
+
+def test_freeze_mutable_roundtrip(m: Model) -> None:
+    frozen = m.constraints["c"]
+    assert isinstance(frozen, Constraint)
+    mc = frozen.mutable()
+    assert isinstance(mc, MutableConstraint)
+    refrozen = Constraint.from_mutable(mc, frozen._cindex)
+    assert_equal(frozen.labels, refrozen.labels)
+    assert_equal(frozen.rhs, refrozen.rhs)
+    assert_equal(frozen.sign, refrozen.sign)
+    np.testing.assert_array_equal(frozen._csr.toarray(), refrozen._csr.toarray())
+    np.testing.assert_array_equal(frozen._con_labels, refrozen._con_labels)
+
+
+def test_freeze_mutable_roundtrip_with_masking() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x")
+    mask = xr.DataArray([True, False, True, False, True], dims=["i"])
+    m.add_constraints(x.where(mask) >= 0, name="c")
+    frozen = m.constraints["c"]
+    mc = frozen.mutable()
+    refrozen = Constraint.from_mutable(mc, frozen._cindex)
+    assert_equal(frozen.labels, refrozen.labels)
+    assert_equal(frozen.rhs, refrozen.rhs)
+    assert frozen.ncons == refrozen.ncons == 3
+
+
+def test_from_mutable_mixed_signs_raises() -> 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)
+
+
+def test_variable_label_index(m: Model) -> None:
+    li = m.variables.label_index
+    assert li.n_active_vars > 0
+    assert len(li.vlabels) == li.n_active_vars
+    assert li.label_to_pos.shape[0] == m._xCounter
+    for lbl in li.vlabels:
+        assert li.label_to_pos[lbl] >= 0
+    assert (li.label_to_pos[li.vlabels] == np.arange(li.n_active_vars)).all()
+
+
+def test_variable_label_index_invalidation(m: Model) -> None:
+    li = m.variables.label_index
+    old_vlabels = li.vlabels.copy()
+    m.add_variables(name="w")
+    li.invalidate()
+    assert len(li.vlabels) > len(old_vlabels)
+
+
+def test_to_matrix_with_rhs(m: Model) -> None:
+    c = m.constraints["c"]
+    li = m.variables.label_index
+    csr, con_labels, b, sense = c.to_matrix_with_rhs(li)
+    assert csr.shape[0] == len(con_labels)
+    assert csr.shape[0] == len(b)
+    assert csr.shape[0] == len(sense)
+    assert all(s in ("<", ">", "=") for s in sense)
+    np.testing.assert_array_equal(b, c._rhs)
+
+
+def test_to_matrix_with_rhs_mutable(m: Model) -> None:
+    mc = m.constraints["c"].mutable()
+    li = m.variables.label_index
+    csr, con_labels, b, sense = mc.to_matrix_with_rhs(li)
+    assert csr.shape[0] == len(con_labels)
+    assert csr.shape[0] == len(b)
+    assert csr.shape[0] == len(sense)
+
+
+def test_constraint_repr_shows_variable_names(m: Model) -> None:
+    c = m.constraints["c"]
+    r = repr(c)
+    assert "x" in r