From 040e05d09613490bd453afd2a0410b3a2b701106 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 08:19:57 +0200 Subject: [PATCH 1/6] 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 3f2bb8f4..81fff260 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 6e1ec74d..473868da 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: @@ -3526,56 +3526,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 67919367..0c03468d 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 @@ -610,9 +611,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") @@ -715,6 +719,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() @@ -1000,3 +1014,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 2f7487f30fdd04c5ec0537683842d521270cb20b Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:44:10 +0200 Subject: [PATCH 2/6] 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/dualization.py | 9 +++--- linopy/expressions.py | 40 +++++++++++++++++++------- test/test_csr.py | 66 +++++++++++++++++++++++++++++++++++++++++-- 4 files changed, 124 insertions(+), 36 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/dualization.py b/linopy/dualization.py index edcfffdf..27bf68e7 100644 --- a/linopy/dualization.py +++ b/linopy/dualization.py @@ -479,12 +479,13 @@ def _add_dual_feasibility_constraints( dual_vars : dict ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``. """ - A = m.matrices.A + M = m.matrices + A = M.A if A is None: raise ValueError("Constraint matrix is None, model has no constraints.") - vlabels = np.asarray(m.matrices.vlabels, dtype=np.int64) - clabels = np.asarray(m.matrices.clabels, dtype=np.int64) + vlabels = np.asarray(M.vlabels, dtype=np.int64) + clabels = np.asarray(M.clabels, dtype=np.int64) flat_con_to_dual = _build_flat_con_to_dual_label_lookup(m, dual_vars) if not len(flat_con_to_dual): @@ -496,7 +497,7 @@ def _add_dual_feasibility_constraints( flat_v, flat_d, nnz_data = _extract_dual_feas_entries( A, vlabels, clabels, flat_con_to_dual ) - c_lookup = _build_obj_coeff_lookup(vlabels, m.matrices.c) + c_lookup = _build_obj_coeff_lookup(vlabels, M.c) logger.debug("Building dual feasibility constraints for each primal variable.") for var_name, var in m.variables.items(): diff --git a/linopy/expressions.py b/linopy/expressions.py index 827a319a..8aad59de 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1737,17 +1737,25 @@ def mask(self) -> None: return None @has_optimized_model + def _label_solution(self) -> np.ndarray: + """ + Solution values indexed by variable label, with a trailing NaN that + label ``-1`` reads. + """ + sol = np.full(self.model._xCounter + 1, np.nan) + for _, var in self.model.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + sol[labels[mask]] = var.solution.values.ravel()[mask] + return sol + def _map_solution(self) -> DataArray: """ Replace variable labels by solution values. """ - m = self.model - M = m.matrices - sol = pd.Series(M.sol, M.vlabels) - sol[-1] = np.nan - idx = np.ravel(self.vars) - values = np.asarray(sol[idx]).reshape(self.vars.shape) - return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) + labels = self.vars + values = self._label_solution()[labels.values] + return xr.DataArray(values, dims=labels.dims, coords=labels.coords) @property def solution(self) -> DataArray: @@ -2496,6 +2504,20 @@ def const(self) -> DataArray: def const(self, value: DataArray) -> None: self._data = assign_multiindex_safe(self.data, const=value) + @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 = np.nan_to_num(self._label_solution()[: csr.csr.shape[1]]) + return csr.grid.dataarray(csr.csr @ sol + csr.const, name="solution") + def _combined_with_constant( self, self_const: DataArray, @@ -3736,9 +3758,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 a3349105..1192887f 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -748,6 +748,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() @@ -894,16 +950,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 ac955d79b823015622d71acb8417bcde9b154293 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:55:27 +0200 Subject: [PATCH 3/6] 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 | 54 +++++----- linopy/matrices.py | 237 ++++++++++++++++++++++++++---------------- linopy/model.py | 23 ++-- test/test_csr.py | 14 ++- test/test_matrices.py | 42 ++++++++ test/test_scaling.py | 5 +- 7 files changed, 256 insertions(+), 140 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 320c863d..ef8c99b4 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -939,31 +939,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. @@ -1002,13 +1010,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. @@ -1448,7 +1449,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(), @@ -1472,7 +1473,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 @@ -1617,14 +1620,12 @@ def from_dense( scaling = con.scaling.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() active_signs = sign_vals[active_mask] - unique_signs = np.unique(active_signs) - if len(unique_signs) == 0: + sign: str | np.ndarray + if not len(active_signs): full_unique_signs = np.unique(sign_vals) - sign: str | np.ndarray = ( - str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "=" - ) - elif len(unique_signs) == 1: - sign = str(unique_signs[0]) + sign = str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "=" + elif (active_signs == active_signs[0]).all(): + sign = str(active_signs[0]) else: sign = active_signs dual = ( @@ -2221,12 +2222,7 @@ def to_matrix_with_rhs( csr.sum_duplicates() b = self.rhs.values.ravel()[row_mask] - sign_flat = self.sign.values.ravel()[row_mask] - unique_signs = np.unique(sign_flat) - if len(unique_signs) == 1: - sense = np.full(len(con_labels), str(unique_signs[0])[0], dtype="U1") - else: - sense = sign_flat.astype("U1") + sense = self.sign.values.ravel()[row_mask].astype("U1") return csr, con_labels, b, sense def sanitize_zeros(self) -> Constraint: diff --git a/linopy/matrices.py b/linopy/matrices.py index bca66d70..35dfa682 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -7,8 +7,9 @@ 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 @@ -16,15 +17,66 @@ from linopy import expressions from linopy.constraints import CSRConstraint -from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup +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 | None: +class _RowBlock(NamedTuple): """ - Vertically stack CSR blocks, or None when there are none. + 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 _unit_or(scaling: ndarray) -> ndarray | None: + return None if (scaling == 1).all() else scaling + + +def _row_block(con: ConstraintBase, label_index: VariableLabelIndex) -> _RowBlock: + 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, + _unit_or(con._scaling), + ) + csr, _, rhs, sense = con.to_matrix_with_rhs(label_index) + scaling = _unit_or(con.scaling.values.ravel()) + return _RowBlock( + lambda: csr, + csr.shape[0], + int(np.count_nonzero(csr.data)), + rhs, + sense, + None if scaling is None else scaling[con.active_row_mask()], + ) + + +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 @@ -32,16 +84,52 @@ def _stack(csrs: list) -> scipy.sparse.csr_array | None: 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 csrs: - return None - stacked = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) - stacked.eliminate_zeros() - return stacked - - -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) + 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: @@ -72,23 +160,17 @@ def __init__(self, model: Model) -> None: def _build_vars(self) -> None: m = self._parent - label_index = m.variables.label_index - self.vlabels: ndarray = label_index.vlabels - var_scaling_by_label = variable_scaling_lookup(m) - self.var_scaling: ndarray = ( - var_scaling_by_label[self.vlabels] - if len(self.vlabels) - else np.array([], dtype=float) - ) - - lb_list = [] - ub_list = [] - vtypes_list = [] - + 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") + + 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: @@ -97,77 +179,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)) - - 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 - con_scaling_by_label = constraint_scaling_lookup(m) - unit_cols = bool((self.var_scaling == 1).all()) - - def scale_rows_and_cols( - csr: scipy.sparse.csr_array, con_labels: np.ndarray, b: np.ndarray - ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - if csr.shape[0] == 0: - return csr, b - row_scaling = con_scaling_by_label[con_labels] - 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, con_labels, b, sense = cc.to_matrix_with_rhs(label_index) - csr, b = scale_rows_and_cols(csr, con_labels, 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, con_labels, b, sense = c.to_matrix_with_rhs(label_index) - csr, b = scale_rows_and_cols(csr, con_labels, 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 1192887f..e89d0c7b 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 @@ -456,12 +457,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 6da3eaf1..da6743d1 100644 --- a/test/test_matrices.py +++ b/test/test_matrices.py @@ -7,6 +7,7 @@ import numpy as np import pandas as pd +import pytest import xarray as xr from linopy import EQUAL, GREATER_EQUAL, Model @@ -98,3 +99,44 @@ def test_matrices_float_c() -> None: c = m.matrices.c assert np.all(c == np.array([1.5, 1.5])) + + +@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 16c31842..5685121d 100644 --- a/test/test_scaling.py +++ b/test/test_scaling.py @@ -180,13 +180,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 f31578813fcb0822ac43a05ad595c04dcfcae918 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 22:59:34 +0200 Subject: [PATCH 4/6] 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 42e8bee4..8a965da6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -73,6 +73,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`` and ``QuadraticExpression.solution`` no longer rebuild the full constraint matrix; the solution is mapped by variable label, and a CSR-backed expression is evaluated on its sparse backing without densifying. Roughly 10x faster and 6x less memory on a large model. ``Model.dualize`` reads ``model.matrices`` once instead of four times. +* ``@``/``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 1bee8684377d787a5137316f6f3b89be7c50a449 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 13:56:43 +0200 Subject: [PATCH 5/6] 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 8a965da6..92e4ae28 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -77,7 +77,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`` and ``QuadraticExpression.solution`` no longer rebuild the full constraint matrix; the solution is mapped by variable label, and a CSR-backed expression is evaluated on its sparse backing without densifying. Roughly 10x faster and 6x less memory on a large model. ``Model.dualize`` reads ``model.matrices`` once instead of four times. -* ``@``/``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 2e33cec51dd31b6a90718c0ea8d6fe6616b93af3 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 26 Sep 2026 22:16:29 +0200 Subject: [PATCH 6/6] 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 0c03468d..0277e325 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -1016,17 +1016,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( @@ -1051,7 +1051,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: @@ -1065,7 +1065,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: