diff --git a/linopy/constraints.py b/linopy/constraints.py index b2c8b372..b2bbab6a 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -935,8 +935,73 @@ def mutable(self) -> Constraint: return Constraint(self.data, self._model, self._name) def to_polars(self) -> pl.DataFrame: - """Convert to polars DataFrame — delegates to mutable().""" - return self.mutable().to_polars() + """Convert frozen constraint to polars DataFrame directly from CSR.""" + csr = self._csr + sign_dtype = pl.Enum(["=", "<=", ">="]) + if csr.nnz == 0: + return pl.DataFrame( + schema={ + "labels": pl.Int64, + "coeffs": pl.Float64, + "vars": pl.Int64, + "sign": sign_dtype, + "rhs": pl.Float64, + } + ) + + rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) + vlabels = self._model.variables.label_index.vlabels + + data: dict[str, Any] = { + "labels": self._con_labels[rows], + "coeffs": csr.data, + "vars": vlabels[csr.indices], + "rhs": self._rhs[rows], + } + if isinstance(self._sign, str): + data["sign"] = pl.Series( + "sign", [self._sign], dtype=sign_dtype + ).new_from_index(0, len(rows)) + else: + data["sign"] = pl.Series("sign", self._sign[rows], dtype=sign_dtype) + return pl.DataFrame(data)[["labels", "coeffs", "vars", "sign", "rhs"]] + + def iterate_slices( + self, + slice_size: int | None = 2_000_000, + slice_dims: list | None = None, + ) -> Iterator[CSRConstraint]: + """Yield row-batched sub-Constraints without Dataset reconstruction.""" + nnz = self._csr.nnz + if slice_size is None or nnz <= slice_size: + yield self + return + + n = self._csr.shape[0] + cumulative = np.cumsum(np.diff(self._csr.indptr)) + batch_start = 0 + for batch_end_nnz in range(slice_size, nnz + slice_size, slice_size): + batch_end = int(np.searchsorted(cumulative, batch_end_nnz, side="right")) + batch_end = max(batch_end, batch_start + 1) + if batch_end >= n: + batch_end = n + sign = ( + self._sign + if isinstance(self._sign, str) + else self._sign[batch_start:batch_end] + ) + yield CSRConstraint( + csr=self._csr[batch_start:batch_end], + con_labels=self._con_labels[batch_start:batch_end], + rhs=self._rhs[batch_start:batch_end], + sign=sign, + coords=self._coords, + model=self._model, + name=self._name, + ) + batch_start = batch_end + if batch_start >= n: + break @classmethod def from_mutable( @@ -955,6 +1020,7 @@ def from_mutable( """ label_index = con.model.variables.label_index csr, con_labels = con.to_matrix(label_index) + csr.eliminate_zeros() coords = [con.indexes[d] for d in con.coord_dims] # Build active_mask aligned with con_labels (rows in csr) # Use same filter as to_matrix: label != -1 AND at least one var != -1 diff --git a/test/test_io.py b/test/test_io.py index 119b4432..c534fd17 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -447,3 +447,63 @@ def test_to_file_lp_mixed_sign_constraints(tmp_path: Path) -> None: assert "<=" in content assert ">=" in content assert "=" in content + + +def test_to_file_lp_frozen_vs_mutable(tmp_path: Path) -> None: + """Test that frozen and mutable constraints produce identical LP output.""" + m_frozen = Model() + N = np.arange(5) + x = m_frozen.add_variables(coords=[N], name="x") + y = m_frozen.add_variables(coords=[N], name="y") + m_frozen.add_constraints(x + y <= 10, name="upper") + m_frozen.add_constraints(x >= 1, name="lower") + m_frozen.add_constraints(2 * x + y == 8, name="eq") + m_frozen.add_objective(x.sum() + 2 * y.sum()) + + m_mutable = Model() + x2 = m_mutable.add_variables(coords=[N], name="x") + y2 = m_mutable.add_variables(coords=[N], name="y") + m_mutable.add_constraints(x2 + y2 <= 10, name="upper", freeze=False) + m_mutable.add_constraints(x2 >= 1, name="lower", freeze=False) + m_mutable.add_constraints(2 * x2 + y2 == 8, name="eq", freeze=False) + m_mutable.add_objective(x2.sum() + 2 * y2.sum()) + + fn_frozen = tmp_path / "frozen.lp" + fn_mutable = tmp_path / "mutable.lp" + m_frozen.to_file(fn_frozen) + m_mutable.to_file(fn_mutable) + + assert fn_frozen.read_text() == fn_mutable.read_text() + + +def test_to_file_lp_frozen_mixed_sign(tmp_path: Path) -> None: + """Test LP writing for frozen constraint with per-row signs.""" + m_frozen = Model() + N = pd.RangeIndex(4, name="i") + x = m_frozen.add_variables(coords=[N], name="x") + + def bound(m: Model, i: int) -> object: + if i % 2: + return x.at[i] >= i + return x.at[i] <= 10 + + m_frozen.add_constraints(bound, coords=[N], name="mixed", freeze=True) + m_frozen.add_objective(x.sum()) + + m_mutable = Model() + x2 = m_mutable.add_variables(coords=[N], name="x") + + def bound2(m: Model, i: int) -> object: + if i % 2: + return x2.at[i] >= i + return x2.at[i] <= 10 + + m_mutable.add_constraints(bound2, coords=[N], name="mixed", freeze=False) + m_mutable.add_objective(x2.sum()) + + fn_frozen = tmp_path / "frozen_mixed.lp" + fn_mutable = tmp_path / "mutable_mixed.lp" + m_frozen.to_file(fn_frozen) + m_mutable.to_file(fn_mutable) + + assert fn_frozen.read_text() == fn_mutable.read_text()