From 37539edfaf3a26021e66f22796f75c1b1490a2e0 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Mon, 23 Mar 2026 17:13:05 +0100 Subject: [PATCH 1/2] perf: direct CSR-to-LP writer for frozen constraints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Override Constraint.to_polars() to expand CSR data directly into a polars DataFrame, bypassing the expensive mutable() → xarray Dataset reconstruction. Also override iterate_slices() to yield CSR row-batches instead of relying on xarray's isel(). Move eliminate_zeros() to freeze time (from_mutable) so the cleanup happens once rather than on every to_polars() call. LP write is now 20-40% faster than master across all benchmark models. Co-Authored-By: Claude Opus 4.6 (1M context) --- linopy/constraints.py | 63 ++++++++++++++++++++++++++++++++++++++++--- test/test_io.py | 27 +++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index fae5a03a..a5869176 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -892,9 +892,65 @@ def mutable(self) -> MutableConstraint: """Convert to a MutableConstraint.""" return MutableConstraint(self.data, self._model, self._name) - def to_polars(self) -> Any: - """Convert to polars DataFrame — delegates to mutable().""" - return self.mutable().to_polars() + def to_polars(self) -> pl.DataFrame: + """Convert frozen constraint to polars DataFrame directly from CSR.""" + csr = self._csr + if csr.nnz == 0: + return pl.DataFrame( + schema={ + "labels": pl.Int64, + "coeffs": pl.Float64, + "vars": pl.Int64, + "sign": pl.Enum(["=", "<=", ">="]), + "rhs": pl.Float64, + } + ) + + rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) + vlabels = self._model.variables.label_index.vlabels + + return pl.DataFrame( + { + "labels": self._con_labels[rows], + "coeffs": csr.data, + "vars": vlabels[csr.indices], + "rhs": self._rhs[rows], + } + ).with_columns( + pl.lit(self._sign).cast(pl.Enum(["=", "<=", ">="])).alias("sign") + )[["labels", "coeffs", "vars", "sign", "rhs"]] + + def iterate_slices( + self, + slice_size: int | None = 2_000_000, + slice_dims: list | None = None, + ) -> Iterator[Constraint]: + """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 + yield Constraint( + csr=self._csr[batch_start:batch_end], + con_labels=self._con_labels[batch_start:batch_end], + rhs=self._rhs[batch_start:batch_end], + sign=self._sign, + coords=self._coords, + model=self._model, + name=self._name, + ) + batch_start = batch_end + if batch_start >= n: + break @classmethod def from_mutable( @@ -913,6 +969,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 3d269636..70ffde20 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -390,3 +390,30 @@ 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() From 01d62cf754d916bebc1f4778410b89e3043bbec9 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Mon, 30 Mar 2026 11:18:01 +0200 Subject: [PATCH 2/2] fix: handle mixed per-row signs in CSR-to-LP writer When _sign is a numpy array (per-row signs from from_rule with mixed <=/>=/= constraints), expand it per-nonzero via _sign[rows] instead of using pl.lit() which only works for scalar strings. Also slice _sign in iterate_slices when it's an array. Add test for frozen mixed-sign constraint LP output equivalence. Co-Authored-By: Claude Opus 4.6 (1M context) --- linopy/constraints.py | 37 +++++++++++++++++++++++-------------- test/test_io.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 14 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 0de32f57..b2bbab6a 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -937,13 +937,14 @@ def mutable(self) -> Constraint: def to_polars(self) -> pl.DataFrame: """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": pl.Enum(["=", "<=", ">="]), + "sign": sign_dtype, "rhs": pl.Float64, } ) @@ -951,22 +952,25 @@ def to_polars(self) -> pl.DataFrame: rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) vlabels = self._model.variables.label_index.vlabels - return pl.DataFrame( - { - "labels": self._con_labels[rows], - "coeffs": csr.data, - "vars": vlabels[csr.indices], - "rhs": self._rhs[rows], - } - ).with_columns( - pl.lit(self._sign).cast(pl.Enum(["=", "<=", ">="])).alias("sign") - )[["labels", "coeffs", "vars", "sign", "rhs"]] + 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[Constraint]: + ) -> Iterator[CSRConstraint]: """Yield row-batched sub-Constraints without Dataset reconstruction.""" nnz = self._csr.nnz if slice_size is None or nnz <= slice_size: @@ -981,11 +985,16 @@ def iterate_slices( batch_end = max(batch_end, batch_start + 1) if batch_end >= n: batch_end = n - yield Constraint( + 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=self._sign, + sign=sign, coords=self._coords, model=self._model, name=self._name, diff --git a/test/test_io.py b/test/test_io.py index 2ac3eb07..c534fd17 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -474,3 +474,36 @@ def test_to_file_lp_frozen_vs_mutable(tmp_path: Path) -> None: 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()