From 5680c19d0e879f0aa7499b87930502f439eacfd3 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 08:19:57 +0200 Subject: [PATCH 1/7] perf(io,solvers): vectorised MOSEK build, faster LP writer, HiGHS names via getModel MOSEK bound keys and names are built vectorised. The polars LP writer skips identity scaling, avoids re-sorting frozen rows and builds fewer string branches; output is byte-identical. HiGHS names, when requested, are attached through getModel/passModel so the Hessian survives. --- linopy/io.py | 147 +++++++++++++++++++++++++++------------------- linopy/solvers.py | 67 ++++++++++----------- test/test_io.py | 87 ++++++++++++++++++++++++++- 3 files changed, 201 insertions(+), 100 deletions(-) diff --git a/linopy/io.py b/linopy/io.py index 205f100b..80568f18 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -119,23 +119,30 @@ def _lookup_positive_labels(lookup: np.ndarray, labels: np.ndarray) -> np.ndarra return values +def _non_identity(lookup: np.ndarray) -> np.ndarray | None: + """Return the scaling lookup, or None if all factors are 1.""" + return None if (lookup == 1).all() else lookup + + def _scale_objective_dataframe( - df: pl.DataFrame, variable_scaling: np.ndarray, objective_scaling: float + df: pl.DataFrame, variable_scaling: np.ndarray | None, objective_scaling: float ) -> pl.DataFrame: """Apply column scaling and row-like objective scaling to objective terms.""" if df.is_empty(): return df - if "vars" in df.columns: - scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) - else: - scale1 = _lookup_positive_labels(variable_scaling, df["vars1"].to_numpy()) - scale2 = _lookup_positive_labels(variable_scaling, df["vars2"].to_numpy()) - scales = scale1 * scale2 + if variable_scaling is not None: + if "vars" in df.columns: + scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) + else: + scale1 = _lookup_positive_labels(variable_scaling, df["vars1"].to_numpy()) + scale2 = _lookup_positive_labels(variable_scaling, df["vars2"].to_numpy()) + scales = scale1 * scale2 + df = df.with_columns(pl.col("coeffs") / pl.Series(scales)) - return df.with_columns( - (pl.col("coeffs") / pl.Series(scales) * objective_scaling).alias("coeffs") - ) + if objective_scaling != 1: + df = df.with_columns(pl.col("coeffs") * objective_scaling) + return df def _scale_bounds_dataframe( @@ -154,19 +161,22 @@ def _scale_bounds_dataframe( def _scale_constraint_dataframe( df: pl.DataFrame, - variable_scaling: np.ndarray, - constraint_scaling: np.ndarray, + variable_scaling: np.ndarray | None, + constraint_scaling: np.ndarray | None, ) -> pl.DataFrame: """Divide by column scaling and multiply by row scaling.""" - if df.is_empty(): - return df - row_scales = _lookup_positive_labels(constraint_scaling, df["labels"].to_numpy()) - var_scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) - row_scale_series = pl.Series(row_scales) - return df.with_columns( - (pl.col("coeffs") / pl.Series(var_scales) * row_scale_series).alias("coeffs"), - (pl.col("rhs") * row_scale_series).alias("rhs"), - ) + if variable_scaling is not None: + var_scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) + df = df.with_columns(pl.col("coeffs") / pl.Series(var_scales)) + if constraint_scaling is not None: + row_scales = _lookup_positive_labels( + constraint_scaling, df["labels"].to_numpy() + ) + row_scale_series = pl.Series(row_scales) + df = df.with_columns( + pl.col("coeffs") * row_scale_series, pl.col("rhs") * row_scale_series + ) + return df def format_coord(coord: str) -> str: @@ -281,6 +291,7 @@ def objective_write_quadratic_terms( def objective_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, progress: bool = False, explicit_coordinate_names: bool = False, ) -> None: @@ -293,7 +304,6 @@ def objective_to_file( print_variable, _ = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) sense = m.objective.sense f.write(f"{sense}\n\nobj:\n\n".encode()) @@ -325,6 +335,7 @@ def _binary_has_nondefault_bounds(var: Variable) -> bool: def bounds_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, progress: bool = False, slice_size: int = 2_000_000, explicit_coordinate_names: bool = False, @@ -348,7 +359,6 @@ def bounds_to_file( print_variable, _ = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) f.write(b"\n\nbounds\n\n") if progress: @@ -362,7 +372,8 @@ def bounds_to_file( var = m.variables[name] for var_slice in var.iterate_slices(slice_size): df = var_slice.to_polars() - df = _scale_bounds_dataframe(df, variable_scaling) + if variable_scaling is not None: + df = _scale_bounds_dataframe(df, variable_scaling) columns = [ *signed_number(pl.col("lower")), @@ -557,6 +568,7 @@ def sos_to_file( def indicator_constraints_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray, explicit_coordinate_names: bool = False, ) -> None: """ @@ -574,7 +586,6 @@ def indicator_constraints_to_file( print_variable_scalar, _ = get_printers_scalar( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) for con in m.constraints.indicator.data.values(): ic_data = con.data @@ -619,6 +630,8 @@ def indicator_constraints_to_file( def constraints_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, + constraint_scaling: np.ndarray | None, progress: bool = False, lazy: bool = False, slice_size: int = 2_000_000, @@ -631,8 +644,7 @@ def constraints_to_file( print_variable, print_constraint = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) - constraint_scaling = constraint_scaling_lookup(m) + scaled = variable_scaling is not None or constraint_scaling is not None f.write(b"\n\ns.t.\n\n") names = list(regular) @@ -643,48 +655,53 @@ def constraints_to_file( colour=TQDM_COLOR, ) - # to make this even faster, we can use polars expression - # https://docs.pola.rs/user-guide/expressions/plugins/#output-data-types for name in names: con = regular[name] for con_slice in con.iterate_slices(slice_size): df = con_slice.to_polars() - df = _scale_constraint_dataframe(df, variable_scaling, constraint_scaling) - if df.height == 0: continue + if scaled: + df = _scale_constraint_dataframe( + df, variable_scaling, constraint_scaling + ) + if not df["labels"].is_sorted(): + df = df.sort("labels", maintain_order=True) + _write_constraint_rows(df, f, print_constraint, print_variable) + - # Sort by labels and mark first/last occurrences - df = df.sort("labels").with_columns( +def _write_constraint_rows( + df: pl.DataFrame, + f: BufferedWriter, + print_constraint: Callable, + print_variable: Callable, +) -> None: + labels = df["labels"].to_numpy() + first = np.empty(len(labels), dtype=bool) + first[0] = True + np.not_equal(labels[1:], labels[:-1], out=first[1:]) + last = np.append(first[1:], True) + columns = [ + pl.when(pl.Series(first)).then( + pl.concat_str( + [*print_constraint(pl.col("labels")), pl.lit(":\n")], + ignore_nulls=True, + ) + ), + *signed_number(pl.col("coeffs")), + *print_variable(pl.col("vars")), + pl.when(pl.Series(last)).then( + pl.concat_str( [ - pl.col("labels").is_first_distinct().alias("is_first_in_group"), - (pl.col("labels") != pl.col("labels").shift(-1)) - .fill_null(True) - .alias("is_last_in_group"), + pl.lit("\n"), + pl.col("sign"), + pl.lit(" "), + pl.col("rhs").cast(pl.String), ] ) - - row_labels = print_constraint(pl.col("labels")) - col_labels = print_variable(pl.col("vars")) - columns = [ - pl.when(pl.col("is_first_in_group")).then(row_labels[0]), - pl.when(pl.col("is_first_in_group")).then(row_labels[1]), - pl.when(pl.col("is_first_in_group")).then(pl.lit(":\n")).alias(":"), - *signed_number(pl.col("coeffs")), - col_labels[0], - col_labels[1], - pl.when(pl.col("is_last_in_group")).then(pl.lit("\n")), - pl.when(pl.col("is_last_in_group")).then(pl.col("sign")), - pl.when(pl.col("is_last_in_group")).then(pl.lit(" ")), - pl.when(pl.col("is_last_in_group")).then(pl.col("rhs").cast(pl.String)), - ] - - _format_and_write(df, columns, f) - - # in the future, we could use lazy dataframes when they support appending - # tp existent files - # formatted = df.lazy().select(pl.concat_str(columns, ignore_nulls=True)) - # formatted.sink_csv(f, **kwargs) + ), + ] + _format_and_write(df, columns, f) def to_lp_file( @@ -697,13 +714,21 @@ def to_lp_file( ) -> None: with open(fn, mode="wb") as f: start = time.time() + variable_scaling = variable_scaling_lookup(m) + active_variable_scaling = _non_identity(variable_scaling) objective_to_file( - m, f, progress=progress, explicit_coordinate_names=explicit_coordinate_names + m, + f, + active_variable_scaling, + progress=progress, + explicit_coordinate_names=explicit_coordinate_names, ) constraints_to_file( m, f=f, + variable_scaling=active_variable_scaling, + constraint_scaling=_non_identity(constraint_scaling_lookup(m)), progress=progress, slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, @@ -711,11 +736,13 @@ def to_lp_file( indicator_constraints_to_file( m, f=f, + variable_scaling=variable_scaling, explicit_coordinate_names=explicit_coordinate_names, ) bounds_to_file( m, f=f, + variable_scaling=active_variable_scaling, progress=progress, slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, diff --git a/linopy/solvers.py b/linopy/solvers.py index 600e3d6d..4df516d3 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -1726,11 +1726,11 @@ def _build_solver_model( print_variables, print_constraints = linopy.io.get_printers_scalar( model, explicit_coordinate_names=explicit_coordinate_names ) - lp = h.getLp() - lp.col_names_ = print_variables(M.vlabels) + mdl = h.getModel() + mdl.lp_.col_names_ = print_variables(M.vlabels) if len(M.clabels): - lp.row_names_ = print_constraints(M.clabels) - h.passModel(lp) + mdl.lp_.row_names_ = print_constraints(M.clabels) + h.passModel(mdl) Q = M.Q if Q is not None: @@ -3524,56 +3524,49 @@ def _build_solver_model( np.arange(0, len(labels)), "%0", [len(labels)], None, [0], labels ) - bkx = [ - ( - ( - (mosek.boundkey.ra if lb < ub else mosek.boundkey.fx) - if ub < np.inf - else mosek.boundkey.lo - ) - if (lb > -np.inf) - else (mosek.boundkey.up if (ub < np.inf) else mosek.boundkey.fr) - ) - for (lb, ub) in zip(M.lb, M.ub) - ] - blx = [b if b > -np.inf else 0.0 for b in M.lb] - bux = [b if b < np.inf else 0.0 for b in M.ub] - task.putvarboundslice(0, model.nvars, bkx, blx, bux) + bk = mosek.boundkey + keys = np.array([bk.fr, bk.lo, bk.up, bk.fx, bk.ra], dtype=object) + lb_fin = M.lb > -np.inf + ub_fin = M.ub < np.inf + bkx_code = np.select( + [lb_fin & ub_fin & (M.lb < M.ub), lb_fin & ub_fin, lb_fin, ub_fin], + [4, 3, 1, 2], + default=0, + ) + blx = np.where(lb_fin, M.lb, 0.0) + bux = np.where(ub_fin, M.ub, 0.0) + task.putvarboundslice(0, model.nvars, keys[bkx_code].tolist(), blx, bux) if len(model.binaries.labels) + len(model.integers.labels) > 0: - idx = [i for (i, v) in enumerate(M.vtypes) if v in ["B", "I"]] + idx = np.flatnonzero(np.isin(M.vtypes, ["B", "I"])).astype(np.int32) task.putvartypelist(idx, [mosek.variabletype.type_int] * len(idx)) if len(model.constraints) > 0: if set_names: names = print_constraints(M.clabels) - for i, n in enumerate(names): - task.putconname(i, n) - bkc = [ - ( - (mosek.boundkey.up if b < np.inf else mosek.boundkey.fr) - if s == "<" - else ( - (mosek.boundkey.lo if b > -np.inf else mosek.boundkey.up) - if s == ">" - else mosek.boundkey.fx - ) + task.generateconnames( + np.arange(0, len(names)), "%0", [len(names)], None, [0], names ) - for s, b in zip(M.sense, M.b) - ] - blc = [b if b > -np.inf else 0.0 for b in M.b] - buc = [b if b < np.inf else 0.0 for b in M.b] + b_lo = M.b > -np.inf + b_up = M.b < np.inf + leq = M.sense == "<" + geq = M.sense == ">" + bkc_code = np.select( + [leq & b_up, leq, geq & b_lo, geq], [2, 0, 1, 2], default=3 + ) + blc = np.where(b_lo, M.b, 0.0) + buc = np.where(b_up, M.b, 0.0) if M.A is not None: A = M.A.tocsr() task.putarowslice( 0, model.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data ) - task.putconboundslice(0, model.ncons, bkc, blc, buc) + task.putconboundslice(0, model.ncons, keys[bkc_code].tolist(), blc, buc) if M.Q is not None: Q = (0.5 * tril(M.Q + M.Q.transpose())).tocoo() task.putqobj(Q.row, Q.col, Q.data) - task.putclist(list(np.arange(model.nvars)), M.c) + task.putclist(np.arange(model.nvars, dtype=np.int32), M.c) if model.objective.sense == "max": task.putobjsense(mosek.objsense.maximize) diff --git a/test/test_io.py b/test/test_io.py index 7a55916d..aa66f90d 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -20,6 +20,7 @@ from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf from linopy.constants import FACTOR_DIM +from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import CONTAINER_ORDER_ATTR, signed_number from linopy.testing import assert_exprequal, assert_model_equal @@ -694,9 +695,12 @@ def test_to_gurobipy(model: Model) -> None: @pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") -def test_to_highspy(model: Model) -> None: - h = model.to_highspy() - assert h.getLp().num_col_ > 0 +@pytest.mark.parametrize("set_names", [True, False]) +def test_to_highspy(model: Model, set_names: bool) -> None: + lp = model.to_highspy(set_names=set_names).getLp() + assert lp.num_col_ > 0 + assert len(lp.col_names_) == (lp.num_col_ if set_names else 0) + assert len(lp.row_names_) == (lp.num_row_ if set_names else 0) @pytest.mark.skipif("mosek" not in available_solvers, reason="Mosek not installed") @@ -799,6 +803,16 @@ def test_model_set_names_in_solver_io(model: Model) -> None: assert model.objective.value == pytest.approx(expected_obj) +@pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") +def test_highs_direct_qp_keeps_hessian_with_names() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(2, name="i")], name="x") + m.add_constraints(x.sum() >= 1) + m.add_objective((x * x).sum() + x.sum()) + m.solve(solver_name="highs", io_api="direct", set_names=True) + assert m.objective.value == pytest.approx(1.5) + + def test_to_blocks(tmp_path: Path) -> None: m: Model = Model() @@ -1084,3 +1098,70 @@ def bound2(m: Model, i: int) -> object: m_mutable.to_file(fn_mutable) assert fn_frozen.read_text() == fn_mutable.read_text() + + +def _lp_constraint_section(m: Model, path: Path, **kwargs: int) -> str: + m.to_file(path, progress=False, **kwargs) + return path.read_text().split("s.t.\n\n")[1].split("\n\nbounds")[0] + + +@pytest.mark.parametrize("freeze", [True, False]) +@pytest.mark.parametrize( + "to_file_kwargs", + [ + pytest.param({}, id="default-slices"), + pytest.param({"slice_size": 1}, id="slice-1"), + ], +) +@pytest.mark.parametrize( + ("scaled", "expected"), + [ + pytest.param( + True, + "c0:\n+2.0 x0\n-4.0 x2\n<= 3.0\n" + "c1:\n+1.0 x1\n-4.0 x3\n<= 3.0\n" + "c2:\n+2.0 x0\n+1.0 x1\n>= -0.0\n", + id="scaled", + ), + pytest.param( + False, + "c0:\n+1.0 x0\n-2.0 x2\n<= 1.5\n" + "c1:\n+1.0 x1\n-2.0 x3\n<= 1.5\n" + "c2:\n+1.0 x0\n+1.0 x1\n>= -0.0\n", + id="unscaled", + ), + ], +) +def test_to_file_lp_constraint_section( + tmp_path: Path, + freeze: bool, + to_file_kwargs: dict[str, int], + scaled: bool, + expected: str, +) -> None: + m = Model() + i = pd.RangeIndex(2, name="i") + row_scaling = 2.0 if scaled else 1.0 + x = m.add_variables(coords=[i], name="x", scaling=[1.0, 2.0] if scaled else 1.0) + y = m.add_variables(coords=[i], name="y") + m.add_constraints(x - 2 * y <= 1.5, name="a", freeze=freeze, scaling=row_scaling) + m.add_constraints(x.sum() >= -0.0, name="b", freeze=freeze, scaling=row_scaling) + m.add_objective(x.sum()) + + fn = tmp_path / "constraints.lp" + assert _lp_constraint_section(m, fn, **to_file_kwargs) == expected + + +def test_to_file_lp_unsorted_constraint_labels(tmp_path: Path) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + x = m.add_variables(coords=[i], name="x") + y = m.add_variables(coords=[i], name="y") + m.add_constraints(x + 2 * y <= 1, name="a", freeze=False) + m.add_objective(x.sum()) + expected = _lp_constraint_section(m, tmp_path / "sorted.lp") + + con = m.constraints["a"] + m.constraints.data["a"] = Constraint(con.data.isel(i=slice(None, None, -1)), m, "a") + assert not m.constraints["a"].to_polars()["labels"].is_sorted() + assert _lp_constraint_section(m, tmp_path / "unsorted.lp") == expected From 772804c334ca6de9a1206b5b1121fb699e3d1255 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:44:10 +0200 Subject: [PATCH 2/7] perf(csr): N-ary merge, matrix-free expression solution, nnz-sized contraction chunks Merge concatenates all CSR operands once instead of folding pairwise. Expression.solution maps labels directly and evaluates CSR rows without densifying or rebuilding the constraint matrix. Contraction chunks are sized by nonzeros instead of a fixed row count. --- linopy/csr.py | 45 ++++++++++++++++------------- linopy/expressions.py | 38 +++++++++++++++++++------ test/test_csr.py | 66 +++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 119 insertions(+), 30 deletions(-) diff --git a/linopy/csr.py b/linopy/csr.py index 2222bd50..2c023fc0 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -28,6 +28,7 @@ from __future__ import annotations +import functools import operator from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field, replace @@ -51,8 +52,8 @@ from linopy.expressions import LinearExpression from linopy.model import Model -CONTRACTION_CHUNK = 64 -"""Kept-axis block size of the chunked Kronecker product in ``contracted``.""" +CONTRACTION_CHUNK_NNZ = 1 << 22 +"""Target nonzeros per kept-axis block of the chunked Kronecker product in ``contracted``.""" AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] """Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" @@ -572,25 +573,30 @@ def same_grid(self, other: CSRLinearExpression) -> bool: """Whether both live on the same cells, auxiliary coordinates aside.""" return self.grid.same_layout(other.grid) - def added(self, other: CSRLinearExpression) -> CSRLinearExpression: + def added(self, *others: CSRLinearExpression) -> CSRLinearExpression: """ - Sparse matrix addition == merge along the term dimension. Goes through - COO so explicit zero coefficients survive (scipy's ``+`` drops them), + Sparse matrix addition == merge along the term dimension, over any + number of operands on the same grid in one COO pass. Goes through COO + so explicit zero coefficients survive (scipy's ``+`` drops them), keeping a cell with only zero-coefficient terms distinguishable from - an empty cell. A cell absent in either operand is absent in the sum - and carries no terms. Auxiliary coordinates propagate and conflicting - ones raise (§11). - """ - const = self.const + other.const - a, b = self.csr.tocoo(), other.csr.tocoo() - shape = (self.n_cells, max(a.shape[1], b.shape[1])) - rows = np.concatenate([a.coords[0], b.coords[0]]) - cols = np.concatenate([a.coords[1], b.coords[1]]) - data = np.concatenate([a.data, b.data]) + an empty cell. A cell absent in any operand is absent in the sum and + carries no terms. Auxiliary coordinates propagate, earlier operands + taking precedence, and conflicting ones raise (§11). + """ + parts = (self, *others) + const = functools.reduce(np.add, (p.const for p in parts)) + coos = [p.csr.tocoo() for p in parts] + shape = (self.n_cells, max(c.shape[1] for c in coos)) + rows = np.concatenate([c.coords[0] for c in coos]) + cols = np.concatenate([c.coords[1] for c in coos]) + data = np.concatenate([c.data for c in coos]) present = ~np.isnan(const)[rows] csr = coo_to_csr(data[present], rows[present], cols[present], shape, self.model) - enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in (self, other)]) - grid = replace(self.grid, aux=other.grid.aux | self.grid.aux) + enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in parts]) + aux: AuxCoords = {} + for p in reversed(parts): + aux |= p.grid.aux + grid = replace(self.grid, aux=aux) return replace(self, csr=csr, const=const, grid=grid) def contracted( @@ -609,7 +615,7 @@ def contracted( off its ``name``. The result lives on the kept grid dims followed by ``new_indexes`` and is ``kron(I_kept, matrix.T) @ csr``, evaluated in chunks of the kept axis - so the operator never grows with the kept size. + sized by nonzero count so the operator never grows with the kept size. The result is in compact canonical form: duplicate variables summed, terms label-ordered and explicit zeros pruned -- unlike :meth:`added`, @@ -630,7 +636,8 @@ def contracted( kept_grid = source.grid.reordered(kept) n_kept = kept_grid.size const = np.nan_to_num(source.const) - chunk = min(CONTRACTION_CHUNK, n_kept) + nnz_per_kept = matrix.nnz + source.csr.nnz // max(n_kept, 1) + chunk = min(max(CONTRACTION_CHUNK_NNZ // max(nnz_per_kept, 1), 1), n_kept) operator = scipy.sparse.kron( scipy.sparse.eye_array(chunk), matrix.T, format="csr" ) diff --git a/linopy/expressions.py b/linopy/expressions.py index d2f08463..bc811abd 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1738,12 +1738,14 @@ def mask(self) -> None: return None @has_optimized_model - def _map_solution(self) -> DataArray: + def _label_solution(self, labels: np.ndarray) -> np.ndarray: """ - Replace variable labels by solution values. + Solution values indexed by variable label, with a trailing NaN that + label ``-1`` reads. + + Raises if ``labels`` reference variables missing from the model. """ m = self.model - labels = self.vars.values known = np.append(m.variables.label_index.label_to_pos != -1, True) if not known[labels].all(): raise KeyError("Expression references variables missing from the model.") @@ -1751,8 +1753,15 @@ def _map_solution(self) -> DataArray: for var in m.variables.data.values(): sol[var.labels.values] = var.solution.values sol[-1] = np.nan - values = sol[labels] - return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) + return sol + + def _map_solution(self) -> DataArray: + """ + Replace variable labels by solution values. + """ + labels = self.vars + values = self._label_solution(labels.values)[labels.values] + return xr.DataArray(values, dims=labels.dims, coords=labels.coords) @property def solution(self) -> DataArray: @@ -2512,6 +2521,21 @@ def has_terms(self) -> DataArray: present & (np.diff(csr.csr.indptr) > 0), name="has_terms" ) + @property + def solution(self) -> DataArray: + """ + Get the optimal values of the expression. + + The function raises an error in case no model is set as a + reference or the model is not optimized. + """ + csr = self._csr + if csr is None: + return super().solution + sol = self._label_solution(csr.csr.indices)[: csr.csr.shape[1]] + sol = np.nan_to_num(sol) + return csr.grid.dataarray(csr.csr @ sol + csr.const, name="solution") + def _combined_with_constant( self, self_const: DataArray, @@ -3769,9 +3793,7 @@ def _try_csr_merge( return None csrs = aligned - combined = csrs[0] - for csr in csrs[1:]: - combined = combined.added(csr) + combined = csrs[0].added(*csrs[1:]) return LinearExpression._from_csr(combined, exprs[0].model) diff --git a/test/test_csr.py b/test/test_csr.py index 513204af..f5fcbb40 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -831,6 +831,62 @@ def test_reindex_stays_csr_and_matches_dense(indexers: dict) -> None: assert_linequal(sparse, dense) +def case_twins( + build: Callable[[Case], LinearExpression], +) -> Callable[[], tuple[LinearExpression, LinearExpression]]: + def twins() -> tuple[LinearExpression, LinearExpression]: + c1, c2 = twin_models() + return build(c1), build(c2) + + return twins + + +def keyed_observed_twins() -> tuple[LinearExpression, LinearExpression]: + dense, sparse = (keyed_model(sparse=s)[1] for s in (False, True)) + keys = ["period", "season"] + return dense.groupby(keys).sum(observed=True), sparse.groupby(keys).sum( + observed=True + ) + + +@pytest.mark.parametrize("nan_every_other", [False, True], ids=["full", "nan"]) +@pytest.mark.parametrize( + "twins", + [ + case_twins(lambda c: c.balance_lhs() + 2.0), + case_twins(lambda c: c.gen_sum().reindex(bus=["bus3", "bus9", "bus0"])), + case_twins( + lambda c: linopy.merge( + cross_grid_parts(c) + cross_grid_parts(c, ("line3", "line4"))[1:], + join="outer", + fill_value=linopy.ABSENT, + cls=LinearExpression, + ) + ), + keyed_observed_twins, + ], + ids=["composed", "absent_cell", "absent_merge", "aux_coords"], +) +def test_csr_solution_matches_dense( + twins: Callable[[], tuple[LinearExpression, LinearExpression]], + nan_every_other: bool, +) -> None: + require_v1() + dense, sparse = twins() + assert sparse._csr is not None + rng = np.random.default_rng(0) + for name, var in dense.model.variables.items(): + values = rng.uniform(-1, 1, var.shape) + if nan_every_other: + values.ravel()[::2] = np.nan + for m in (dense.model, sparse.model): + m.variables[name].solution = xr.DataArray(values, coords=var.labels.coords) + m._status = "ok" + sol = sparse.solution + assert sparse._csr is not None + xr.testing.assert_allclose(sol, dense.solution) + + def test_reindex_falls_back_to_dense_for_unsupported_kwargs() -> None: require_v1() c1, c2 = twin_models() @@ -977,16 +1033,20 @@ def test_transposed_grid_exact_merge_stays_csr_and_matches_dense() -> None: assert_terms_equal(res, linopy.merge(dense, cls=LinearExpression)) +@pytest.mark.parametrize("fill_value", [None, linopy.ABSENT], ids=["fill", "absent"]) @pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) -def test_three_operand_cross_grid_merge_matches_dense(join: JoinOptions) -> None: +def test_three_operand_cross_grid_merge_matches_dense( + join: JoinOptions, fill_value: Any +) -> None: require_v1() c1, c2 = twin_models() third_lines = ("line3", "line4") sparse = cross_grid_parts(c2) + cross_grid_parts(c2, third_lines)[1:] dense = cross_grid_parts(c1) + cross_grid_parts(c1, third_lines)[1:] - res = linopy.merge(sparse, join=join, cls=LinearExpression) + kwargs: dict[str, Any] = {"join": join, "fill_value": fill_value} + res = linopy.merge(sparse, **kwargs, cls=LinearExpression) assert res._csr is not None - assert_terms_equal(res, linopy.merge(dense, join=join, cls=LinearExpression)) + assert_terms_equal(res, linopy.merge(dense, **kwargs, cls=LinearExpression)) def test_cross_grid_merge_absent_fill_matches_dense() -> None: From 3b47545e51559c4c686fd6fa25d7546335de7e6f Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:55:27 +0200 Subject: [PATCH 3/7] perf(matrices): single-pass preallocated assembly, direct frozen dual read-back, one freeze gather Constraint blocks are gathered into preallocated CSR buffers without vstack or eliminate_zeros; scaling is skipped when identity and label lookups use the label dtype. Frozen constraints receive duals on active rows directly. Mask and non-empty filtering share one gather and sense conversion is vectorised. --- linopy/common.py | 21 ++-- linopy/constraints.py | 35 ++++--- linopy/matrices.py | 232 +++++++++++++++++++++++++++--------------- linopy/model.py | 23 +++-- test/test_csr.py | 14 ++- test/test_matrices.py | 41 ++++++++ test/test_scaling.py | 5 +- 7 files changed, 254 insertions(+), 117 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 79291050..16b1da0c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -930,13 +930,16 @@ def label_to_pos(self) -> np.ndarray: """ Mapping from variable label to dense position, shape (_xCounter,). + Positions share the model's label dtype, since they never exceed a label. + Position i in the active variable array corresponds to label vlabels[i]. Masked or unused labels map to -1. """ vlabels = self.vlabels - n = self._variables.model._xCounter - label_to_pos = np.full(n, -1, dtype=np.intp) - label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp) + model = self._variables.model + dtype = model._dtypes["labels"] + label_to_pos = np.full(model._xCounter, -1, dtype=dtype) + label_to_pos[vlabels] = np.arange(len(vlabels), dtype=dtype) return label_to_pos @property @@ -969,17 +972,17 @@ def clabels(self) -> np.ndarray: for c in self._constraints.data.values() if not c.is_indicator ] - return ( - np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp) - ) + dtype = self._constraints.model._dtypes["labels"] + return np.concatenate([np.array([], dtype=dtype), *label_lists], dtype=dtype) @cached_property def label_to_pos(self) -> np.ndarray: """Mapping from constraint label to dense position, shape (_cCounter,).""" clabels = self.clabels - n = self._constraints.model._cCounter - label_to_pos = np.full(n, -1, dtype=np.intp) - label_to_pos[clabels] = np.arange(len(clabels), dtype=np.intp) + model = self._constraints.model + dtype = model._dtypes["labels"] + label_to_pos = np.full(model._cCounter, -1, dtype=dtype) + label_to_pos[clabels] = np.arange(len(clabels), dtype=dtype) return label_to_pos @property diff --git a/linopy/constraints.py b/linopy/constraints.py index 9f8466c4..fa42411c 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -950,31 +950,39 @@ def _replace(self, **changes: Any) -> CSRConstraint: return new def assign_labels( - self, cindex: int, name: str, scaling: float | DataArray = 1.0 + self, + cindex: int, + name: str, + scaling: float | DataArray = 1.0, + mask: np.ndarray | None = None, ) -> CSRConstraint: """ Return a copy labelled from ``cindex`` and named ``name``. Rows without terms are dropped, as when freezing a dense constraint; - a zero coefficient counts as a term. ``scaling`` is a scalar or a row - scaling broadcast on the grid; its distinct values are validated - without expanding a broadcast view. + a zero coefficient counts as a term. Active rows where the boolean + ``mask`` is False are dropped in the same gather. ``scaling`` is a + scalar or a row scaling broadcast on the grid; its distinct values are + validated without expanding a broadcast view. """ values = np.asarray(scaling) distinct = values[tuple(slice(None) if s else 0 for s in values.strides)] validate_scaling(distinct, "constraint scaling") - kept = self._kept(np.diff(self._csr.indptr) > 0) + keep = np.diff(self._csr.indptr) > 0 + if mask is not None: + keep &= mask + kept = self._kept(keep) csr = kept._csr if not csr.data.all(): csr = csr.copy() if csr is self._csr else csr csr.eliminate_zeros() if isinstance(scaling, DataArray): - row_scaling = kept._active_values(scaling) + row_scaling = kept.active_values(scaling) else: row_scaling = np.full(csr.shape[0], float(scaling)) return kept._replace(csr=csr, cindex=cindex, name=name, scaling=row_scaling) - def _active_values(self, values: DataArray) -> np.ndarray: + def active_values(self, values: DataArray) -> np.ndarray: """Values of ``values``, broadcast on the grid, at the active rows.""" return _take_flat( values.transpose(*self._grid.dims).values, self._active_positions @@ -1003,13 +1011,6 @@ def rows(values: Any) -> Any: binval=rows(self._binval), ) - def masked(self, mask: DataArray) -> CSRConstraint: - """ - Copy with the cells where the boolean ``mask`` is False made inactive, - without the dense rectangle. ``mask`` must lie on the constraint grid. - """ - return self._kept(self._active_values(mask).astype(bool)) - def _assign_coords(self, **coords: Any) -> CSRConstraint: """ Reassign coordinate values on the constraint, keeping the shape. @@ -1449,7 +1450,7 @@ def to_matrix_with_rhs( if isinstance(self._sign, str): sense = np.full(len(self._rhs), self._sign[0]) else: - sense = np.array([s[0] for s in self._sign]) + sense = self._sign.astype("U1") return ( self._to_positional_csr(label_index), self.active_labels(), @@ -1473,7 +1474,9 @@ def sanitize_zeros(self) -> CSRConstraint: external holders of the previous arrays (e.g. a ModelSnapshot sharing them) keep a valid baseline. """ - zeros = np.abs(self._csr.data) <= 1e-10 + data = self._csr.data + zeros = data <= 1e-10 + zeros &= data >= -1e-10 if zeros.any(): csr = self._csr.copy() csr.data[zeros] = 0 diff --git a/linopy/matrices.py b/linopy/matrices.py index 08e328e4..f7beaaf0 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -7,30 +7,40 @@ from __future__ import annotations -from functools import cached_property -from typing import TYPE_CHECKING, cast +from collections.abc import Callable +from functools import cached_property, partial +from typing import TYPE_CHECKING, NamedTuple, cast import numpy as np import scipy.sparse from numpy import ndarray from linopy import expressions -from linopy.constraints import ConstraintBase, CSRConstraint +from linopy.constraints import CSRConstraint +from linopy.csr import index_dtype if TYPE_CHECKING: + from linopy.common import VariableLabelIndex + from linopy.constraints import ConstraintBase from linopy.model import Model -def _stack(csrs: list[scipy.sparse.csr_array]) -> scipy.sparse.csr_array | None: - """Vertically stack CSR blocks, or None when there are none.""" - if not csrs: - return None - return cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) +class _RowBlock(NamedTuple): + """ + Rows of one constraint. ``positional`` returns its positional CSR; for a + frozen constraint it is built only while the rows are gathered. + """ + + positional: Callable[[], scipy.sparse.csr_array] + n_rows: int + nnz: int + rhs: ndarray + sign: str | ndarray + row_scaling: ndarray | None -def _concat(arrays: list, dtype: type | None = None) -> ndarray: - """Concatenate arrays, or an empty array when there are none.""" - return np.concatenate(arrays) if arrays else np.array([], dtype=dtype) +def _unit_or(scaling: ndarray) -> ndarray | None: + return None if (scaling == 1).all() else scaling def _row_scaling(c: ConstraintBase) -> ndarray: @@ -40,6 +50,95 @@ def _row_scaling(c: ConstraintBase) -> ndarray: return c.scaling.values.ravel()[c.active_row_mask()] +def _row_block(con: ConstraintBase, label_index: VariableLabelIndex) -> _RowBlock: + row_scaling = _unit_or(_row_scaling(con)) + if isinstance(con, CSRConstraint): + stored = con._csr + return _RowBlock( + partial(con._to_positional_csr, label_index), + stored.shape[0], + int(np.count_nonzero(stored.data)), + con._rhs, + con._sign, + row_scaling, + ) + csr, _, rhs, sense = con.to_matrix_with_rhs(label_index) + return _RowBlock( + lambda: csr, + csr.shape[0], + int(np.count_nonzero(csr.data)), + rhs, + sense, + row_scaling, + ) + + +def _stack( + blocks: list[_RowBlock], + label_index: VariableLabelIndex, + col_scaling: ndarray | None, + model: Model, +) -> tuple[scipy.sparse.csr_array | None, ndarray, ndarray]: + """ + Gather the scaled row blocks into one preallocated CSR, rhs and sense. + + With solver variables y = Scol * x, constraints A x = b become + Srow * A * Scol^-1 * y = Srow * b. + + Explicit zeros are dropped: expressions that broadcast against a dense + coordinate store one coefficient per pair, most of them zero, and a zero + coefficient never changes a constraint. Keeping them only inflates the + stored nnz handed to the solvers/writers (e.g. ``highspy.addRows`` scales + with stored nnz), so we prune them once, centrally, for every backend. + """ + if not blocks: + return None, np.array([]), np.array([], dtype=object) + nnz = sum(block.nnz for block in blocks) + n_rows = sum(block.n_rows for block in blocks) + n_cols = label_index.n_active_vars + dtype = index_dtype(nnz, (n_rows, n_cols), model) + data = np.empty(nnz, dtype=float) + indices = np.empty(nnz, dtype=dtype) + indptr = np.empty(n_rows + 1, dtype=dtype) + b = np.empty(n_rows, dtype=float) + sense = np.empty(n_rows, dtype="U1") + indptr[0] = 0 + pos = row = 0 + for block in blocks: + source = block.positional() + count = block.nnz + block_data, block_indices, block_indptr = ( + source.data, + source.indices, + source.indptr, + ) + if count < source.nnz: + keep = block_data != 0 + block_data, block_indices = block_data[keep], block_indices[keep] + block_indptr = np.concatenate([[0], np.cumsum(keep)])[block_indptr] + rows = slice(row, row + source.shape[0]) + entries = slice(pos, pos + count) + if block.row_scaling is None: + data[entries] = block_data + b[rows] = block.rhs + else: + row_counts = np.diff(block_indptr) + np.multiply( + block_data, np.repeat(block.row_scaling, row_counts), out=data[entries] + ) + np.multiply(block.rhs, block.row_scaling, out=b[rows]) + if col_scaling is not None: + data[entries] /= col_scaling[block_indices] + indices[entries] = block_indices + block_offsets = indptr[rows.start + 1 : rows.stop + 1] + block_offsets[:] = block_indptr[1:] + block_offsets += pos + sense[rows] = block.sign + pos, row = entries.stop, rows.stop + A = scipy.sparse.csr_array((data, indices, indptr), shape=(n_rows, n_cols)) + return A, b, sense + + def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray: """Broadcast an indicator triggering value to one entry per active row.""" if np.ndim(binval) == 0: @@ -69,16 +168,16 @@ def __init__(self, model: Model) -> None: def _build_vars(self) -> None: m = self._parent self.vlabels: ndarray = m.variables.label_index.vlabels + n = len(self.vlabels) + self.var_scaling: ndarray = np.empty(n, dtype=float) + self.lb: ndarray = np.empty(n, dtype=float) + self.ub: ndarray = np.empty(n, dtype=float) + self.vtypes: ndarray = np.empty(n, dtype="U1") - lb_list = [] - ub_list = [] - vtypes_list = [] - scaling_list = [] - + pos = 0 for name, var in m.variables.items(): - labels = var.labels.values.ravel() - mask = labels != -1 - + mask = var.labels.values.ravel() != -1 + cols = slice(pos, pos + int(np.count_nonzero(mask))) if name in m.binaries: vtype = "B" elif name in m.integers: @@ -87,79 +186,48 @@ def _build_vars(self) -> None: vtype = "S" else: vtype = "C" + self.vtypes[cols] = vtype + self.var_scaling[cols] = var.solver_scaling.values.ravel()[mask] + self.lb[cols] = var.lower.values.ravel()[mask] + self.ub[cols] = var.upper.values.ravel()[mask] + pos = cols.stop - lb_list.append(var.lower.values.ravel()[mask]) - ub_list.append(var.upper.values.ravel()[mask]) - vtypes_list.append(np.full(mask.sum(), vtype)) - scaling_list.append(var.solver_scaling.values.ravel()[mask]) - - self.var_scaling: ndarray = _concat(scaling_list, dtype=float) - if lb_list: - self.lb: ndarray = np.concatenate(lb_list) * self.var_scaling - self.ub: ndarray = np.concatenate(ub_list) * self.var_scaling - self.vtypes: ndarray = np.concatenate(vtypes_list) - else: - self.lb = np.array([]) - self.ub = np.array([]) - self.vtypes = np.array([], dtype=object) + if not (self.var_scaling == 1).all(): + self.lb *= self.var_scaling + self.ub *= self.var_scaling def _build_cons(self) -> None: m = self._parent label_index = m.variables.label_index label_to_pos = label_index.label_to_pos - unit_cols = bool((self.var_scaling == 1).all()) - - def scale_rows_and_cols( - csr: scipy.sparse.csr_array, row_scaling: np.ndarray, b: np.ndarray - ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - if csr.shape[0] == 0: - return csr, b - unit_rows = bool((row_scaling == 1).all()) - if unit_rows and unit_cols: - return csr, b - # With solver variables y = Scol * x, constraints A x = b become - # Srow * A * Scol^-1 * y = Srow * b. - data = csr.data - if not unit_rows: - data = data * np.repeat(row_scaling, np.diff(csr.indptr)) - if not unit_cols: - data = data / self.var_scaling[csr.indices] - scaled = scipy.sparse.csr_array( - (data, csr.indices, csr.indptr), shape=csr.shape - ) - return scaled, b * row_scaling + col_scaling = None if (self.var_scaling == 1).all() else self.var_scaling - reg_csrs, reg_b, reg_sense = [], [], [] - ind_csrs, ind_b, ind_sense, ind_binvar, ind_binval = [], [], [], [], [] + regular, indicator = [], [] + binvar, binval = [], [] for c in m.constraints.data.values(): - if c.is_indicator: - cc = c if isinstance(c, CSRConstraint) else c.freeze() - csr, _, b, sense = cc.to_matrix_with_rhs(label_index) - csr, b = scale_rows_and_cols(csr, cc._scaling, b) - ind_csrs.append(csr) - ind_b.append(b) - ind_sense.append(sense) - ind_binvar.append(label_to_pos[cc._binvar_labels]) - binval = cast("int | np.ndarray", cc._binval) - ind_binval.append(_binval_per_row(binval, len(b))) - else: - csr, _, b, sense = c.to_matrix_with_rhs(label_index) - if not isinstance(c, CSRConstraint): - csr.eliminate_zeros() - csr, b = scale_rows_and_cols(csr, _row_scaling(c), b) - reg_csrs.append(csr) - reg_b.append(b) - reg_sense.append(sense) + if not c.is_indicator: + regular.append(_row_block(c, label_index)) + continue + cc = c if isinstance(c, CSRConstraint) else c.freeze() + indicator.append(_row_block(cc, label_index)) + binvar.append(label_to_pos[cc._binvar_labels]) + cc_binval = cast("int | np.ndarray", cc._binval) + binval.append(_binval_per_row(cc_binval, len(cc._rhs))) self.clabels: ndarray = m.constraints.label_index.clabels - self.A: scipy.sparse.csr_array | None = _stack(reg_csrs) - self.b: ndarray = _concat(reg_b) - self.sense: ndarray = _concat(reg_sense, dtype=object) - self.indicator_A: scipy.sparse.csr_array | None = _stack(ind_csrs) - self.indicator_b: ndarray = _concat(ind_b) - self.indicator_sense: ndarray = _concat(ind_sense, dtype=object) - self.indicator_binvar: ndarray = _concat(ind_binvar, dtype=np.intp) - self.indicator_binval: ndarray = _concat(ind_binval, dtype=np.intp) + self.A: scipy.sparse.csr_array | None + self.A, self.b, self.sense = _stack(regular, label_index, col_scaling, m) + self.indicator_A: scipy.sparse.csr_array | None + self.indicator_A, self.indicator_b, self.indicator_sense = _stack( + indicator, label_index, col_scaling, m + ) + label_dtype = m._dtypes["labels"] + self.indicator_binvar: ndarray = np.concatenate( + [np.array([], dtype=label_dtype), *binvar] + ) + self.indicator_binval: ndarray = ( + np.concatenate(binval) if binval else np.array([], dtype=np.intp) + ) @cached_property def c(self) -> ndarray: diff --git a/linopy/model.py b/linopy/model.py index c873f684..8c041441 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -127,11 +127,13 @@ DtypeKey = Literal["labels"] -def _check_infinities(sign: Any, rhs: Any, name: str) -> None: +def _check_infinities( + sign: Any, rhs: Any, name: str, where: np.ndarray | None = None +) -> None: invalid = ((sign == LESS_EQUAL) & (rhs == -np.inf)) | ( (sign == GREATER_EQUAL) & (rhs == np.inf) ) - if np.any(invalid): + if np.any(invalid if where is None else invalid & where): raise ValueError(f"Constraint {name} contains incorrect infinite values.") @@ -1371,9 +1373,14 @@ def add_constraints( con = self._constraint_from_lhs(lhs, sign, rhs, coords) if isinstance(con, CSRConstraint) and freeze: - if mask is not None: - con = con.masked(broadcast_to_coords(mask, con.coords, label="mask")) - _check_infinities(con._sign, con._rhs, name) + row_mask = ( + None + if mask is None + else con.active_values( + broadcast_to_coords(mask, con.coords, label="mask") + ).astype(bool) + ) + _check_infinities(con._sign, con._rhs, name, row_mask) self.check_force_dim_names(con.coords.to_dataset()) enforce_no_multiindex(con, context=f"constraint {name!r}") row_scaling = ( @@ -1385,7 +1392,7 @@ def add_constraints( ) cindex = self._cCounter self._cCounter += con.full_size - con = con.assign_labels(cindex, name, row_scaling) + con = con.assign_labels(cindex, name, row_scaling, row_mask) return self._soften_added(self.constraints.add(con), penalty) if isinstance(con, CSRConstraint): if chunked: @@ -2523,6 +2530,10 @@ def assign_result( if con.is_indicator: continue start, end = con.range + if isinstance(con, CSRConstraint): + active = dual[start:end][con.active_positions] + con._dual = active * con._scaling / self.objective.scaling + continue coords = {dim: con.coords[dim] for dim in con.coord_dims} values = ( dual[start:end].reshape(con.shape) diff --git a/test/test_csr.py b/test/test_csr.py index f5fcbb40..5ddb8a2d 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -11,6 +11,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass, replace +from functools import partial from pathlib import Path from typing import Any @@ -539,12 +540,21 @@ def build(sparse: bool) -> ConstraintBase: np.testing.assert_array_equal(dense.labels.values, sparse.labels.values) +@pytest.mark.parametrize("masked", [False, True], ids=["unmasked", "masked"]) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_invalid_infinite_rhs_raises(sparse: bool) -> None: +def test_frozen_invalid_infinite_rhs_raises(sparse: bool, masked: bool) -> None: require_v1() c = base_model(sparse=sparse) + valid = c.load.bus != "bus0" + rhs = c.load.where(valid, -np.inf) + mask = valid if masked else None + add = partial(c.m.add_constraints, name="bal", freeze=True, mask=mask) + if masked and sparse: + con = add(c.balance_lhs() <= rhs) + assert con.ncons == int(valid.sum()) * c.load.sizes["snapshot"] + return with pytest.raises(ValueError, match="incorrect infinite values"): - c.m.add_constraints(c.balance_lhs() <= -np.inf, name="bal", freeze=True) + add(c.balance_lhs() <= rhs) ROW_SCALINGS: dict[str, Callable[[xr.DataArray], Any]] = { diff --git a/test/test_matrices.py b/test/test_matrices.py index 04317f8a..eeb86f22 100644 --- a/test/test_matrices.py +++ b/test/test_matrices.py @@ -118,3 +118,44 @@ def test_matrices_sol_aligned_with_vlabels() -> None: M = m.matrices np.testing.assert_array_equal(M.vlabels, [0, 2, 3, 4, 5]) np.testing.assert_array_equal(M.sol, [1.0, 3.0, 4.0, 5.0, 6.0]) + + +@pytest.mark.parametrize("freeze", [False, True]) +def test_matrices_scaled_masked_mixed_signs(freeze: bool) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + x = m.add_variables(0, 4, coords=[i], name="x", scaling=[1.0, 2.0, 4.0]) + y = m.add_variables(coords=[i], name="y") + z = m.add_variables(coords=[i], name="z", binary=True) + sign = xr.DataArray(["<=", ">=", "="], coords=[i]) + m.add_constraints( + 2 * x + y - y, + sign, + xr.DataArray([1.0, 2.0, 3.0], coords=[i]), + name="c", + scaling=xr.DataArray([10.0, 1.0, 5.0], coords=[i]), + mask=xr.DataArray([True, False, True], coords=[i]), + freeze=freeze, + ) + m.add_constraints(x + 3 * y >= 1, name="d", freeze=not freeze) + m.add_indicator_constraints(z, 1, x <= 2, name="ind") + M = m.matrices + + assert M.A is not None and M.indicator_A is not None + col_scaling = np.array([1.0, 0.5, 0.25]) + expected = np.zeros((5, 9)) + expected[0, 0], expected[1, 2] = 20.0, 2.5 + expected[2:, :3] = np.diag(col_scaling) + expected[2:, 3:6] = 3 * np.eye(3) + np.testing.assert_array_equal(M.A.toarray(), expected) + assert M.A.nnz == 8 + np.testing.assert_array_equal(M.b, [10.0, 15.0, 1.0, 1.0, 1.0]) + np.testing.assert_array_equal(M.sense, ["<", "=", ">", ">", ">"]) + np.testing.assert_array_equal(M.clabels, [0, 2, 3, 4, 5]) + expected_ind = np.zeros((3, 9)) + expected_ind[:, :3] = np.diag(col_scaling) + np.testing.assert_array_equal(M.indicator_A.toarray(), expected_ind) + np.testing.assert_array_equal(M.indicator_b, [2.0, 2.0, 2.0]) + np.testing.assert_array_equal(M.indicator_binvar, [6, 7, 8]) + np.testing.assert_array_equal(M.lb[:3], [0.0, 0.0, 0.0]) + np.testing.assert_array_equal(M.ub[:3], [4.0, 8.0, 16.0]) diff --git a/test/test_scaling.py b/test/test_scaling.py index 709172c9..f4e6dced 100644 --- a/test/test_scaling.py +++ b/test/test_scaling.py @@ -212,13 +212,14 @@ def test_indicator_constraint_lp_export_uses_scaled_values(tmp_path: Path) -> No assert f"<= {40.0 * 4}" in text -def test_assign_result_unscales_solution_objective_and_dual() -> None: +@pytest.mark.parametrize("freeze", [False, True]) +def test_assign_result_unscales_solution_objective_and_dual(freeze: bool) -> None: m = Model() i = pd.Index(["a", "b"], name="i") x = m.add_variables(coords=[i], name="x", scaling=[10.0, 100.0]) b = m.add_variables(binary=True, name="b", scaling=50.0) - m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0]) + m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0], freeze=freeze) m.add_objective(x.sum() + b, scaling=10.0) primal = np.full(m._xCounter, np.nan) From 9614ae16e654812675c2208d9d14949e6b4121c3 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:59:34 +0200 Subject: [PATCH 4/7] docs: release notes for sparse pipeline performance work --- doc/release_notes.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index c11f209e..3f1bdfba 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -75,6 +75,13 @@ Upcoming Version * ``LinearExpression.flat`` and ``LinearExpression.to_polars`` on a CSR-backed expression are emitted directly from the sparse backing instead of expanding the dense rectangle; the rows equal those of the dense path, with absent cells and zero coefficients dropped. (`#968 `__) * A CSR-backed linear objective stays sparse when set with ``Model.add_objective`` and when the model is exported or solved: the objective vector ``matrices.c``, LP/MPS files, netcdf output, ``Model.copy``, persistent snapshots and direct solver APIs all read it without expanding the dense rectangle, with results equal to the dense objective. The objective name is now owned by the ``Objective`` itself (``Objective.name`` and ``Objective.attrs``) instead of being written into the expression's attributes. (`#967 `__) * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 `__) +* The MOSEK direct API builds bound keys, bounds and constraint names vectorised instead of looping over every row and column in Python. +* The LP writer computes the scaling lookups once per file, skips scaling when every factor is 1, does not re-sort rows that are already grouped (as frozen constraints deliver them) and builds each constraint line from fewer string branches. Constraint writing is about 2x faster; the output is byte-identical. +* ``linopy.merge`` of CSR-backed expressions concatenates all operands once instead of folding them pairwise, so an N-way merge scales with the total number of nonzeros instead of quadratically with the operand count. +* ``LinearExpression.solution`` on a CSR-backed expression is evaluated on its sparse backing without densifying. +* ``@``/``dot`` on a CSR-backed expression sizes its contraction chunks by nonzero count instead of a fixed 64 rows, about 20x faster on wide kept axes. +* ``model.matrices`` assembles the constraint blocks into preallocated CSR buffers in one pass, without ``vstack``, ``eliminate_zeros`` or a full-grid scaling round trip; scaling is skipped when it is identity. Label lookups use the model's label dtype, so ``matrices.clabels`` and ``matrices.indicator_binvar`` are now ``int32`` by default. On an 8M-row model this is 45% faster with 200 MB less peak memory, with bit-identical output. +* Duals of frozen constraints are read back on the active rows directly instead of through three full-grid arrays, and freezing a masked constraint gathers rows once instead of up to three times. Mixed-sign constraints convert their sense vectorised instead of in a Python loop. **Bug fixes** From 437631b4a39213229488c80adeede7e782a1f61b Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 13:56:43 +0200 Subject: [PATCH 5/7] perf(csr): keep fixed contraction chunk, superseded by column compaction (#990) --- doc/release_notes.rst | 1 - linopy/csr.py | 9 ++++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 3f1bdfba..a1805c27 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -79,7 +79,6 @@ Upcoming Version * The LP writer computes the scaling lookups once per file, skips scaling when every factor is 1, does not re-sort rows that are already grouped (as frozen constraints deliver them) and builds each constraint line from fewer string branches. Constraint writing is about 2x faster; the output is byte-identical. * ``linopy.merge`` of CSR-backed expressions concatenates all operands once instead of folding them pairwise, so an N-way merge scales with the total number of nonzeros instead of quadratically with the operand count. * ``LinearExpression.solution`` on a CSR-backed expression is evaluated on its sparse backing without densifying. -* ``@``/``dot`` on a CSR-backed expression sizes its contraction chunks by nonzero count instead of a fixed 64 rows, about 20x faster on wide kept axes. * ``model.matrices`` assembles the constraint blocks into preallocated CSR buffers in one pass, without ``vstack``, ``eliminate_zeros`` or a full-grid scaling round trip; scaling is skipped when it is identity. Label lookups use the model's label dtype, so ``matrices.clabels`` and ``matrices.indicator_binvar`` are now ``int32`` by default. On an 8M-row model this is 45% faster with 200 MB less peak memory, with bit-identical output. * Duals of frozen constraints are read back on the active rows directly instead of through three full-grid arrays, and freezing a masked constraint gathers rows once instead of up to three times. Mixed-sign constraints convert their sense vectorised instead of in a Python loop. diff --git a/linopy/csr.py b/linopy/csr.py index 2c023fc0..b7bad36a 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -52,8 +52,8 @@ from linopy.expressions import LinearExpression from linopy.model import Model -CONTRACTION_CHUNK_NNZ = 1 << 22 -"""Target nonzeros per kept-axis block of the chunked Kronecker product in ``contracted``.""" +CONTRACTION_CHUNK = 64 +"""Kept-axis block size of the chunked Kronecker product in ``contracted``.""" AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] """Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" @@ -615,7 +615,7 @@ def contracted( off its ``name``. The result lives on the kept grid dims followed by ``new_indexes`` and is ``kron(I_kept, matrix.T) @ csr``, evaluated in chunks of the kept axis - sized by nonzero count so the operator never grows with the kept size. + so the operator never grows with the kept size. The result is in compact canonical form: duplicate variables summed, terms label-ordered and explicit zeros pruned -- unlike :meth:`added`, @@ -636,8 +636,7 @@ def contracted( kept_grid = source.grid.reordered(kept) n_kept = kept_grid.size const = np.nan_to_num(source.const) - nnz_per_kept = matrix.nnz + source.csr.nnz // max(n_kept, 1) - chunk = min(max(CONTRACTION_CHUNK_NNZ // max(nnz_per_kept, 1), 1), n_kept) + chunk = min(CONTRACTION_CHUNK, n_kept) operator = scipy.sparse.kron( scipy.sparse.eye_array(chunk), matrix.T, format="csr" ) From 313eb408af20a48baf94279153422f9804b25c1b Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 22:16:29 +0200 Subject: [PATCH 6/7] test(io): type slice_size explicitly in LP constraint section helper --- test/test_io.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/test/test_io.py b/test/test_io.py index aa66f90d..dabee6f5 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -1100,17 +1100,17 @@ def bound2(m: Model, i: int) -> object: assert fn_frozen.read_text() == fn_mutable.read_text() -def _lp_constraint_section(m: Model, path: Path, **kwargs: int) -> str: - m.to_file(path, progress=False, **kwargs) +def _lp_constraint_section(m: Model, path: Path, slice_size: int = 2_000_000) -> str: + m.to_file(path, progress=False, slice_size=slice_size) return path.read_text().split("s.t.\n\n")[1].split("\n\nbounds")[0] @pytest.mark.parametrize("freeze", [True, False]) @pytest.mark.parametrize( - "to_file_kwargs", + "slice_size", [ - pytest.param({}, id="default-slices"), - pytest.param({"slice_size": 1}, id="slice-1"), + pytest.param(2_000_000, id="default-slices"), + pytest.param(1, id="slice-1"), ], ) @pytest.mark.parametrize( @@ -1135,7 +1135,7 @@ def _lp_constraint_section(m: Model, path: Path, **kwargs: int) -> str: def test_to_file_lp_constraint_section( tmp_path: Path, freeze: bool, - to_file_kwargs: dict[str, int], + slice_size: int, scaled: bool, expected: str, ) -> None: @@ -1149,7 +1149,7 @@ def test_to_file_lp_constraint_section( m.add_objective(x.sum()) fn = tmp_path / "constraints.lp" - assert _lp_constraint_section(m, fn, **to_file_kwargs) == expected + assert _lp_constraint_section(m, fn, slice_size) == expected def test_to_file_lp_unsorted_constraint_labels(tmp_path: Path) -> None: From 908a59bae8c0f6e760f29c1d60f92b2e2665e05c Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 7 Oct 2026 10:11:54 +0200 Subject: [PATCH 7/7] test(csr): CSR expression solution raises on removed variables --- test/test_csr.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/test_csr.py b/test/test_csr.py index 5ddb8a2d..4c249669 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -897,6 +897,17 @@ def test_csr_solution_matches_dense( xr.testing.assert_allclose(sol, dense.solution) +def test_csr_solution_raises_on_removed_variable() -> None: + require_v1() + _, sparse = case_twins(lambda c: c.balance_lhs())() + assert sparse._csr is not None + m = sparse.model + m._mock_solve() + m.remove_variables(next(iter(m.variables))) + with pytest.raises(KeyError, match="missing from the model"): + sparse.solution + + def test_reindex_falls_back_to_dense_for_unsupported_kwargs() -> None: require_v1() c1, c2 = twin_models()