Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 68 additions & 2 deletions linopy/constraints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand Down
60 changes: 60 additions & 0 deletions test/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()