From 2af0867899bcbf31587a1cb48e6036dafb732a4d Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:11:24 +0200 Subject: [PATCH 1/5] proto: deferred groupby-sum realized directly as CSRConstraint Prototype for #745/#756 option 1 (deferred groupby): hold (expr, grouper) unmaterialized and realize the balance constraint straight from ungrouped long triplets via scipy COO->CSR (duplicate summation == the group sum). No padded _term rectangle ever exists; CSRConstraint plugs into the existing LP/matrix export unchanged. Equivalence: identical polars term rows and LP files vs the dense groupby path (incl. permuted group order). memray peaks on the #745 hub scenario (120 buses, 24 snapshots, build-only, setup baseline ~102MB): hub gens dense deferred 8000 846.6MB 196.8MB (constraint part: ~745MB vs ~95MB) 16000 1213MB 223.0MB dev-scripts is gitignored; files force-added to preserve the prototype on this branch only. Co-Authored-By: Claude Fable 5 --- dev-scripts/proto_deferred_bench.py | 70 ++++++++++ dev-scripts/proto_deferred_check.py | 133 ++++++++++++++++++ dev-scripts/proto_deferred_groupby.py | 193 ++++++++++++++++++++++++++ 3 files changed, 396 insertions(+) create mode 100644 dev-scripts/proto_deferred_bench.py create mode 100644 dev-scripts/proto_deferred_check.py create mode 100644 dev-scripts/proto_deferred_groupby.py diff --git a/dev-scripts/proto_deferred_bench.py b/dev-scripts/proto_deferred_bench.py new file mode 100644 index 00000000..dea79027 --- /dev/null +++ b/dev-scripts/proto_deferred_bench.py @@ -0,0 +1,70 @@ +""" +Benchmark: dense groupby-sum balance constraint vs deferred CSR realization. + +Scenario from issue #745 (hub-skewed nodal balance): 120 buses, one hub with +8000 generators, 100 on each other bus (19,900 total), ring of 120 lines, +24 snapshots. Build-only. + +Run under memray, one variant per process: + memray run -o dense.bin proto_deferred_bench.py dense + memray run -o deferred.bin proto_deferred_bench.py deferred +""" + +from __future__ import annotations + +import sys +import time + +from proto_deferred_check import build_base_model +from proto_deferred_groupby import DeferredGroupbySum, add_deferred_constraints + +N_BUS = 120 +HUB = int(sys.argv[2]) if len(sys.argv) > 2 else 8000 +GENS_PER_BUS = [HUB] + [100] * (N_BUS - 1) +N_SNAP = 24 + + +def main() -> None: + variant = sys.argv[1] + m, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model( + N_BUS, GENS_PER_BUS, N_SNAP + ) + start = time.perf_counter() + + if variant == "setup": + print(f"setup only: {time.perf_counter() - start:.2f}s") + return + if variant == "dense": + lhs = ( + (eff * gen_p).groupby(gbus).sum() + + (1.0 * flow).groupby(bus0).sum() + - (1.0 * flow).groupby(bus1).sum() + ) + con = m.add_constraints(lhs == load, name="balance") + nterm = con.data.sizes["_term"] + elif variant == "deferred": + parts = [ + DeferredGroupbySum(eff * gen_p, gbus), + DeferredGroupbySum(1.0 * flow, bus0), + DeferredGroupbySum(-1.0 * flow, bus1), + ] + con = add_deferred_constraints(m, parts, "=", load, "balance", buses) + nterm = con.nterm + else: + raise SystemExit(f"unknown variant {variant!r}") + + build_s = time.perf_counter() - start + + # prove the export seam works: the LP writer consumes this frame + start = time.perf_counter() + n_rows = con.to_polars().height + export_s = time.perf_counter() - start + + print( + f"{variant}: build {build_s:.2f}s, to_polars {export_s:.2f}s, " + f"nterm={nterm}, polars term rows={n_rows}" + ) + + +if __name__ == "__main__": + main() diff --git a/dev-scripts/proto_deferred_check.py b/dev-scripts/proto_deferred_check.py new file mode 100644 index 00000000..a6a2c34e --- /dev/null +++ b/dev-scripts/proto_deferred_check.py @@ -0,0 +1,133 @@ +""" +Equivalence check: deferred groupby-sum constraint == dense groupby path. + +Builds the same nodal-balance constraint twice on identical models: + dense : m.add_constraints(gen.groupby(bus).sum() + flow.groupby(bus0).sum() + - flow.groupby(bus1).sum() == load) + deferred : add_deferred_constraints(...) via CSRConstraint, no padded rectangle + +and compares labels, terms (coeffs per (label, var)), rhs and sign. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import polars as pl +import xarray as xr +from proto_deferred_groupby import DeferredGroupbySum, add_deferred_constraints + +import linopy + + +def build_base_model(n_bus: int, gens_per_bus: list[int], n_snap: int, seed: int = 0): + """Model with gen_p (gen, snapshot) and flow (line, snapshot) on a ring.""" + rng = np.random.default_rng(seed) + buses = pd.Index([f"bus{i}" for i in range(n_bus)], name="bus") + gen_bus = np.repeat(np.arange(n_bus), gens_per_bus) + gens = pd.Index([f"gen{i}" for i in range(len(gen_bus))], name="gen") + lines = pd.Index([f"line{i}" for i in range(n_bus)], name="line") # ring + snaps = pd.Index(range(n_snap), name="snapshot") + + m = linopy.Model() + gen_p = m.add_variables(coords=[gens, snaps], name="gen_p") + flow = m.add_variables(coords=[lines, snaps], name="flow") + + gbus = pd.Series(buses[gen_bus], index=gens, name="bus") + bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") + bus1 = pd.Series(buses[(np.arange(n_bus) + 1) % n_bus], index=lines, name="bus") + load = xr.DataArray( + rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" + ) + # per-gen efficiency coefficient so coefficients are not all 1 + eff = xr.DataArray(rng.uniform(0.5, 1.5, len(gens)), coords=[gens]) + return m, gen_p, flow, eff, gbus, bus0, bus1, load, buses + + +def canon(df: pl.DataFrame) -> pl.DataFrame: + return ( + df.group_by(["labels", "vars"]) + .agg(pl.col("coeffs").sum(), pl.col("sign").first(), pl.col("rhs").first()) + .sort(["labels", "vars"]) + ) + + +def check(n_bus: int, gens_per_bus: list[int], n_snap: int) -> None: + args = dict(n_bus=n_bus, gens_per_bus=gens_per_bus, n_snap=n_snap) + + # dense path + m1, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model(**args) + lhs = ( + (eff * gen_p).groupby(gbus).sum() + + (1.0 * flow).groupby(bus0).sum() + - (1.0 * flow).groupby(bus1).sum() + ) + c1 = m1.add_constraints(lhs == load, name="balance") + m1.add_objective((1.0 * gen_p).sum()) + print("dense constraint dims:", dict(c1.data.sizes)) + + # deferred path; the dense groupby sorts group labels, so use a sorted + # group index to get an identical constraint grid + m2, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model(**args) + parts = [ + DeferredGroupbySum(eff * gen_p, gbus), + DeferredGroupbySum(1.0 * flow, bus0), + DeferredGroupbySum(-1.0 * flow, bus1), + ] + c2 = add_deferred_constraints(m2, parts, "=", load, "balance", buses.sort_values()) + m2.add_objective((1.0 * gen_p).sum()) + print("deferred nterm (max row nnz):", c2.nterm) + + d1, d2 = canon(c1.to_polars()), canon(c2.to_polars()) + assert d1.shape == d2.shape, (d1.shape, d2.shape) + assert (d1["labels"] == d2["labels"]).all() + assert (d1["vars"] == d2["vars"]).all() + assert np.allclose(d1["coeffs"], d2["coeffs"]) + assert (d1["sign"] == d2["sign"]).all() + assert np.allclose(d1["rhs"], d2["rhs"]) + + # label layout: same grid, same order + assert c1.labels.dims == tuple(c2.coord_names), (c1.labels.dims, c2.coord_names) + assert np.array_equal( + np.sort(c1.labels.values.ravel()), np.sort(c2.active_labels()) + ) + + # end to end: identical LP files (term order within a row is not + # semantic -- CSR emits column-sorted, dense emits build-order -- so + # sort term lines within each block before comparing) + import re + import tempfile + from pathlib import Path + + term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") + + def canon_lp(text: str) -> list[str]: + out: list[str] = [] + buf: list[str] = [] + for line in text.splitlines(): + if term_line.match(line): + buf.append(line) + else: + out += sorted(buf) + [line] + buf = [] + return out + sorted(buf) + + with tempfile.TemporaryDirectory() as tmp: + f1, f2 = Path(tmp) / "dense.lp", Path(tmp) / "deferred.lp" + m1.to_file(f1) + m2.to_file(f2) + assert canon_lp(f1.read_text()) == canon_lp(f2.read_text()), "LP files differ" + + print(f"OK: {d1.height} term rows and LP files identical (dense vs deferred)") + + +def main() -> None: + # small; bus names sort like creation order + check(n_bus=5, gens_per_bus=[7, 1, 3, 1, 2], n_snap=3) + # 12 buses: lexicographic group order ('bus10' < 'bus2') differs from + # creation order, exercising label-based rhs alignment + check(n_bus=12, gens_per_bus=[7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1], n_snap=3) + + +if __name__ == "__main__": + main() diff --git a/dev-scripts/proto_deferred_groupby.py b/dev-scripts/proto_deferred_groupby.py new file mode 100644 index 00000000..1bceffe0 --- /dev/null +++ b/dev-scripts/proto_deferred_groupby.py @@ -0,0 +1,193 @@ +""" +Prototype: deferred groupby().sum() realized directly as a CSRConstraint. + +Idea (issue #745 / #756): the padded rectangle produced by +``expr.groupby(g).sum()`` is pure intermediate waste whenever the result ends +up as a constraint — the LP/matrix export un-pads it again. Instead of +materializing the grouped expression, hold (ungrouped expr, grouper) as a +deferred node and realize the constraint straight from the ungrouped long +triplets: + + row = flat position of (group, *rest_dims) in the constraint grid + col = label_to_pos[var_label] + data = coeff + +scipy's COO->CSR conversion sums duplicate (row, col) entries, which *is* the +group sum. No `_term` padding ever exists; peak memory is O(nnz), independent +of the group-size distribution. + +This module deliberately builds a `CSRConstraint` — the existing second +constraint backend (#630 / freeze_constraints) — so the whole downstream +pipeline (LP writer via `to_polars`, matrix export via `to_matrix_with_rhs`) +works unchanged. + +Prototype limitations (documented, not fundamental): +- all deferred parts must share the same non-grouped ("rest") dims, +- the grouped result is only usable as a constraint LHS (no further + expression arithmetic; that path would fall back to the dense kernel), +- signs are uniform per constraint, masks are not supported. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pandas as pd +import scipy.sparse +import xarray as xr + +from linopy.constants import TERM_DIM +from linopy.constraints import CSRConstraint +from linopy.expressions import LinearExpression +from linopy.model import Model + + +@dataclass +class DeferredGroupbySum: + """ + An unmaterialized ``expr.groupby(grouper).sum()``. + + ``grouper`` maps the labels of one of ``expr``'s dims (its index, whose + ``name`` must be that dim) to group labels. + """ + + expr: LinearExpression + grouper: pd.Series + + @property + def member_dim(self) -> str: + return str(self.grouper.index.name) + + @property + def rest_dims(self) -> list[str]: + return [d for d in self.expr.coord_dims if d != self.member_dim] + + def materialize(self) -> LinearExpression: + """Fall back to the dense scatter kernel (today's behavior).""" + return self.expr.groupby(self.grouper).sum() + + +def _part_triplets( + part: DeferredGroupbySum, + group_index: pd.Index, + rest_dims: list[str], + label_to_pos: np.ndarray, + n_rest: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Long-form (row, col, coeff) triplets of one part, plus its grouped const. + + Rows are flat positions in the (group, *rest_dims) constraint grid. + """ + ds = part.expr.data + member_dim = part.member_dim + + grouper = part.grouper.reindex(ds.indexes[member_dim]) + if grouper.isna().any(): + raise ValueError(f"grouper does not cover all {member_dim!r} labels") + gcodes = group_index.get_indexer(grouper.to_numpy()) + if (gcodes == -1).any(): + raise ValueError("grouper contains labels outside the group index") + + coeffs = ds.coeffs.transpose(member_dim, *rest_dims, TERM_DIM).to_numpy() + vars_ = ds.vars.transpose(member_dim, *rest_dims, TERM_DIM).to_numpy() + nterm = ds.sizes[TERM_DIM] + + # flat row id of each (member, rest...) cell in the constraint grid + cell_rows = gcodes[:, None] * n_rest + np.arange(n_rest)[None, :] + rows = np.repeat(cell_rows.reshape(-1), nterm) + cols_label = vars_.reshape(-1) + data = coeffs.reshape(-1) + + keep = (cols_label != -1) & (data != 0) & ~np.isnan(data) + rows, cols_label, data = rows[keep], cols_label[keep], data[keep] + cols = label_to_pos[cols_label] + + # group-sum of the constant (a true numeric reduction, already cheap) + const = ( + ds.const.transpose(member_dim, *rest_dims) + .to_numpy() + .reshape(len(gcodes), n_rest) + ) + const_grid = np.zeros((len(group_index), n_rest)) + np.add.at(const_grid, gcodes, np.where(np.isnan(const), 0.0, const)) + + return rows, cols, data, const_grid.reshape(-1) + + +def add_deferred_constraints( + model: Model, + parts: list[DeferredGroupbySum], + sign: str, + rhs: float | xr.DataArray, + name: str, + group_index: pd.Index, +) -> CSRConstraint: + """ + Realize ``sum(parts) sign rhs`` as a CSRConstraint, no dense rectangle. + + Equivalent to + ``model.add_constraints(sum(p.materialize() for p in parts), sign, rhs, name)`` + for constraint rows over ``(group_index, *rest_dims)``. + """ + rest_dims = parts[0].rest_dims + for p in parts[1:]: + if p.rest_dims != rest_dims: + raise ValueError("all parts must share the same non-grouped dims") + rest_indexes = [parts[0].expr.data.indexes[d] for d in rest_dims] + n_rest = int(np.prod([len(ix) for ix in rest_indexes])) if rest_indexes else 1 + full_size = len(group_index) * n_rest + + label_index = model.variables.label_index + label_to_pos = label_index.label_to_pos + + rows_l, cols_l, data_l = [], [], [] + const_flat = np.zeros(full_size) + for part in parts: + rows, cols, data, const = _part_triplets( + part, group_index, rest_dims, label_to_pos, n_rest + ) + rows_l.append(rows) + cols_l.append(cols) + data_l.append(data) + const_flat += const + + coo = scipy.sparse.coo_array( + (np.concatenate(data_l), (np.concatenate(rows_l), np.concatenate(cols_l))), + shape=(full_size, label_index.n_active_vars), + ) + csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum + csr.eliminate_zeros() + + # rhs over the grid; expression constants move to the right-hand side + if isinstance(rhs, xr.DataArray): + grid = {str(group_index.name): group_index} | dict(zip(rest_dims, rest_indexes)) + rhs_flat = ( + rhs.reindex(grid) # align by label to the constraint grid + .transpose(str(group_index.name), *rest_dims) + .to_numpy() + .reshape(-1) + ) + else: + rhs_flat = np.full(full_size, float(rhs)) + rhs_flat = rhs_flat - const_flat + + # label allocation, mirroring Model._allocate_constraint_labels + cindex = model._cCounter + model._cCounter += full_size + active = np.diff(csr.indptr) > 0 + con_labels = np.arange(cindex, cindex + full_size)[active] + + con = CSRConstraint( + csr[active], + con_labels, + rhs_flat[active], + sign, + coords=[group_index, *rest_indexes], + model=model, + name=name, + cindex=cindex, + ) + model.constraints.add(con) + return con From e225992bf858473a8d3207eda181653a4299c92c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:49:20 +0200 Subject: [PATCH 2/5] feat(v1): lazy groupby-sum with direct CSR realization under freeze Internal-state redesign of the deferred-groupby prototype: no new public primitive. groupby(g).sum(lazy=True) (or options['lazy_groupby'] under v1) returns an ordinary LinearExpression whose payload is a LazyGroupSum (ungrouped parts + groupers) instead of the materialized dense dataset, modeled on dask-backed xarray. The .data property materializes through today's kernel, so any operation without a lazy branch transparently falls back to exactly today's result. Lazy branches: neg, scalar mul, merge along _term (covers +/-), comparison with a constant rhs, and Model.add_constraints with freeze=True, which realizes the constraint directly as a CSRConstraint from long triplets (COO->CSR duplicate summation is the group sum) - the #745 padded rectangle never exists. Gated behind v1 semantics; under legacy, lazy=True raises and the option is ignored. v1 parity kept: NaN rhs raises (par.5), reordered rhs raises (par.8), absent const rows realize as masked (par.12). memray, #745 hub scenario (120 buses, 24 snapshots, setup ~107MB): hub gens eager lazy 8000 851.9MB 208.4MB 16000 1219MB 223.3MB Full suite: 6336 passed, 557 skipped (test/remote failures pre-exist on the branch). Includes test/test_lazy_groupby.py (10 cases x legacy/v1). Co-Authored-By: Claude Fable 5 --- dev-scripts/proto_deferred_bench.py | 50 ++-- dev-scripts/proto_deferred_check.py | 133 ---------- dev-scripts/proto_deferred_groupby.py | 193 --------------- linopy/config.py | 1 + linopy/constraints.py | 22 +- linopy/expressions.py | 69 +++++- linopy/lazy.py | 337 ++++++++++++++++++++++++++ linopy/model.py | 13 + test/test_lazy_groupby.py | 258 ++++++++++++++++++++ 9 files changed, 717 insertions(+), 359 deletions(-) delete mode 100644 dev-scripts/proto_deferred_check.py delete mode 100644 dev-scripts/proto_deferred_groupby.py create mode 100644 linopy/lazy.py create mode 100644 test/test_lazy_groupby.py diff --git a/dev-scripts/proto_deferred_bench.py b/dev-scripts/proto_deferred_bench.py index dea79027..df558740 100644 --- a/dev-scripts/proto_deferred_bench.py +++ b/dev-scripts/proto_deferred_bench.py @@ -1,13 +1,14 @@ """ -Benchmark: dense groupby-sum balance constraint vs deferred CSR realization. +Benchmark: eager groupby-sum balance constraint vs lazy CSR realization. Scenario from issue #745 (hub-skewed nodal balance): 120 buses, one hub with -8000 generators, 100 on each other bus (19,900 total), ring of 120 lines, -24 snapshots. Build-only. +HUB generators, 100 on each other bus, ring of 120 lines, 24 snapshots. +Build-only, using the real API: ``groupby(g).sum(lazy=...)`` + +``add_constraints(lhs == load, freeze=...)``. Run under memray, one variant per process: - memray run -o dense.bin proto_deferred_bench.py dense - memray run -o deferred.bin proto_deferred_bench.py deferred + memray run -o eager.bin proto_deferred_bench.py eager [hub] + memray run -o lazy.bin proto_deferred_bench.py lazy [hub] """ from __future__ import annotations @@ -15,54 +16,43 @@ import sys import time -from proto_deferred_check import build_base_model -from proto_deferred_groupby import DeferredGroupbySum, add_deferred_constraints +sys.path.insert( + 0, str(__import__("pathlib").Path(__file__).resolve().parent.parent / "test") +) +from test_lazy_groupby import balance_lhs, base_model N_BUS = 120 HUB = int(sys.argv[2]) if len(sys.argv) > 2 else 8000 -GENS_PER_BUS = [HUB] + [100] * (N_BUS - 1) +GENS_PER_BUS = tuple([HUB] + [100] * (N_BUS - 1)) N_SNAP = 24 def main() -> None: + import linopy + + linopy.options["semantics"] = "v1" # the lazy path is gated behind v1 variant = sys.argv[1] - m, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model( - N_BUS, GENS_PER_BUS, N_SNAP + m, gen_p, flow, eff, gbus, bus0, bus1, load, buses = base_model( + gens_per_bus=GENS_PER_BUS, n_snap=N_SNAP ) start = time.perf_counter() if variant == "setup": print(f"setup only: {time.perf_counter() - start:.2f}s") return - if variant == "dense": - lhs = ( - (eff * gen_p).groupby(gbus).sum() - + (1.0 * flow).groupby(bus0).sum() - - (1.0 * flow).groupby(bus1).sum() - ) - con = m.add_constraints(lhs == load, name="balance") - nterm = con.data.sizes["_term"] - elif variant == "deferred": - parts = [ - DeferredGroupbySum(eff * gen_p, gbus), - DeferredGroupbySum(1.0 * flow, bus0), - DeferredGroupbySum(-1.0 * flow, bus1), - ] - con = add_deferred_constraints(m, parts, "=", load, "balance", buses) - nterm = con.nterm - else: - raise SystemExit(f"unknown variant {variant!r}") + lazy = variant == "lazy" + lhs = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=lazy) + con = m.add_constraints(lhs == load, name="balance", freeze=lazy) build_s = time.perf_counter() - start - # prove the export seam works: the LP writer consumes this frame start = time.perf_counter() n_rows = con.to_polars().height export_s = time.perf_counter() - start print( f"{variant}: build {build_s:.2f}s, to_polars {export_s:.2f}s, " - f"nterm={nterm}, polars term rows={n_rows}" + f"type={type(con).__name__}, polars term rows={n_rows}" ) diff --git a/dev-scripts/proto_deferred_check.py b/dev-scripts/proto_deferred_check.py deleted file mode 100644 index a6a2c34e..00000000 --- a/dev-scripts/proto_deferred_check.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -Equivalence check: deferred groupby-sum constraint == dense groupby path. - -Builds the same nodal-balance constraint twice on identical models: - dense : m.add_constraints(gen.groupby(bus).sum() + flow.groupby(bus0).sum() - - flow.groupby(bus1).sum() == load) - deferred : add_deferred_constraints(...) via CSRConstraint, no padded rectangle - -and compares labels, terms (coeffs per (label, var)), rhs and sign. -""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import polars as pl -import xarray as xr -from proto_deferred_groupby import DeferredGroupbySum, add_deferred_constraints - -import linopy - - -def build_base_model(n_bus: int, gens_per_bus: list[int], n_snap: int, seed: int = 0): - """Model with gen_p (gen, snapshot) and flow (line, snapshot) on a ring.""" - rng = np.random.default_rng(seed) - buses = pd.Index([f"bus{i}" for i in range(n_bus)], name="bus") - gen_bus = np.repeat(np.arange(n_bus), gens_per_bus) - gens = pd.Index([f"gen{i}" for i in range(len(gen_bus))], name="gen") - lines = pd.Index([f"line{i}" for i in range(n_bus)], name="line") # ring - snaps = pd.Index(range(n_snap), name="snapshot") - - m = linopy.Model() - gen_p = m.add_variables(coords=[gens, snaps], name="gen_p") - flow = m.add_variables(coords=[lines, snaps], name="flow") - - gbus = pd.Series(buses[gen_bus], index=gens, name="bus") - bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") - bus1 = pd.Series(buses[(np.arange(n_bus) + 1) % n_bus], index=lines, name="bus") - load = xr.DataArray( - rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" - ) - # per-gen efficiency coefficient so coefficients are not all 1 - eff = xr.DataArray(rng.uniform(0.5, 1.5, len(gens)), coords=[gens]) - return m, gen_p, flow, eff, gbus, bus0, bus1, load, buses - - -def canon(df: pl.DataFrame) -> pl.DataFrame: - return ( - df.group_by(["labels", "vars"]) - .agg(pl.col("coeffs").sum(), pl.col("sign").first(), pl.col("rhs").first()) - .sort(["labels", "vars"]) - ) - - -def check(n_bus: int, gens_per_bus: list[int], n_snap: int) -> None: - args = dict(n_bus=n_bus, gens_per_bus=gens_per_bus, n_snap=n_snap) - - # dense path - m1, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model(**args) - lhs = ( - (eff * gen_p).groupby(gbus).sum() - + (1.0 * flow).groupby(bus0).sum() - - (1.0 * flow).groupby(bus1).sum() - ) - c1 = m1.add_constraints(lhs == load, name="balance") - m1.add_objective((1.0 * gen_p).sum()) - print("dense constraint dims:", dict(c1.data.sizes)) - - # deferred path; the dense groupby sorts group labels, so use a sorted - # group index to get an identical constraint grid - m2, gen_p, flow, eff, gbus, bus0, bus1, load, buses = build_base_model(**args) - parts = [ - DeferredGroupbySum(eff * gen_p, gbus), - DeferredGroupbySum(1.0 * flow, bus0), - DeferredGroupbySum(-1.0 * flow, bus1), - ] - c2 = add_deferred_constraints(m2, parts, "=", load, "balance", buses.sort_values()) - m2.add_objective((1.0 * gen_p).sum()) - print("deferred nterm (max row nnz):", c2.nterm) - - d1, d2 = canon(c1.to_polars()), canon(c2.to_polars()) - assert d1.shape == d2.shape, (d1.shape, d2.shape) - assert (d1["labels"] == d2["labels"]).all() - assert (d1["vars"] == d2["vars"]).all() - assert np.allclose(d1["coeffs"], d2["coeffs"]) - assert (d1["sign"] == d2["sign"]).all() - assert np.allclose(d1["rhs"], d2["rhs"]) - - # label layout: same grid, same order - assert c1.labels.dims == tuple(c2.coord_names), (c1.labels.dims, c2.coord_names) - assert np.array_equal( - np.sort(c1.labels.values.ravel()), np.sort(c2.active_labels()) - ) - - # end to end: identical LP files (term order within a row is not - # semantic -- CSR emits column-sorted, dense emits build-order -- so - # sort term lines within each block before comparing) - import re - import tempfile - from pathlib import Path - - term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") - - def canon_lp(text: str) -> list[str]: - out: list[str] = [] - buf: list[str] = [] - for line in text.splitlines(): - if term_line.match(line): - buf.append(line) - else: - out += sorted(buf) + [line] - buf = [] - return out + sorted(buf) - - with tempfile.TemporaryDirectory() as tmp: - f1, f2 = Path(tmp) / "dense.lp", Path(tmp) / "deferred.lp" - m1.to_file(f1) - m2.to_file(f2) - assert canon_lp(f1.read_text()) == canon_lp(f2.read_text()), "LP files differ" - - print(f"OK: {d1.height} term rows and LP files identical (dense vs deferred)") - - -def main() -> None: - # small; bus names sort like creation order - check(n_bus=5, gens_per_bus=[7, 1, 3, 1, 2], n_snap=3) - # 12 buses: lexicographic group order ('bus10' < 'bus2') differs from - # creation order, exercising label-based rhs alignment - check(n_bus=12, gens_per_bus=[7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1], n_snap=3) - - -if __name__ == "__main__": - main() diff --git a/dev-scripts/proto_deferred_groupby.py b/dev-scripts/proto_deferred_groupby.py deleted file mode 100644 index 1bceffe0..00000000 --- a/dev-scripts/proto_deferred_groupby.py +++ /dev/null @@ -1,193 +0,0 @@ -""" -Prototype: deferred groupby().sum() realized directly as a CSRConstraint. - -Idea (issue #745 / #756): the padded rectangle produced by -``expr.groupby(g).sum()`` is pure intermediate waste whenever the result ends -up as a constraint — the LP/matrix export un-pads it again. Instead of -materializing the grouped expression, hold (ungrouped expr, grouper) as a -deferred node and realize the constraint straight from the ungrouped long -triplets: - - row = flat position of (group, *rest_dims) in the constraint grid - col = label_to_pos[var_label] - data = coeff - -scipy's COO->CSR conversion sums duplicate (row, col) entries, which *is* the -group sum. No `_term` padding ever exists; peak memory is O(nnz), independent -of the group-size distribution. - -This module deliberately builds a `CSRConstraint` — the existing second -constraint backend (#630 / freeze_constraints) — so the whole downstream -pipeline (LP writer via `to_polars`, matrix export via `to_matrix_with_rhs`) -works unchanged. - -Prototype limitations (documented, not fundamental): -- all deferred parts must share the same non-grouped ("rest") dims, -- the grouped result is only usable as a constraint LHS (no further - expression arithmetic; that path would fall back to the dense kernel), -- signs are uniform per constraint, masks are not supported. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np -import pandas as pd -import scipy.sparse -import xarray as xr - -from linopy.constants import TERM_DIM -from linopy.constraints import CSRConstraint -from linopy.expressions import LinearExpression -from linopy.model import Model - - -@dataclass -class DeferredGroupbySum: - """ - An unmaterialized ``expr.groupby(grouper).sum()``. - - ``grouper`` maps the labels of one of ``expr``'s dims (its index, whose - ``name`` must be that dim) to group labels. - """ - - expr: LinearExpression - grouper: pd.Series - - @property - def member_dim(self) -> str: - return str(self.grouper.index.name) - - @property - def rest_dims(self) -> list[str]: - return [d for d in self.expr.coord_dims if d != self.member_dim] - - def materialize(self) -> LinearExpression: - """Fall back to the dense scatter kernel (today's behavior).""" - return self.expr.groupby(self.grouper).sum() - - -def _part_triplets( - part: DeferredGroupbySum, - group_index: pd.Index, - rest_dims: list[str], - label_to_pos: np.ndarray, - n_rest: int, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """ - Long-form (row, col, coeff) triplets of one part, plus its grouped const. - - Rows are flat positions in the (group, *rest_dims) constraint grid. - """ - ds = part.expr.data - member_dim = part.member_dim - - grouper = part.grouper.reindex(ds.indexes[member_dim]) - if grouper.isna().any(): - raise ValueError(f"grouper does not cover all {member_dim!r} labels") - gcodes = group_index.get_indexer(grouper.to_numpy()) - if (gcodes == -1).any(): - raise ValueError("grouper contains labels outside the group index") - - coeffs = ds.coeffs.transpose(member_dim, *rest_dims, TERM_DIM).to_numpy() - vars_ = ds.vars.transpose(member_dim, *rest_dims, TERM_DIM).to_numpy() - nterm = ds.sizes[TERM_DIM] - - # flat row id of each (member, rest...) cell in the constraint grid - cell_rows = gcodes[:, None] * n_rest + np.arange(n_rest)[None, :] - rows = np.repeat(cell_rows.reshape(-1), nterm) - cols_label = vars_.reshape(-1) - data = coeffs.reshape(-1) - - keep = (cols_label != -1) & (data != 0) & ~np.isnan(data) - rows, cols_label, data = rows[keep], cols_label[keep], data[keep] - cols = label_to_pos[cols_label] - - # group-sum of the constant (a true numeric reduction, already cheap) - const = ( - ds.const.transpose(member_dim, *rest_dims) - .to_numpy() - .reshape(len(gcodes), n_rest) - ) - const_grid = np.zeros((len(group_index), n_rest)) - np.add.at(const_grid, gcodes, np.where(np.isnan(const), 0.0, const)) - - return rows, cols, data, const_grid.reshape(-1) - - -def add_deferred_constraints( - model: Model, - parts: list[DeferredGroupbySum], - sign: str, - rhs: float | xr.DataArray, - name: str, - group_index: pd.Index, -) -> CSRConstraint: - """ - Realize ``sum(parts) sign rhs`` as a CSRConstraint, no dense rectangle. - - Equivalent to - ``model.add_constraints(sum(p.materialize() for p in parts), sign, rhs, name)`` - for constraint rows over ``(group_index, *rest_dims)``. - """ - rest_dims = parts[0].rest_dims - for p in parts[1:]: - if p.rest_dims != rest_dims: - raise ValueError("all parts must share the same non-grouped dims") - rest_indexes = [parts[0].expr.data.indexes[d] for d in rest_dims] - n_rest = int(np.prod([len(ix) for ix in rest_indexes])) if rest_indexes else 1 - full_size = len(group_index) * n_rest - - label_index = model.variables.label_index - label_to_pos = label_index.label_to_pos - - rows_l, cols_l, data_l = [], [], [] - const_flat = np.zeros(full_size) - for part in parts: - rows, cols, data, const = _part_triplets( - part, group_index, rest_dims, label_to_pos, n_rest - ) - rows_l.append(rows) - cols_l.append(cols) - data_l.append(data) - const_flat += const - - coo = scipy.sparse.coo_array( - (np.concatenate(data_l), (np.concatenate(rows_l), np.concatenate(cols_l))), - shape=(full_size, label_index.n_active_vars), - ) - csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum - csr.eliminate_zeros() - - # rhs over the grid; expression constants move to the right-hand side - if isinstance(rhs, xr.DataArray): - grid = {str(group_index.name): group_index} | dict(zip(rest_dims, rest_indexes)) - rhs_flat = ( - rhs.reindex(grid) # align by label to the constraint grid - .transpose(str(group_index.name), *rest_dims) - .to_numpy() - .reshape(-1) - ) - else: - rhs_flat = np.full(full_size, float(rhs)) - rhs_flat = rhs_flat - const_flat - - # label allocation, mirroring Model._allocate_constraint_labels - cindex = model._cCounter - model._cCounter += full_size - active = np.diff(csr.indptr) > 0 - con_labels = np.arange(cindex, cindex + full_size)[active] - - con = CSRConstraint( - csr[active], - con_labels, - rhs_flat[active], - sign, - coords=[group_index, *rest_indexes], - model=model, - name=name, - cindex=cindex, - ) - model.constraints.add(con) - return con diff --git a/linopy/config.py b/linopy/config.py index 92cd0945..ef929534 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -82,4 +82,5 @@ def __repr__(self) -> str: display_max_rows=14, display_max_terms=6, semantics=LEGACY_SEMANTICS, + lazy_groupby=False, ) diff --git a/linopy/constraints.py b/linopy/constraints.py index d323a4de..acfcb39a 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1144,7 +1144,7 @@ class Constraint(ConstraintBase): Supports setters, xarray operations via conwrap, and from_rule construction. """ - __slots__ = ("_data", "_model", "_assigned", "_coef_dirty") + __slots__ = ("_data", "_model", "_assigned", "_coef_dirty", "_lazy") def __init__( self, @@ -1174,9 +1174,29 @@ def __init__( self._data = data self._model = model self._coef_dirty = False + self._lazy = None + + @classmethod + def _from_lazy( + cls, lhs: expressions.LinearExpression, sign: str, rhs: Any, model: Model + ) -> Constraint: + """Anonymous constraint over a still-lazy lhs (see linopy.lazy).""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None + obj._assigned = False + obj._coef_dirty = False + obj._lazy = (lhs, sign, rhs) + return obj @property def data(self) -> Dataset: + if self._data is None and self._lazy is not None: + lhs, sign, rhs = self._lazy + lhs.data # noqa: B018 materialize; clears lhs's lazy payload + self._data = lhs.to_constraint(sign, rhs).data + self._assigned = "labels" in self._data + self._lazy = None return self._data @property diff --git a/linopy/expressions.py b/linopy/expressions.py index 9c957d69..099f94d6 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -477,7 +477,10 @@ def map( return LinearExpression(self.groupby.map(func, args=args, **kwargs), self.model) def sum( - self, use_fallback: bool = False, observed: bool = False + self, + use_fallback: bool = False, + observed: bool = False, + lazy: bool | None = None, ) -> LinearExpression: """ Sum the expression over each group. @@ -491,6 +494,17 @@ def sum( Fall back to the previous, slower groupby-sum implementation, kept as an escape hatch. Leave at False unless the default misbehaves. Defaults to False. + lazy : bool, optional + Return the grouped sum unmaterialized: a long-format payload + behind the ordinary LinearExpression type, akin to dask-backed + xarray objects. Any operation without a lazy branch + transparently materializes today's dense result; a still-lazy + left-hand side reaching ``Model.add_constraints`` with + ``freeze=True`` is realized directly as a CSRConstraint, + skipping the group-size-padded term dimension. Requires v1 + semantics. Defaults to ``linopy.options["lazy_groupby"]`` + (only honored under v1). Only single-key groupers over an + existing dimension are held lazily. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -515,6 +529,24 @@ def sum( multikey_frame = ( None if use_fallback else _multikey_value_frame(group, self.data) ) + + if lazy is None: + lazy = is_v1() and options["lazy_groupby"] + elif lazy and not is_v1(): + raise ValueError( + "lazy groupby-sum requires v1 semantics; opt in with " + "linopy.options['semantics'] = 'v1'." + ) + if lazy and not use_fallback and not observed and multikey_frame is None: + series = group.to_pandas() if isinstance(group, DataArray) else group + if isinstance(series, pd.Series) and series.index.name in self.data.dims: + from linopy.lazy import LazyGroupSum + + expr = LinearExpression(self.data, self.model) + group_name = str(getattr(series, "name", None) or "group") + payload = LazyGroupSum.from_grouper(expr, series, group_name) + return LinearExpression._from_lazy(payload, self.model) + if multikey_frame is not None: group = multikey_frame @@ -711,7 +743,7 @@ def sum(self, **kwargs: Any) -> LinearExpression: class BaseExpression(ABC): - __slots__ = ("_data", "_model") + __slots__ = ("_data", "_model", "_lazy") __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 @@ -785,6 +817,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: self._model = model self._data = cast(Dataset, data) + self._lazy = None def __repr__(self) -> str: """ @@ -885,6 +918,8 @@ def __neg__(self) -> Self: """ Get the negative of the expression. """ + if self._lazy is not None: + return self._from_lazy(self._lazy.scaled(-1.0), self._model) return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) def _multiply_by_linear_expression( @@ -1460,8 +1495,21 @@ def type(self) -> str: @property def data(self) -> Dataset: + if self._data is None and self._lazy is not None: + # materialize through today's dense kernel; laziness is spent + self._data = self._lazy.materialize().data + self._lazy = None return self._data + @classmethod + def _from_lazy(cls, lazy: Any, model: Model) -> Self: + """Construct an expression backed by an unmaterialized LazyGroupSum.""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None # type: ignore[assignment] + obj._lazy = lazy + return obj + @property def model(self) -> Model: return self._model @@ -1473,6 +1521,9 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: + if self._data is None and self._lazy is not None: + # answerable from the lazy payload without materializing + return tuple(self._lazy.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @property @@ -1661,6 +1712,11 @@ def to_constraint( which are moved to the left-hand-side and constant values which are moved to the right-hand side. """ + if self._lazy is not None and isinstance(sign, str) and is_constant(rhs): + # carry the lazy payload on the constraint; its .data property + # materializes through this method's dense path if needed + return constraints.Constraint._from_lazy(self, sign, rhs, self.model) + rhs = as_constant(rhs) if self.is_constant and is_constant(rhs): raise ValueError( @@ -2298,6 +2354,8 @@ def __mul__( """ Multiply the expr by a factor. """ + if self._lazy is not None and isinstance(other, int | float | np.number): + return type(self)._from_lazy(self._lazy.scaled(float(other)), self._model) other = as_constant(other) if isinstance(other, QuadraticExpression): return other.__rmul__(self) @@ -3069,6 +3127,13 @@ def merge( model = exprs[0].model + if issubclass(cls, LinearExpression) and not has_quad_expression: + from linopy.lazy import try_lazy_merge + + lazy_result = try_lazy_merge(exprs, dim=dim, join=join, kwargs=kwargs) + if lazy_result is not None: + return lazy_result + data = [e.data if isinstance(e, linopy_types) else e for e in exprs] data = [fill_missing_coords(ds, fill_helper_dims=True) for ds in data] diff --git a/linopy/lazy.py b/linopy/lazy.py new file mode 100644 index 00000000..023a3321 --- /dev/null +++ b/linopy/lazy.py @@ -0,0 +1,337 @@ +""" +Lazy groupby-sum: an internal long-format representation for grouped sums. + +``expr.groupby(g).sum(lazy=True)`` (or ``linopy.options["lazy_groupby"] = +True``) returns an ordinary :class:`~linopy.expressions.LinearExpression` +whose payload is a :class:`LazyGroupSum` — a list of ungrouped parts plus +their groupers — instead of the materialized dense dataset. Modeled on +dask-backed xarray objects: same public type, different backing, and any +operation without a lazy branch transparently materializes through the +``.data`` property (yielding exactly today's result, so laziness can only +delay the dense rectangle, never change semantics). + +The payoff comes at ``Model.add_constraints`` with ``freeze=True`` (or +``Model.freeze_constraints``): a still-lazy left-hand side is realized +directly as a :class:`~linopy.constraints.CSRConstraint` from the ungrouped +long triplets — scipy's COO->CSR duplicate summation *is* the group sum — +so the ``group-size-padded`` ``_term`` rectangle (issue #745) never exists. + +Operations with a lazy branch (everything else materializes): + +- unary minus and multiplication by a scalar (scales the ungrouped parts), +- ``merge`` along ``_term`` — and therefore ``+``/``-`` — with other lazy + expressions on the same grid or with dense expressions already on the + result grid, +- comparison with a constant rhs (``==``, ``<=``, ``>=``), carried lazily + on the resulting :class:`~linopy.constraints.Constraint`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import scipy.sparse +from xarray import DataArray + +from linopy.constants import HELPER_DIMS, TERM_DIM + +if TYPE_CHECKING: + from linopy.constraints import CSRConstraint + from linopy.expressions import LinearExpression + from linopy.model import Model + + +@dataclass(frozen=True) +class LazyPart: + """ + One summand: an ungrouped expression plus its grouper. + + ``grouper`` maps the labels of ``expr``'s member dim (the grouper's + index, named after that dim) to group labels. ``None`` marks an + identity part that already lives on the result grid. + """ + + expr: LinearExpression + grouper: pd.Series | None + + +@dataclass(frozen=True) +class LazyGroupSum: + """ + A sum of grouped/dense parts over a fixed result grid. + + ``grid_dims`` is the result's coordinate dim order (the grouped dim in + the member dim's position, as the dense kernel produces it) and + ``indexes`` maps each grid dim to its pandas index. + """ + + parts: tuple[LazyPart, ...] + group_dim: str + grid_dims: tuple[str, ...] + indexes: dict[str, pd.Index] + + @classmethod + def from_grouper( + cls, expr: LinearExpression, grouper: pd.Series, group_dim: str + ) -> LazyGroupSum: + member_dim = str(grouper.index.name) + if member_dim in expr.data.indexes: + # labels checked equal upstream; conform the order + grouper = grouper.reindex(expr.data.indexes[member_dim]) + elif len(grouper) != expr.data.sizes[member_dim]: + raise ValueError(f"grouper length does not match dimension {member_dim!r}") + codes, uniques = pd.factorize(grouper, sort=True) + if (codes == -1).any(): + raise ValueError( + "Cannot group by a pandas object containing NaN values. " + "Drop or fill the corresponding entries before grouping." + ) + grid_dims = tuple( + group_dim if d == member_dim else str(d) for d in expr.coord_dims + ) + indexes = { + d: expr.data.get_index(d).rename(d) + for d in expr.coord_dims + if d != member_dim + } + indexes[group_dim] = pd.Index(uniques, name=group_dim) + return cls((LazyPart(expr, grouper),), group_dim, grid_dims, indexes) + + @property + def group_index(self) -> pd.Index: + return self.indexes[self.group_dim] + + def scaled(self, factor: float) -> LazyGroupSum: + parts = tuple(LazyPart(p.expr * factor, p.grouper) for p in self.parts) + return replace(self, parts=parts) + + def same_grid(self, other: LazyGroupSum) -> bool: + return ( + self.group_dim == other.group_dim + and self.grid_dims == other.grid_dims + and all(self.indexes[d].equals(other.indexes[d]) for d in self.grid_dims) + ) + + def wrap_dense(self, expr: LinearExpression) -> LazyPart | None: + """Wrap a dense expression already on the result grid as identity part.""" + if set(expr.coord_dims) != set(self.grid_dims): + return None + for d in expr.coord_dims: + if not expr.data.get_index(d).equals(self.indexes[str(d)]): + return None + return LazyPart(expr, None) + + def materialize(self) -> LinearExpression: + """Left-fold the parts through today's dense kernel.""" + result: LinearExpression | None = None + for part in self.parts: + if part.grouper is None: + term = part.expr + else: + term = part.expr.groupby(part.grouper).sum(lazy=False) + result = term if result is None else result + term + assert result is not None + return result + + +def try_lazy_merge( + exprs: Any, dim: str, join: Any, kwargs: dict +) -> LinearExpression | None: + """ + Lazy branch of :func:`linopy.expressions.merge`. + + Returns a lazy combined expression when every input is a plain + LinearExpression on the same result grid (lazy, or dense wrappable as + an identity part) — where any join produces the identical result — and + None otherwise to fall through to the dense path. + """ + from linopy.expressions import LinearExpression + + if dim != TERM_DIM or kwargs: + return None + if not all(type(e) is LinearExpression for e in exprs): + return None + lazies = [e._lazy for e in exprs if e._lazy is not None] + if not lazies: + return None + template = lazies[0] + if not all(template.same_grid(lz) for lz in lazies[1:]): + return None + + parts: list[LazyPart] = [] + for e in exprs: + if e._lazy is not None: + parts.extend(e._lazy.parts) + else: + part = template.wrap_dense(e) + if part is None: + return None + parts.append(part) + combined = replace(template, parts=tuple(parts)) + return LinearExpression._from_lazy(combined, exprs[0].model) + + +def extract_lazy(lhs: Any, sign: Any, rhs: Any) -> tuple[LazyGroupSum, str, Any] | None: + """Return (payload, sign, rhs) if lhs is a realizable lazy constraint.""" + from linopy.common import is_constant + from linopy.constraints import Constraint + from linopy.expressions import LinearExpression + + if ( + isinstance(lhs, Constraint) + and lhs._lazy is not None + and sign is None + and rhs is None + ): + expr, sign, rhs = lhs._lazy + if expr._lazy is None: # already materialized elsewhere + return None + lhs = expr + if not (isinstance(lhs, LinearExpression) and lhs._lazy is not None): + return None + if not isinstance(sign, str) or rhs is None or not is_constant(rhs): + return None + lazy = lhs._lazy + rhs_da = _as_rhs_dataarray(rhs) + if rhs_da is None or not set(rhs_da.dims) <= set(lazy.grid_dims): + return None + return lazy, sign, rhs + + +def _as_rhs_dataarray(rhs: Any) -> DataArray | None: + from linopy.alignment import as_dataarray + + try: + da = as_dataarray(rhs) + except (TypeError, ValueError): + return None + if set(da.dims) & set(HELPER_DIMS): + return None + return da + + +def _part_entries( + part: LazyPart, lazy: LazyGroupSum, label_to_pos: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Long-form (rows, cols, coeffs) triplets of one part plus its const. + + Rows are flat positions in the ``grid_dims`` constraint grid (C order). + Returns (rows_per_term, cols, coeffs, cell_rows_flat, const_flat_values). + """ + ds = part.expr.data + shape = tuple(len(lazy.indexes[d]) for d in lazy.grid_dims) + strides = np.array( + [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] + ) + + # per grid dim: the position of each of the part's cells along that dim + transposed: list[str] = [] + axis_positions: list[np.ndarray] = [] + for d, stride in zip(lazy.grid_dims, strides): + if part.grouper is not None and d == lazy.group_dim: + member_dim = str(part.grouper.index.name) + transposed.append(member_dim) + codes = lazy.group_index.get_indexer(part.grouper.to_numpy()) + axis_positions.append(codes * stride) + else: + transposed.append(d) + axis_positions.append(np.arange(shape[lazy.grid_dims.index(d)]) * stride) + + cell_rows = axis_positions[0] + for pos in axis_positions[1:]: + cell_rows = cell_rows[..., None] + pos + cell_rows = cell_rows.reshape(-1) + + coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + nterm = ds.sizes[TERM_DIM] + rows = np.repeat(cell_rows, nterm) + + keep = (vars_ != -1) & (coeffs != 0) & ~np.isnan(coeffs) + rows, vars_, coeffs = rows[keep], vars_[keep], coeffs[keep] + + const = ds.const.transpose(*transposed).to_numpy().reshape(-1) + return rows, label_to_pos[vars_], coeffs, cell_rows, const + + +def realize_lazy_constraint( + model: Model, lazy: LazyGroupSum, sign: str, rhs: Any, name: str +) -> CSRConstraint: + """ + Realize ``lazy sign rhs`` directly as a CSRConstraint. + + Equivalent to materializing the lazy sum, adding the constraint and + freezing it — but the padded dense ``_term`` rectangle never exists; + peak memory is proportional to the number of nonzero terms. + """ + from linopy.common import maybe_replace_sign + from linopy.constraints import CSRConstraint + from linopy.semantics import check_user_nan, is_v1 + + sign = maybe_replace_sign(sign) + shape = tuple(len(lazy.indexes[d]) for d in lazy.grid_dims) + full_size = int(np.prod(shape, dtype=np.int64)) if shape else 1 + + label_index = model.variables.label_index + label_to_pos = label_index.label_to_pos + + rows_l, cols_l, data_l = [], [], [] + const_flat = np.zeros(full_size) + for part in lazy.parts: + rows, cols, coeffs, cell_rows, const = _part_entries(part, lazy, label_to_pos) + rows_l.append(rows) + cols_l.append(cols) + data_l.append(coeffs) + np.add.at(const_flat, cell_rows, np.where(np.isnan(const), 0.0, const)) + + coo = scipy.sparse.coo_array( + (np.concatenate(data_l), (np.concatenate(rows_l), np.concatenate(cols_l))), + shape=(full_size, label_index.n_active_vars), + ) + csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum + csr.eliminate_zeros() + + # rhs aligned by label to the grid; expression constants move to the rhs + rhs_da = _as_rhs_dataarray(rhs) + assert rhs_da is not None + if is_v1() and bool(rhs_da.isnull().any()): + check_user_nan() # §5: NaN in a user constant raises under v1 + if is_v1(): + # §8 parity with the dense path: a reordered or differing rhs index + # raises; the user aligns explicitly with .sel/.reindex + for d in rhs_da.dims: + if not rhs_da.get_index(d).equals(lazy.indexes[str(d)]): + raise ValueError( + f"Coordinate mismatch on shared dimension {d!r} between " + "the rhs and the grouped result. Align the rhs with " + ".sel(...) / .reindex(...) before combining (§8)." + ) + rhs_da = rhs_da.reindex( + {d: lazy.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan + ) + missing = {d: lazy.indexes[d] for d in lazy.grid_dims if d not in rhs_da.dims} + if missing: + rhs_da = rhs_da.expand_dims(missing) + rhs_flat = rhs_da.transpose(*lazy.grid_dims).to_numpy().reshape(-1) - const_flat + + # label allocation, mirroring Model._allocate_constraint_labels; rows + # without terms or with NaN rhs are inactive, as in the frozen dense path + cindex = model._cCounter + model._cCounter += full_size + active = (np.diff(csr.indptr) > 0) & ~np.isnan(rhs_flat) + con_labels = np.arange(cindex, cindex + full_size)[active] + + return CSRConstraint( + csr[active], + con_labels, + rhs_flat[active], + sign, + coords=[lazy.indexes[d] for d in lazy.grid_dims], + model=model, + name=name, + cindex=cindex, + ) diff --git a/linopy/model.py b/linopy/model.py index f15c086b..8c7e2fc6 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1138,6 +1138,19 @@ def add_constraints( """ name = self._resolve_constraint_name(name) + + # a still-lazy grouped lhs with freeze on realizes directly as a + # CSRConstraint from long triplets, skipping the padded dense form + resolved_freeze = self.freeze_constraints if freeze is None else freeze + if resolved_freeze and mask is None and not self.chunk: + from linopy.lazy import extract_lazy, realize_lazy_constraint + + extracted = extract_lazy(lhs, sign, rhs) + if extracted is not None: + lazy_sum, lazy_sign, lazy_rhs = extracted + con = realize_lazy_constraint(self, lazy_sum, lazy_sign, lazy_rhs, name) + return self.constraints.add(con) + if sign is not None: sign = maybe_replace_signs(as_dataarray(sign)) diff --git a/test/test_lazy_groupby.py b/test/test_lazy_groupby.py new file mode 100644 index 00000000..23e3461f --- /dev/null +++ b/test/test_lazy_groupby.py @@ -0,0 +1,258 @@ +""" +Tests for lazy groupby-sum (linopy.lazy): type stability, transparent +materialization, and direct CSR realization under freeze. v1-only feature. +""" + +from __future__ import annotations + +import re + +import numpy as np +import pandas as pd +import polars as pl +import pytest +import xarray as xr + +import linopy +from linopy import LinearExpression +from linopy.constraints import Constraint, CSRConstraint +from linopy.semantics import is_v1 +from linopy.testing import assert_conequal, assert_linequal + + +def require_v1() -> None: + if not is_v1(): + pytest.skip("lazy groupby-sum is gated behind v1 semantics") + + +def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): + """Model with gen_p (gen, snapshot) and flow (line, snapshot) on a ring.""" + rng = np.random.default_rng(seed) + n_bus = len(gens_per_bus) + buses = pd.Index([f"bus{i}" for i in range(n_bus)], name="bus") + gen_bus = np.repeat(np.arange(n_bus), gens_per_bus) + gens = pd.Index([f"gen{i}" for i in range(len(gen_bus))], name="gen") + lines = pd.Index([f"line{i}" for i in range(n_bus)], name="line") + snaps = pd.Index(range(n_snap), name="snapshot") + + m = linopy.Model() + gen_p = m.add_variables(coords=[gens, snaps], name="gen_p") + flow = m.add_variables(coords=[lines, snaps], name="flow") + + gbus = pd.Series(buses[gen_bus], index=gens, name="bus") + bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") + bus1 = pd.Series(buses[(np.arange(n_bus) + 1) % n_bus], index=lines, name="bus") + # v1 requires explicit alignment: order the load like the (sorted) groups + load = xr.DataArray( + rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" + ).sortby("bus") + eff = xr.DataArray(rng.uniform(0.5, 1.5, len(gens)), coords=[gens]) + return m, gen_p, flow, eff, gbus, bus0, bus1, load, buses + + +def balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy): + return ( + (eff * gen_p).groupby(gbus).sum(lazy=lazy) + + (1.0 * flow).groupby(bus0).sum(lazy=lazy) + - (1.0 * flow).groupby(bus1).sum(lazy=lazy) + ) + + +def canon(df: pl.DataFrame) -> pl.DataFrame: + return ( + df.group_by(["labels", "vars"]) + .agg(pl.col("coeffs").sum(), pl.col("sign").first(), pl.col("rhs").first()) + .sort(["labels", "vars"]) + ) + + +def test_lazy_requires_v1(): + m, gen_p, flow, eff, gbus, *_ = base_model() + if is_v1(): + res = (eff * gen_p).groupby(gbus).sum(lazy=True) + assert type(res) is LinearExpression + return + with pytest.raises(ValueError, match="requires v1 semantics"): + (eff * gen_p).groupby(gbus).sum(lazy=True) + # the option is not honored under legacy: result is eager and dense + linopy.options["lazy_groupby"] = True + try: + res = (eff * gen_p).groupby(gbus).sum() + finally: + linopy.options["lazy_groupby"] = False + assert res._lazy is None + + +def test_lazy_is_plain_linear_expression_and_materializes_identically(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + lazy = (eff * gen_p).groupby(gbus).sum(lazy=True) + eager = (eff * gen_p).groupby(gbus).sum() + assert type(lazy) is LinearExpression + # comparing data materializes; the result must be exactly today's + assert_linequal(lazy, eager) + + +def test_lazy_composition_materializes_identically(): + require_v1() + m, gen_p, flow, eff, gbus, bus0, bus1, *_ = base_model() + lazy = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=True) + eager = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=False) + assert type(lazy) is LinearExpression + assert_linequal(lazy, eager) + + +def test_scalar_ops_stay_lazy(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + lazy = -2.0 * (eff * gen_p).groupby(gbus).sum(lazy=True) + assert lazy._lazy is not None # still unmaterialized after neg/scale + eager = -2.0 * (eff * gen_p).groupby(gbus).sum() + # scaling after grouping normalizes padded-slot fills while scaling + # before grouping keeps fresh pads; masked-slot fill is not contractual, + # so compare with masked slots blanked on both sides + from linopy.testing import _sort_by_vars_along_term + + a, b = _sort_by_vars_along_term(lazy), _sort_by_vars_along_term(eager) + assert (a.vars == b.vars).all() + xr.testing.assert_allclose( + a.coeffs.where(a.vars != -1), b.coeffs.where(b.vars != -1) + ) + xr.testing.assert_equal(a.const, b.const) + + +def test_freeze_realizes_csr_without_dense_rectangle(): + require_v1() + m1, *rest = base_model() + lhs1 = balance_lhs(*rest[:6], lazy=False) + c1 = m1.add_constraints(lhs1 == rest[6], name="balance") + + m2, *rest2 = base_model() + lhs2 = balance_lhs(*rest2[:6], lazy=True) + c2 = m2.add_constraints(lhs2 == rest2[6], name="balance", freeze=True) + + assert isinstance(c2, CSRConstraint) + d1, d2 = canon(c1.to_polars()), canon(c2.to_polars()) + assert d1["labels"].equals(d2["labels"]) + assert d1["vars"].equals(d2["vars"]) + assert np.allclose(d1["coeffs"], d2["coeffs"]) + assert (d1["sign"] == d2["sign"]).all() + assert np.allclose(d1["rhs"], d2["rhs"]) + assert np.array_equal( + np.sort(c1.labels.values.ravel()), np.sort(c2.active_labels()) + ) + + +def test_freeze_false_falls_back_to_identical_dense_constraint(): + require_v1() + m1, *rest = base_model() + c1 = m1.add_constraints( + balance_lhs(*rest[:6], lazy=False) == rest[6], name="balance" + ) + + m2, *rest2 = base_model() + c2 = m2.add_constraints( + balance_lhs(*rest2[:6], lazy=True) == rest2[6], name="balance" + ) + assert isinstance(c2, Constraint) + assert_conequal(c1, c2) + + +def test_option_gates_lazy_and_freeze_model_default(): + require_v1() + m, *rest = base_model() + m.freeze_constraints = True + linopy.options["lazy_groupby"] = True + try: + lhs = balance_lhs(*rest[:6], lazy=None) + con = m.add_constraints(lhs == rest[6], name="balance") + finally: + linopy.options["lazy_groupby"] = False + assert isinstance(con, CSRConstraint) + + +def test_materialized_lazy_still_freezes_via_dense_path(): + require_v1() + m, *rest = base_model() + lhs = balance_lhs(*rest[:6], lazy=True) + _ = lhs.nterm # force materialization before constraint creation + con = m.add_constraints(lhs == rest[6], name="balance", freeze=True) + assert isinstance(con, CSRConstraint) + + +def test_nan_rhs_raises_on_both_paths(): + require_v1() + m1, *rest = base_model() + load = rest[6].copy() + load[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN"): + m1.add_constraints( + balance_lhs(*rest[:6], lazy=False) == load, name="bal", freeze=True + ) + + m2, *rest2 = base_model() + load2 = rest2[6].copy() + load2[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN"): + m2.add_constraints( + balance_lhs(*rest2[:6], lazy=True) == load2, name="bal", freeze=True + ) + + +def test_reordered_rhs_raises_on_both_paths(): + require_v1() + m1, *rest = base_model() + load = rest[6].isel(bus=slice(None, None, -1)) # same labels, reversed + with pytest.raises(ValueError, match="[Cc]oordinate"): + m1.add_constraints( + balance_lhs(*rest[:6], lazy=False) == load, name="bal", freeze=True + ) + + m2, *rest2 = base_model() + load2 = rest2[6].isel(bus=slice(None, None, -1)) + with pytest.raises(ValueError, match="[Cc]oordinate"): + m2.add_constraints( + balance_lhs(*rest2[:6], lazy=True) == load2, name="bal", freeze=True + ) + + +def test_nan_grouper_raises_eagerly(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + gbus = gbus.copy() + gbus.iloc[0] = np.nan + with pytest.raises(ValueError, match="NaN values"): + (1.0 * gen_p).groupby(gbus).sum(lazy=True) + + +def test_lp_files_identical(tmp_path): + require_v1() + sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) + + m1, *rest = base_model(gens_per_bus=sizes) + m1.add_constraints(balance_lhs(*rest[:6], lazy=False) == rest[6], name="balance") + m1.add_objective((1.0 * rest[0]).sum()) + + m2, *rest2 = base_model(gens_per_bus=sizes) + m2.add_constraints( + balance_lhs(*rest2[:6], lazy=True) == rest2[6], name="balance", freeze=True + ) + m2.add_objective((1.0 * rest2[0]).sum()) + + term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") + + def canon_lp(text: str) -> list[str]: + out: list[str] = [] + buf: list[str] = [] + for line in text.splitlines(): + if term_line.match(line): + buf.append(line) + else: + out += sorted(buf) + [line] + buf = [] + return out + sorted(buf) + + f1, f2 = tmp_path / "eager.lp", tmp_path / "lazy.lp" + m1.to_file(f1) + m2.to_file(f2) + assert canon_lp(f1.read_text()) == canon_lp(f2.read_text()) From 2518c6778d8da5392eaa24aed7d4bd5a9e24e375 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:02:02 +0200 Subject: [PATCH 3/5] feat(v1): CSR-backed sparse groupby-sum (payload swap from lazy recipe) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the LazyGroupSum recipe payload with CSRPayload (linopy/csr.py): groupby(g).sum(sparse=True) (or options['sparse_groupby'] under v1) now builds the grouped sum eagerly in CSR form — rows are flat grid cells, columns raw variable labels, duplicate variables summed, terms label-ordered — behind the unchanged LinearExpression type. Operations become sparse linear algebra: grouping scatters into rows (G @ A), merge along _term (and thus +/-) is sparse addition, neg/scalar-mul scale values, and add_constraints(freeze=True) staples sign/rhs on to form a CSRConstraint directly. Anything else expands to the dense rectangle in canonical form via .data (mathematically identical, term layout canonicalized — the reason the feature stays v1-gated). Vs the recipe: ops are real algebra with immediate errors instead of a deferred parts list, chains stay compact, and the payload is the natural seam for future sparse ops (dot #748, ragged merge #749). Cost: the bit-identical fallback is replaced by the canonical-form contract. memray, #745 hub scenario (120 buses, 24 snapshots, setup ~107MB): hub gens eager sparse (CSR) [recipe was] 8000 850.9MB 216.2MB 208.4MB 16000 1219MB 234.1MB 223.3MB Full suite: 6336 passed, 557 skipped (unchanged). Co-Authored-By: Claude Fable 5 --- dev-scripts/proto_deferred_bench.py | 16 +- linopy/config.py | 2 +- linopy/constraints.py | 18 +- linopy/csr.py | 388 ++++++++++++++++++ linopy/expressions.py | 88 ++-- linopy/lazy.py | 337 --------------- linopy/model.py | 12 +- ...lazy_groupby.py => test_sparse_groupby.py} | 91 ++-- 8 files changed, 504 insertions(+), 448 deletions(-) create mode 100644 linopy/csr.py delete mode 100644 linopy/lazy.py rename test/{test_lazy_groupby.py => test_sparse_groupby.py} (69%) diff --git a/dev-scripts/proto_deferred_bench.py b/dev-scripts/proto_deferred_bench.py index df558740..ce20b28d 100644 --- a/dev-scripts/proto_deferred_bench.py +++ b/dev-scripts/proto_deferred_bench.py @@ -1,14 +1,14 @@ """ -Benchmark: eager groupby-sum balance constraint vs lazy CSR realization. +Benchmark: eager groupby-sum balance constraint vs sparse CSR realization. Scenario from issue #745 (hub-skewed nodal balance): 120 buses, one hub with HUB generators, 100 on each other bus, ring of 120 lines, 24 snapshots. -Build-only, using the real API: ``groupby(g).sum(lazy=...)`` + +Build-only, using the real API: ``groupby(g).sum(sparse=...)`` + ``add_constraints(lhs == load, freeze=...)``. Run under memray, one variant per process: memray run -o eager.bin proto_deferred_bench.py eager [hub] - memray run -o lazy.bin proto_deferred_bench.py lazy [hub] + memray run -o sparse.bin proto_deferred_bench.py sparse [hub] """ from __future__ import annotations @@ -19,7 +19,7 @@ sys.path.insert( 0, str(__import__("pathlib").Path(__file__).resolve().parent.parent / "test") ) -from test_lazy_groupby import balance_lhs, base_model +from test_sparse_groupby import balance_lhs, base_model N_BUS = 120 HUB = int(sys.argv[2]) if len(sys.argv) > 2 else 8000 @@ -30,7 +30,7 @@ def main() -> None: import linopy - linopy.options["semantics"] = "v1" # the lazy path is gated behind v1 + linopy.options["semantics"] = "v1" # the sparse path is gated behind v1 variant = sys.argv[1] m, gen_p, flow, eff, gbus, bus0, bus1, load, buses = base_model( gens_per_bus=GENS_PER_BUS, n_snap=N_SNAP @@ -41,9 +41,9 @@ def main() -> None: print(f"setup only: {time.perf_counter() - start:.2f}s") return - lazy = variant == "lazy" - lhs = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=lazy) - con = m.add_constraints(lhs == load, name="balance", freeze=lazy) + sparse = variant == "sparse" + lhs = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=sparse) + con = m.add_constraints(lhs == load, name="balance", freeze=sparse) build_s = time.perf_counter() - start start = time.perf_counter() diff --git a/linopy/config.py b/linopy/config.py index ef929534..590a4cab 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -82,5 +82,5 @@ def __repr__(self) -> str: display_max_rows=14, display_max_terms=6, semantics=LEGACY_SEMANTICS, - lazy_groupby=False, + sparse_groupby=False, ) diff --git a/linopy/constraints.py b/linopy/constraints.py index acfcb39a..b73b872f 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1144,7 +1144,7 @@ class Constraint(ConstraintBase): Supports setters, xarray operations via conwrap, and from_rule construction. """ - __slots__ = ("_data", "_model", "_assigned", "_coef_dirty", "_lazy") + __slots__ = ("_data", "_model", "_assigned", "_coef_dirty", "_pending") def __init__( self, @@ -1174,29 +1174,29 @@ def __init__( self._data = data self._model = model self._coef_dirty = False - self._lazy = None + self._pending = None @classmethod - def _from_lazy( + def _from_pending( cls, lhs: expressions.LinearExpression, sign: str, rhs: Any, model: Model ) -> Constraint: - """Anonymous constraint over a still-lazy lhs (see linopy.lazy).""" + """Anonymous constraint over a still-sparse lhs (see linopy.csr).""" obj = cls.__new__(cls) obj._model = model obj._data = None obj._assigned = False obj._coef_dirty = False - obj._lazy = (lhs, sign, rhs) + obj._pending = (lhs, sign, rhs) return obj @property def data(self) -> Dataset: - if self._data is None and self._lazy is not None: - lhs, sign, rhs = self._lazy - lhs.data # noqa: B018 materialize; clears lhs's lazy payload + if self._data is None and self._pending is not None: + lhs, sign, rhs = self._pending + lhs.data # noqa: B018 materialize; clears lhs's sparse payload self._data = lhs.to_constraint(sign, rhs).data self._assigned = "labels" in self._data - self._lazy = None + self._pending = None return self._data @property diff --git a/linopy/csr.py b/linopy/csr.py new file mode 100644 index 00000000..474c9a6c --- /dev/null +++ b/linopy/csr.py @@ -0,0 +1,388 @@ +""" +CSR-backed expressions: a sparse internal representation for grouped sums. + +``expr.groupby(g).sum(sparse=True)`` (or ``linopy.options["sparse_groupby"] += True`` under v1) returns an ordinary +:class:`~linopy.expressions.LinearExpression` whose payload is a +:class:`CSRPayload` — the expression as a scipy CSR matrix (rows = flat +cells of the coordinate grid, columns = variable labels, values = +coefficients) plus a dense const vector — instead of the materialized +dense dataset. Modeled on dask-backed xarray objects: same public type, +different backing; any operation without a sparse branch transparently +converts to the dense form through the ``.data`` property. + +The CSR form is canonical: duplicate variables within a cell are summed +(``2x + 3x -> 5x``) and terms are ordered by variable label, so the +``_term`` axis is ragged by construction — the group-size padding of issue +#745 has no analog. Operations become sparse linear algebra: + +- grouping scatters members into group rows (conceptually ``G @ A`` with a + 0/1 grouping matrix), +- ``merge`` along ``_term`` — and therefore ``+``/``-`` — is sparse matrix + addition, +- unary minus and scalar multiplication scale the values, +- ``== rhs`` with a constant is carried as a pending payload on the + :class:`~linopy.constraints.Constraint`, and ``Model.add_constraints`` + with ``freeze=True`` staples sign and rhs on to produce a + :class:`~linopy.constraints.CSRConstraint` directly. + +Converting back to the dense rectangle yields the mathematically identical +expression in canonical form (not the eager kernel's exact term layout); +this is why the feature is gated behind v1 semantics, where the term +layout is explicitly non-contractual. + +Columns are raw variable labels (not positions), so payloads stay valid +when variables are added to the model later; realization maps labels to +dense positions via the model's label index. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import scipy.sparse +from xarray import DataArray, Dataset + +from linopy.constants import HELPER_DIMS, TERM_DIM + +if TYPE_CHECKING: + from linopy.constraints import CSRConstraint + from linopy.expressions import LinearExpression + from linopy.model import Model + + +@dataclass(frozen=True) +class CSRPayload: + """ + An expression as ``A @ x + c`` over a fixed coordinate grid. + + ``csr`` has one row per flat grid cell (C order over ``grid_dims``) + and one column per variable label; ``const`` is the per-cell constant. + """ + + csr: scipy.sparse.csr_array + const: np.ndarray + grid_dims: tuple[str, ...] + indexes: dict[str, pd.Index] + model: Model + + @property + def shape(self) -> tuple[int, ...]: + return tuple(len(self.indexes[d]) for d in self.grid_dims) + + @property + def n_cells(self) -> int: + return self.csr.shape[0] + + @classmethod + def from_grouper( + cls, expr: LinearExpression, grouper: pd.Series, group_dim: str + ) -> CSRPayload: + """Build the grouped sum directly in CSR form (no padded rectangle).""" + member_dim = str(grouper.index.name) + if member_dim in expr.data.indexes: + # labels checked equal upstream; conform the order + grouper = grouper.reindex(expr.data.indexes[member_dim]) + elif len(grouper) != expr.data.sizes[member_dim]: + raise ValueError(f"grouper length does not match dimension {member_dim!r}") + codes, uniques = pd.factorize(grouper, sort=True) + if (codes == -1).any(): + raise ValueError( + "Cannot group by a pandas object containing NaN values. " + "Drop or fill the corresponding entries before grouping." + ) + grid_dims = tuple( + group_dim if d == member_dim else str(d) for d in expr.coord_dims + ) + indexes = { + d: expr.data.get_index(d).rename(d) + for d in expr.coord_dims + if d != member_dim + } + indexes[group_dim] = pd.Index(uniques, name=group_dim) + return _assemble(expr, grid_dims, indexes, group_dim, member_dim, codes) + + @classmethod + def from_expression( + cls, expr: LinearExpression, template: CSRPayload + ) -> CSRPayload | None: + """Convert a dense expression on the template's grid to CSR form.""" + if set(expr.coord_dims) != set(template.grid_dims): + return None + for d in expr.coord_dims: + if not expr.data.get_index(d).equals(template.indexes[str(d)]): + return None + first = template.grid_dims[0] + codes = np.arange(len(template.indexes[first])) + return _assemble( + expr, template.grid_dims, template.indexes, first, first, codes + ) + + def scaled(self, factor: float) -> CSRPayload: + return replace(self, csr=self.csr * factor, const=self.const * factor) + + def same_grid(self, other: CSRPayload) -> bool: + return self.grid_dims == other.grid_dims and all( + self.indexes[d].equals(other.indexes[d]) for d in self.grid_dims + ) + + def add(self, other: CSRPayload) -> CSRPayload: + """Sparse matrix addition == merge along the term dimension.""" + a, b = self.csr, other.csr + width = max(a.shape[1], b.shape[1]) + a, b = _widen(a, width), _widen(b, width) + return replace(self, csr=a + b, const=self.const + other.const) + + def materialize(self) -> LinearExpression: + """ + Expand to the dense rectangle in canonical form. + + Terms are variable-label-ordered and duplicates are summed; the + term axis is padded to the widest cell with the usual fill + (vars=-1, coeffs=NaN). + """ + from linopy.expressions import LinearExpression + + csr = self.csr.copy() + csr.eliminate_zeros() + csr.sort_indices() + lengths = np.diff(csr.indptr) + nterm = max(int(lengths.max()) if len(lengths) else 0, 1) + + labels_dtype = self.model._dtypes["labels"] + vars_flat = np.full((self.n_cells, nterm), -1, dtype=labels_dtype) + coeffs_flat = np.full((self.n_cells, nterm), np.nan) + rows = np.repeat(np.arange(self.n_cells), lengths) + pos = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], lengths) + vars_flat[rows, pos] = csr.indices + coeffs_flat[rows, pos] = csr.data + + shape = self.shape + dims = (*self.grid_dims, TERM_DIM) + coords = {d: self.indexes[d] for d in self.grid_dims} + ds = Dataset( + { + "coeffs": (dims, coeffs_flat.reshape(*shape, nterm)), + "vars": (dims, vars_flat.reshape(*shape, nterm)), + "const": (self.grid_dims, self.const.reshape(shape)), + }, + coords=coords, + ) + return LinearExpression(ds, self.model) + + +def _widen(csr: scipy.sparse.csr_array, width: int) -> scipy.sparse.csr_array: + if csr.shape[1] == width: + return csr + csr = csr.copy() + csr.resize((csr.shape[0], width)) + return csr + + +def _assemble( + expr: LinearExpression, + grid_dims: tuple[str, ...], + indexes: dict[str, pd.Index], + scatter_dim: str, + member_dim: str, + codes: np.ndarray, +) -> CSRPayload: + """ + Scatter an expression's terms into grid rows (conceptually G @ A). + + ``expr``'s ``member_dim`` scatters into the grid dim ``scatter_dim`` + at the row positions given by ``codes``; every other grid dim maps + one-to-one. + """ + ds = expr.data + shape = tuple(len(indexes[d]) for d in grid_dims) + strides = np.array( + [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] + ) + + transposed: list[str] = [] + axis_positions: list[np.ndarray] = [] + for i, (d, stride) in enumerate(zip(grid_dims, strides)): + if d == scatter_dim: + transposed.append(member_dim) + axis_positions.append(codes * stride) + else: + transposed.append(d) + axis_positions.append(np.arange(shape[i]) * stride) + + cell_rows = axis_positions[0] + for pos in axis_positions[1:]: + cell_rows = cell_rows[..., None] + pos + cell_rows = cell_rows.reshape(-1) + + coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + nterm = ds.sizes[TERM_DIM] + rows = np.repeat(cell_rows, nterm) + + keep = (vars_ != -1) & (coeffs != 0) & ~np.isnan(coeffs) + full_size = int(np.prod(shape, dtype=np.int64)) if shape else 1 + width = expr.model._xCounter + coo = scipy.sparse.coo_array( + (coeffs[keep], (rows[keep], vars_[keep])), shape=(full_size, width) + ) + csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum + + const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) + const = np.zeros(full_size) + np.add.at(const, cell_rows, np.where(np.isnan(const_vals), 0.0, const_vals)) + + return CSRPayload(csr, const, grid_dims, indexes, expr.model) + + +def try_csr_merge( + exprs: Any, dim: str, join: Any, kwargs: dict +) -> LinearExpression | None: + """ + Sparse branch of :func:`linopy.expressions.merge`. + + Returns a CSR-backed combined expression when every input is a plain + LinearExpression on the same result grid (CSR-backed, or dense and + convertible) — where any join produces the identical result — and None + otherwise to fall through to the dense path. + """ + from linopy.expressions import LinearExpression + + if dim != TERM_DIM or kwargs: + return None + if not all(type(e) is LinearExpression for e in exprs): + return None + payloads = [e._csr for e in exprs if e._csr is not None] + if not payloads: + return None + template = payloads[0] + if not all(template.same_grid(p) for p in payloads[1:]): + return None + + combined: CSRPayload | None = None + for e in exprs: + payload = e._csr + if payload is None: + payload = CSRPayload.from_expression(e, template) + if payload is None: + return None + combined = payload if combined is None else combined.add(payload) + assert combined is not None + return LinearExpression._from_csr(combined, exprs[0].model) + + +def extract_pending( + lhs: Any, sign: Any, rhs: Any +) -> tuple[CSRPayload, str, Any] | None: + """Return (payload, sign, rhs) if lhs is a realizable CSR constraint.""" + from linopy.common import is_constant + from linopy.constraints import Constraint + from linopy.expressions import LinearExpression + + if ( + isinstance(lhs, Constraint) + and lhs._pending is not None + and sign is None + and rhs is None + ): + expr, sign, rhs = lhs._pending + if expr._csr is None: # already materialized elsewhere + return None + lhs = expr + if not (isinstance(lhs, LinearExpression) and lhs._csr is not None): + return None + if not isinstance(sign, str) or rhs is None or not is_constant(rhs): + return None + payload = lhs._csr + rhs_da = _as_rhs_dataarray(rhs) + if rhs_da is None or not set(rhs_da.dims) <= set(payload.grid_dims): + return None + return payload, sign, rhs + + +def _as_rhs_dataarray(rhs: Any) -> DataArray | None: + from linopy.alignment import as_dataarray + + try: + da = as_dataarray(rhs) + except (TypeError, ValueError): + return None + if set(da.dims) & set(HELPER_DIMS): + return None + return da + + +def realize_csr_constraint( + model: Model, payload: CSRPayload, sign: str, rhs: Any, name: str +) -> CSRConstraint: + """ + Staple sign and rhs onto a CSR-backed lhs to form a CSRConstraint. + + Equivalent to materializing the sparse sum, adding the constraint and + freezing it — but the padded dense ``_term`` rectangle never exists; + peak memory is proportional to the number of nonzero terms. + """ + from linopy.common import maybe_replace_sign + from linopy.constraints import CSRConstraint + from linopy.semantics import check_user_nan, is_v1 + + sign = maybe_replace_sign(sign) + full_size = payload.n_cells + + # map label columns to dense positions in the active variable array + label_index = model.variables.label_index + label_to_pos = label_index.label_to_pos + coo = payload.csr.tocoo() + csr = scipy.sparse.csr_array( + scipy.sparse.coo_array( + (coo.data, (coo.coords[0], label_to_pos[coo.coords[1]])), + shape=(full_size, label_index.n_active_vars), + ) + ) + csr.eliminate_zeros() + + # rhs aligned by label to the grid; expression constants move to the rhs + rhs_da = _as_rhs_dataarray(rhs) + assert rhs_da is not None + if is_v1() and bool(rhs_da.isnull().any()): + check_user_nan() # §5: NaN in a user constant raises under v1 + if is_v1(): + # §8 parity with the dense path: a reordered or differing rhs index + # raises; the user aligns explicitly with .sel/.reindex + for d in rhs_da.dims: + if not rhs_da.get_index(d).equals(payload.indexes[str(d)]): + raise ValueError( + f"Coordinate mismatch on shared dimension {d!r} between " + "the rhs and the grouped result. Align the rhs with " + ".sel(...) / .reindex(...) before combining (§8)." + ) + rhs_da = rhs_da.reindex( + {d: payload.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan + ) + missing = {d: payload.indexes[d] for d in payload.grid_dims if d not in rhs_da.dims} + if missing: + rhs_da = rhs_da.expand_dims(missing) + rhs_flat = ( + rhs_da.transpose(*payload.grid_dims).to_numpy().reshape(-1) - payload.const + ) + + # label allocation, mirroring Model._allocate_constraint_labels; rows + # without terms or with NaN rhs are inactive, as in the frozen dense path + cindex = model._cCounter + model._cCounter += full_size + active = (np.diff(csr.indptr) > 0) & ~np.isnan(rhs_flat) + con_labels = np.arange(cindex, cindex + full_size)[active] + + return CSRConstraint( + csr[active], + con_labels, + rhs_flat[active], + sign, + coords=[payload.indexes[d] for d in payload.grid_dims], + model=model, + name=name, + cindex=cindex, + ) diff --git a/linopy/expressions.py b/linopy/expressions.py index 099f94d6..4725dcf6 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -480,7 +480,7 @@ def sum( self, use_fallback: bool = False, observed: bool = False, - lazy: bool | None = None, + sparse: bool | None = None, ) -> LinearExpression: """ Sum the expression over each group. @@ -494,17 +494,19 @@ def sum( Fall back to the previous, slower groupby-sum implementation, kept as an escape hatch. Leave at False unless the default misbehaves. Defaults to False. - lazy : bool, optional - Return the grouped sum unmaterialized: a long-format payload - behind the ordinary LinearExpression type, akin to dask-backed - xarray objects. Any operation without a lazy branch - transparently materializes today's dense result; a still-lazy - left-hand side reaching ``Model.add_constraints`` with - ``freeze=True`` is realized directly as a CSRConstraint, - skipping the group-size-padded term dimension. Requires v1 - semantics. Defaults to ``linopy.options["lazy_groupby"]`` - (only honored under v1). Only single-key groupers over an - existing dimension are held lazily. + sparse : bool, optional + Build the grouped sum in CSR form behind the ordinary + LinearExpression type (akin to dask-backed xarray objects): one + row per result cell, one column per variable, duplicate + variables summed and terms label-ordered — the group-size + padding of the dense kernel has no analog. Operations without a + sparse branch transparently expand to the dense rectangle in + this canonical form; a still-sparse left-hand side reaching + ``Model.add_constraints`` with ``freeze=True`` becomes a + CSRConstraint directly. Requires v1 semantics. Defaults to + ``linopy.options["sparse_groupby"]`` (only honored under v1). + Only single-key groupers over an existing dimension are held + sparse. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -530,22 +532,22 @@ def sum( None if use_fallback else _multikey_value_frame(group, self.data) ) - if lazy is None: - lazy = is_v1() and options["lazy_groupby"] - elif lazy and not is_v1(): + if sparse is None: + sparse = is_v1() and options["sparse_groupby"] + elif sparse and not is_v1(): raise ValueError( - "lazy groupby-sum requires v1 semantics; opt in with " + "sparse groupby-sum requires v1 semantics; opt in with " "linopy.options['semantics'] = 'v1'." ) - if lazy and not use_fallback and not observed and multikey_frame is None: + if sparse and not use_fallback and not observed and multikey_frame is None: series = group.to_pandas() if isinstance(group, DataArray) else group if isinstance(series, pd.Series) and series.index.name in self.data.dims: - from linopy.lazy import LazyGroupSum + from linopy.csr import CSRPayload expr = LinearExpression(self.data, self.model) group_name = str(getattr(series, "name", None) or "group") - payload = LazyGroupSum.from_grouper(expr, series, group_name) - return LinearExpression._from_lazy(payload, self.model) + payload = CSRPayload.from_grouper(expr, series, group_name) + return LinearExpression._from_csr(payload, self.model) if multikey_frame is not None: group = multikey_frame @@ -743,7 +745,7 @@ def sum(self, **kwargs: Any) -> LinearExpression: class BaseExpression(ABC): - __slots__ = ("_data", "_model", "_lazy") + __slots__ = ("_data", "_model", "_csr") __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 @@ -817,7 +819,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: self._model = model self._data = cast(Dataset, data) - self._lazy = None + self._csr = None def __repr__(self) -> str: """ @@ -918,8 +920,8 @@ def __neg__(self) -> Self: """ Get the negative of the expression. """ - if self._lazy is not None: - return self._from_lazy(self._lazy.scaled(-1.0), self._model) + if self._csr is not None: + return self._from_csr(self._csr.scaled(-1.0), self._model) return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) def _multiply_by_linear_expression( @@ -1495,19 +1497,19 @@ def type(self) -> str: @property def data(self) -> Dataset: - if self._data is None and self._lazy is not None: - # materialize through today's dense kernel; laziness is spent - self._data = self._lazy.materialize().data - self._lazy = None + if self._data is None and self._csr is not None: + # expand to the dense rectangle (canonical form); sparsity is spent + self._data = self._csr.materialize().data + self._csr = None return self._data @classmethod - def _from_lazy(cls, lazy: Any, model: Model) -> Self: - """Construct an expression backed by an unmaterialized LazyGroupSum.""" + def _from_csr(cls, payload: Any, model: Model) -> Self: + """Construct an expression backed by a CSRPayload.""" obj = cls.__new__(cls) obj._model = model obj._data = None # type: ignore[assignment] - obj._lazy = lazy + obj._csr = payload return obj @property @@ -1521,9 +1523,9 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: - if self._data is None and self._lazy is not None: - # answerable from the lazy payload without materializing - return tuple(self._lazy.grid_dims) + if self._data is None and self._csr is not None: + # answerable from the sparse payload without materializing + return tuple(self._csr.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @property @@ -1712,10 +1714,10 @@ def to_constraint( which are moved to the left-hand-side and constant values which are moved to the right-hand side. """ - if self._lazy is not None and isinstance(sign, str) and is_constant(rhs): - # carry the lazy payload on the constraint; its .data property + if self._csr is not None and isinstance(sign, str) and is_constant(rhs): + # carry the sparse payload on the constraint; its .data property # materializes through this method's dense path if needed - return constraints.Constraint._from_lazy(self, sign, rhs, self.model) + return constraints.Constraint._from_pending(self, sign, rhs, self.model) rhs = as_constant(rhs) if self.is_constant and is_constant(rhs): @@ -2354,8 +2356,8 @@ def __mul__( """ Multiply the expr by a factor. """ - if self._lazy is not None and isinstance(other, int | float | np.number): - return type(self)._from_lazy(self._lazy.scaled(float(other)), self._model) + if self._csr is not None and isinstance(other, int | float | np.number): + return type(self)._from_csr(self._csr.scaled(float(other)), self._model) other = as_constant(other) if isinstance(other, QuadraticExpression): return other.__rmul__(self) @@ -3128,11 +3130,11 @@ def merge( model = exprs[0].model if issubclass(cls, LinearExpression) and not has_quad_expression: - from linopy.lazy import try_lazy_merge + from linopy.csr import try_csr_merge - lazy_result = try_lazy_merge(exprs, dim=dim, join=join, kwargs=kwargs) - if lazy_result is not None: - return lazy_result + csr_result = try_csr_merge(exprs, dim=dim, join=join, kwargs=kwargs) + if csr_result is not None: + return csr_result data = [e.data if isinstance(e, linopy_types) else e for e in exprs] data = [fill_missing_coords(ds, fill_helper_dims=True) for ds in data] diff --git a/linopy/lazy.py b/linopy/lazy.py deleted file mode 100644 index 023a3321..00000000 --- a/linopy/lazy.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -Lazy groupby-sum: an internal long-format representation for grouped sums. - -``expr.groupby(g).sum(lazy=True)`` (or ``linopy.options["lazy_groupby"] = -True``) returns an ordinary :class:`~linopy.expressions.LinearExpression` -whose payload is a :class:`LazyGroupSum` — a list of ungrouped parts plus -their groupers — instead of the materialized dense dataset. Modeled on -dask-backed xarray objects: same public type, different backing, and any -operation without a lazy branch transparently materializes through the -``.data`` property (yielding exactly today's result, so laziness can only -delay the dense rectangle, never change semantics). - -The payoff comes at ``Model.add_constraints`` with ``freeze=True`` (or -``Model.freeze_constraints``): a still-lazy left-hand side is realized -directly as a :class:`~linopy.constraints.CSRConstraint` from the ungrouped -long triplets — scipy's COO->CSR duplicate summation *is* the group sum — -so the ``group-size-padded`` ``_term`` rectangle (issue #745) never exists. - -Operations with a lazy branch (everything else materializes): - -- unary minus and multiplication by a scalar (scales the ungrouped parts), -- ``merge`` along ``_term`` — and therefore ``+``/``-`` — with other lazy - expressions on the same grid or with dense expressions already on the - result grid, -- comparison with a constant rhs (``==``, ``<=``, ``>=``), carried lazily - on the resulting :class:`~linopy.constraints.Constraint`. -""" - -from __future__ import annotations - -from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any - -import numpy as np -import pandas as pd -import scipy.sparse -from xarray import DataArray - -from linopy.constants import HELPER_DIMS, TERM_DIM - -if TYPE_CHECKING: - from linopy.constraints import CSRConstraint - from linopy.expressions import LinearExpression - from linopy.model import Model - - -@dataclass(frozen=True) -class LazyPart: - """ - One summand: an ungrouped expression plus its grouper. - - ``grouper`` maps the labels of ``expr``'s member dim (the grouper's - index, named after that dim) to group labels. ``None`` marks an - identity part that already lives on the result grid. - """ - - expr: LinearExpression - grouper: pd.Series | None - - -@dataclass(frozen=True) -class LazyGroupSum: - """ - A sum of grouped/dense parts over a fixed result grid. - - ``grid_dims`` is the result's coordinate dim order (the grouped dim in - the member dim's position, as the dense kernel produces it) and - ``indexes`` maps each grid dim to its pandas index. - """ - - parts: tuple[LazyPart, ...] - group_dim: str - grid_dims: tuple[str, ...] - indexes: dict[str, pd.Index] - - @classmethod - def from_grouper( - cls, expr: LinearExpression, grouper: pd.Series, group_dim: str - ) -> LazyGroupSum: - member_dim = str(grouper.index.name) - if member_dim in expr.data.indexes: - # labels checked equal upstream; conform the order - grouper = grouper.reindex(expr.data.indexes[member_dim]) - elif len(grouper) != expr.data.sizes[member_dim]: - raise ValueError(f"grouper length does not match dimension {member_dim!r}") - codes, uniques = pd.factorize(grouper, sort=True) - if (codes == -1).any(): - raise ValueError( - "Cannot group by a pandas object containing NaN values. " - "Drop or fill the corresponding entries before grouping." - ) - grid_dims = tuple( - group_dim if d == member_dim else str(d) for d in expr.coord_dims - ) - indexes = { - d: expr.data.get_index(d).rename(d) - for d in expr.coord_dims - if d != member_dim - } - indexes[group_dim] = pd.Index(uniques, name=group_dim) - return cls((LazyPart(expr, grouper),), group_dim, grid_dims, indexes) - - @property - def group_index(self) -> pd.Index: - return self.indexes[self.group_dim] - - def scaled(self, factor: float) -> LazyGroupSum: - parts = tuple(LazyPart(p.expr * factor, p.grouper) for p in self.parts) - return replace(self, parts=parts) - - def same_grid(self, other: LazyGroupSum) -> bool: - return ( - self.group_dim == other.group_dim - and self.grid_dims == other.grid_dims - and all(self.indexes[d].equals(other.indexes[d]) for d in self.grid_dims) - ) - - def wrap_dense(self, expr: LinearExpression) -> LazyPart | None: - """Wrap a dense expression already on the result grid as identity part.""" - if set(expr.coord_dims) != set(self.grid_dims): - return None - for d in expr.coord_dims: - if not expr.data.get_index(d).equals(self.indexes[str(d)]): - return None - return LazyPart(expr, None) - - def materialize(self) -> LinearExpression: - """Left-fold the parts through today's dense kernel.""" - result: LinearExpression | None = None - for part in self.parts: - if part.grouper is None: - term = part.expr - else: - term = part.expr.groupby(part.grouper).sum(lazy=False) - result = term if result is None else result + term - assert result is not None - return result - - -def try_lazy_merge( - exprs: Any, dim: str, join: Any, kwargs: dict -) -> LinearExpression | None: - """ - Lazy branch of :func:`linopy.expressions.merge`. - - Returns a lazy combined expression when every input is a plain - LinearExpression on the same result grid (lazy, or dense wrappable as - an identity part) — where any join produces the identical result — and - None otherwise to fall through to the dense path. - """ - from linopy.expressions import LinearExpression - - if dim != TERM_DIM or kwargs: - return None - if not all(type(e) is LinearExpression for e in exprs): - return None - lazies = [e._lazy for e in exprs if e._lazy is not None] - if not lazies: - return None - template = lazies[0] - if not all(template.same_grid(lz) for lz in lazies[1:]): - return None - - parts: list[LazyPart] = [] - for e in exprs: - if e._lazy is not None: - parts.extend(e._lazy.parts) - else: - part = template.wrap_dense(e) - if part is None: - return None - parts.append(part) - combined = replace(template, parts=tuple(parts)) - return LinearExpression._from_lazy(combined, exprs[0].model) - - -def extract_lazy(lhs: Any, sign: Any, rhs: Any) -> tuple[LazyGroupSum, str, Any] | None: - """Return (payload, sign, rhs) if lhs is a realizable lazy constraint.""" - from linopy.common import is_constant - from linopy.constraints import Constraint - from linopy.expressions import LinearExpression - - if ( - isinstance(lhs, Constraint) - and lhs._lazy is not None - and sign is None - and rhs is None - ): - expr, sign, rhs = lhs._lazy - if expr._lazy is None: # already materialized elsewhere - return None - lhs = expr - if not (isinstance(lhs, LinearExpression) and lhs._lazy is not None): - return None - if not isinstance(sign, str) or rhs is None or not is_constant(rhs): - return None - lazy = lhs._lazy - rhs_da = _as_rhs_dataarray(rhs) - if rhs_da is None or not set(rhs_da.dims) <= set(lazy.grid_dims): - return None - return lazy, sign, rhs - - -def _as_rhs_dataarray(rhs: Any) -> DataArray | None: - from linopy.alignment import as_dataarray - - try: - da = as_dataarray(rhs) - except (TypeError, ValueError): - return None - if set(da.dims) & set(HELPER_DIMS): - return None - return da - - -def _part_entries( - part: LazyPart, lazy: LazyGroupSum, label_to_pos: np.ndarray -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """ - Long-form (rows, cols, coeffs) triplets of one part plus its const. - - Rows are flat positions in the ``grid_dims`` constraint grid (C order). - Returns (rows_per_term, cols, coeffs, cell_rows_flat, const_flat_values). - """ - ds = part.expr.data - shape = tuple(len(lazy.indexes[d]) for d in lazy.grid_dims) - strides = np.array( - [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] - ) - - # per grid dim: the position of each of the part's cells along that dim - transposed: list[str] = [] - axis_positions: list[np.ndarray] = [] - for d, stride in zip(lazy.grid_dims, strides): - if part.grouper is not None and d == lazy.group_dim: - member_dim = str(part.grouper.index.name) - transposed.append(member_dim) - codes = lazy.group_index.get_indexer(part.grouper.to_numpy()) - axis_positions.append(codes * stride) - else: - transposed.append(d) - axis_positions.append(np.arange(shape[lazy.grid_dims.index(d)]) * stride) - - cell_rows = axis_positions[0] - for pos in axis_positions[1:]: - cell_rows = cell_rows[..., None] + pos - cell_rows = cell_rows.reshape(-1) - - coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) - vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) - nterm = ds.sizes[TERM_DIM] - rows = np.repeat(cell_rows, nterm) - - keep = (vars_ != -1) & (coeffs != 0) & ~np.isnan(coeffs) - rows, vars_, coeffs = rows[keep], vars_[keep], coeffs[keep] - - const = ds.const.transpose(*transposed).to_numpy().reshape(-1) - return rows, label_to_pos[vars_], coeffs, cell_rows, const - - -def realize_lazy_constraint( - model: Model, lazy: LazyGroupSum, sign: str, rhs: Any, name: str -) -> CSRConstraint: - """ - Realize ``lazy sign rhs`` directly as a CSRConstraint. - - Equivalent to materializing the lazy sum, adding the constraint and - freezing it — but the padded dense ``_term`` rectangle never exists; - peak memory is proportional to the number of nonzero terms. - """ - from linopy.common import maybe_replace_sign - from linopy.constraints import CSRConstraint - from linopy.semantics import check_user_nan, is_v1 - - sign = maybe_replace_sign(sign) - shape = tuple(len(lazy.indexes[d]) for d in lazy.grid_dims) - full_size = int(np.prod(shape, dtype=np.int64)) if shape else 1 - - label_index = model.variables.label_index - label_to_pos = label_index.label_to_pos - - rows_l, cols_l, data_l = [], [], [] - const_flat = np.zeros(full_size) - for part in lazy.parts: - rows, cols, coeffs, cell_rows, const = _part_entries(part, lazy, label_to_pos) - rows_l.append(rows) - cols_l.append(cols) - data_l.append(coeffs) - np.add.at(const_flat, cell_rows, np.where(np.isnan(const), 0.0, const)) - - coo = scipy.sparse.coo_array( - (np.concatenate(data_l), (np.concatenate(rows_l), np.concatenate(cols_l))), - shape=(full_size, label_index.n_active_vars), - ) - csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum - csr.eliminate_zeros() - - # rhs aligned by label to the grid; expression constants move to the rhs - rhs_da = _as_rhs_dataarray(rhs) - assert rhs_da is not None - if is_v1() and bool(rhs_da.isnull().any()): - check_user_nan() # §5: NaN in a user constant raises under v1 - if is_v1(): - # §8 parity with the dense path: a reordered or differing rhs index - # raises; the user aligns explicitly with .sel/.reindex - for d in rhs_da.dims: - if not rhs_da.get_index(d).equals(lazy.indexes[str(d)]): - raise ValueError( - f"Coordinate mismatch on shared dimension {d!r} between " - "the rhs and the grouped result. Align the rhs with " - ".sel(...) / .reindex(...) before combining (§8)." - ) - rhs_da = rhs_da.reindex( - {d: lazy.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan - ) - missing = {d: lazy.indexes[d] for d in lazy.grid_dims if d not in rhs_da.dims} - if missing: - rhs_da = rhs_da.expand_dims(missing) - rhs_flat = rhs_da.transpose(*lazy.grid_dims).to_numpy().reshape(-1) - const_flat - - # label allocation, mirroring Model._allocate_constraint_labels; rows - # without terms or with NaN rhs are inactive, as in the frozen dense path - cindex = model._cCounter - model._cCounter += full_size - active = (np.diff(csr.indptr) > 0) & ~np.isnan(rhs_flat) - con_labels = np.arange(cindex, cindex + full_size)[active] - - return CSRConstraint( - csr[active], - con_labels, - rhs_flat[active], - sign, - coords=[lazy.indexes[d] for d in lazy.grid_dims], - model=model, - name=name, - cindex=cindex, - ) diff --git a/linopy/model.py b/linopy/model.py index 8c7e2fc6..62fbcfbe 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1139,16 +1139,16 @@ def add_constraints( name = self._resolve_constraint_name(name) - # a still-lazy grouped lhs with freeze on realizes directly as a - # CSRConstraint from long triplets, skipping the padded dense form + # a still-sparse grouped lhs with freeze on becomes a CSRConstraint + # directly (sign and rhs stapled on), skipping the padded dense form resolved_freeze = self.freeze_constraints if freeze is None else freeze if resolved_freeze and mask is None and not self.chunk: - from linopy.lazy import extract_lazy, realize_lazy_constraint + from linopy.csr import extract_pending, realize_csr_constraint - extracted = extract_lazy(lhs, sign, rhs) + extracted = extract_pending(lhs, sign, rhs) if extracted is not None: - lazy_sum, lazy_sign, lazy_rhs = extracted - con = realize_lazy_constraint(self, lazy_sum, lazy_sign, lazy_rhs, name) + payload, csr_sign, csr_rhs = extracted + con = realize_csr_constraint(self, payload, csr_sign, csr_rhs, name) return self.constraints.add(con) if sign is not None: diff --git a/test/test_lazy_groupby.py b/test/test_sparse_groupby.py similarity index 69% rename from test/test_lazy_groupby.py rename to test/test_sparse_groupby.py index 23e3461f..12e06ec0 100644 --- a/test/test_lazy_groupby.py +++ b/test/test_sparse_groupby.py @@ -1,5 +1,5 @@ """ -Tests for lazy groupby-sum (linopy.lazy): type stability, transparent +Tests for sparse groupby-sum (linopy.csr): type stability, transparent materialization, and direct CSR realization under freeze. v1-only feature. """ @@ -22,7 +22,7 @@ def require_v1() -> None: if not is_v1(): - pytest.skip("lazy groupby-sum is gated behind v1 semantics") + pytest.skip("sparse groupby-sum is gated behind v1 semantics") def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): @@ -50,11 +50,11 @@ def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): return m, gen_p, flow, eff, gbus, bus0, bus1, load, buses -def balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy): +def balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse): return ( - (eff * gen_p).groupby(gbus).sum(lazy=lazy) - + (1.0 * flow).groupby(bus0).sum(lazy=lazy) - - (1.0 * flow).groupby(bus1).sum(lazy=lazy) + (eff * gen_p).groupby(gbus).sum(sparse=sparse) + + (1.0 * flow).groupby(bus0).sum(sparse=sparse) + - (1.0 * flow).groupby(bus1).sum(sparse=sparse) ) @@ -66,54 +66,54 @@ def canon(df: pl.DataFrame) -> pl.DataFrame: ) -def test_lazy_requires_v1(): +def test_csr_requires_v1(): m, gen_p, flow, eff, gbus, *_ = base_model() if is_v1(): - res = (eff * gen_p).groupby(gbus).sum(lazy=True) + res = (eff * gen_p).groupby(gbus).sum(sparse=True) assert type(res) is LinearExpression return with pytest.raises(ValueError, match="requires v1 semantics"): - (eff * gen_p).groupby(gbus).sum(lazy=True) + (eff * gen_p).groupby(gbus).sum(sparse=True) # the option is not honored under legacy: result is eager and dense - linopy.options["lazy_groupby"] = True + linopy.options["sparse_groupby"] = True try: res = (eff * gen_p).groupby(gbus).sum() finally: - linopy.options["lazy_groupby"] = False - assert res._lazy is None + linopy.options["sparse_groupby"] = False + assert res._csr is None -def test_lazy_is_plain_linear_expression_and_materializes_identically(): +def test_csr_is_plain_linear_expression_and_materializes_equivalently(): require_v1() m, gen_p, flow, eff, gbus, *_ = base_model() - lazy = (eff * gen_p).groupby(gbus).sum(lazy=True) + sparse = (eff * gen_p).groupby(gbus).sum(sparse=True) eager = (eff * gen_p).groupby(gbus).sum() - assert type(lazy) is LinearExpression + assert type(sparse) is LinearExpression # comparing data materializes; the result must be exactly today's - assert_linequal(lazy, eager) + assert_linequal(sparse, eager) -def test_lazy_composition_materializes_identically(): +def test_csr_composition_materializes_equivalently(): require_v1() m, gen_p, flow, eff, gbus, bus0, bus1, *_ = base_model() - lazy = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=True) - eager = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, lazy=False) - assert type(lazy) is LinearExpression - assert_linequal(lazy, eager) + sparse = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=True) + eager = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=False) + assert type(sparse) is LinearExpression + assert_linequal(sparse, eager) -def test_scalar_ops_stay_lazy(): +def test_scalar_ops_stay_csr(): require_v1() m, gen_p, flow, eff, gbus, *_ = base_model() - lazy = -2.0 * (eff * gen_p).groupby(gbus).sum(lazy=True) - assert lazy._lazy is not None # still unmaterialized after neg/scale + sparse = -2.0 * (eff * gen_p).groupby(gbus).sum(sparse=True) + assert sparse._csr is not None # still unmaterialized after neg/scale eager = -2.0 * (eff * gen_p).groupby(gbus).sum() # scaling after grouping normalizes padded-slot fills while scaling # before grouping keeps fresh pads; masked-slot fill is not contractual, # so compare with masked slots blanked on both sides from linopy.testing import _sort_by_vars_along_term - a, b = _sort_by_vars_along_term(lazy), _sort_by_vars_along_term(eager) + a, b = _sort_by_vars_along_term(sparse), _sort_by_vars_along_term(eager) assert (a.vars == b.vars).all() xr.testing.assert_allclose( a.coeffs.where(a.vars != -1), b.coeffs.where(b.vars != -1) @@ -124,11 +124,11 @@ def test_scalar_ops_stay_lazy(): def test_freeze_realizes_csr_without_dense_rectangle(): require_v1() m1, *rest = base_model() - lhs1 = balance_lhs(*rest[:6], lazy=False) + lhs1 = balance_lhs(*rest[:6], sparse=False) c1 = m1.add_constraints(lhs1 == rest[6], name="balance") m2, *rest2 = base_model() - lhs2 = balance_lhs(*rest2[:6], lazy=True) + lhs2 = balance_lhs(*rest2[:6], sparse=True) c2 = m2.add_constraints(lhs2 == rest2[6], name="balance", freeze=True) assert isinstance(c2, CSRConstraint) @@ -147,34 +147,37 @@ def test_freeze_false_falls_back_to_identical_dense_constraint(): require_v1() m1, *rest = base_model() c1 = m1.add_constraints( - balance_lhs(*rest[:6], lazy=False) == rest[6], name="balance" + balance_lhs(*rest[:6], sparse=False) == rest[6], name="balance" ) m2, *rest2 = base_model() c2 = m2.add_constraints( - balance_lhs(*rest2[:6], lazy=True) == rest2[6], name="balance" + balance_lhs(*rest2[:6], sparse=True) == rest2[6], name="balance" ) assert isinstance(c2, Constraint) - assert_conequal(c1, c2) + # the sparse lhs materializes in canonical form (terms label-sorted, + # padding trailing), so compare mathematically, not layout-strictly + assert_conequal(c1, c2, strict=False) + assert np.array_equal(c1.labels.values, c2.labels.values) -def test_option_gates_lazy_and_freeze_model_default(): +def test_option_gates_csr_and_freeze_model_default(): require_v1() m, *rest = base_model() m.freeze_constraints = True - linopy.options["lazy_groupby"] = True + linopy.options["sparse_groupby"] = True try: - lhs = balance_lhs(*rest[:6], lazy=None) + lhs = balance_lhs(*rest[:6], sparse=None) con = m.add_constraints(lhs == rest[6], name="balance") finally: - linopy.options["lazy_groupby"] = False + linopy.options["sparse_groupby"] = False assert isinstance(con, CSRConstraint) -def test_materialized_lazy_still_freezes_via_dense_path(): +def test_materialized_csr_still_freezes_via_dense_path(): require_v1() m, *rest = base_model() - lhs = balance_lhs(*rest[:6], lazy=True) + lhs = balance_lhs(*rest[:6], sparse=True) _ = lhs.nterm # force materialization before constraint creation con = m.add_constraints(lhs == rest[6], name="balance", freeze=True) assert isinstance(con, CSRConstraint) @@ -187,7 +190,7 @@ def test_nan_rhs_raises_on_both_paths(): load[0, 0] = np.nan with pytest.raises(ValueError, match="NaN"): m1.add_constraints( - balance_lhs(*rest[:6], lazy=False) == load, name="bal", freeze=True + balance_lhs(*rest[:6], sparse=False) == load, name="bal", freeze=True ) m2, *rest2 = base_model() @@ -195,7 +198,7 @@ def test_nan_rhs_raises_on_both_paths(): load2[0, 0] = np.nan with pytest.raises(ValueError, match="NaN"): m2.add_constraints( - balance_lhs(*rest2[:6], lazy=True) == load2, name="bal", freeze=True + balance_lhs(*rest2[:6], sparse=True) == load2, name="bal", freeze=True ) @@ -205,14 +208,14 @@ def test_reordered_rhs_raises_on_both_paths(): load = rest[6].isel(bus=slice(None, None, -1)) # same labels, reversed with pytest.raises(ValueError, match="[Cc]oordinate"): m1.add_constraints( - balance_lhs(*rest[:6], lazy=False) == load, name="bal", freeze=True + balance_lhs(*rest[:6], sparse=False) == load, name="bal", freeze=True ) m2, *rest2 = base_model() load2 = rest2[6].isel(bus=slice(None, None, -1)) with pytest.raises(ValueError, match="[Cc]oordinate"): m2.add_constraints( - balance_lhs(*rest2[:6], lazy=True) == load2, name="bal", freeze=True + balance_lhs(*rest2[:6], sparse=True) == load2, name="bal", freeze=True ) @@ -222,7 +225,7 @@ def test_nan_grouper_raises_eagerly(): gbus = gbus.copy() gbus.iloc[0] = np.nan with pytest.raises(ValueError, match="NaN values"): - (1.0 * gen_p).groupby(gbus).sum(lazy=True) + (1.0 * gen_p).groupby(gbus).sum(sparse=True) def test_lp_files_identical(tmp_path): @@ -230,12 +233,12 @@ def test_lp_files_identical(tmp_path): sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) m1, *rest = base_model(gens_per_bus=sizes) - m1.add_constraints(balance_lhs(*rest[:6], lazy=False) == rest[6], name="balance") + m1.add_constraints(balance_lhs(*rest[:6], sparse=False) == rest[6], name="balance") m1.add_objective((1.0 * rest[0]).sum()) m2, *rest2 = base_model(gens_per_bus=sizes) m2.add_constraints( - balance_lhs(*rest2[:6], lazy=True) == rest2[6], name="balance", freeze=True + balance_lhs(*rest2[:6], sparse=True) == rest2[6], name="balance", freeze=True ) m2.add_objective((1.0 * rest2[0]).sum()) @@ -252,7 +255,7 @@ def canon_lp(text: str) -> list[str]: buf = [] return out + sorted(buf) - f1, f2 = tmp_path / "eager.lp", tmp_path / "lazy.lp" + f1, f2 = tmp_path / "eager.lp", tmp_path / "sparse.lp" m1.to_file(f1) m2.to_file(f2) assert canon_lp(f1.read_text()) == canon_lp(f2.read_text()) From f083482aa599b299fa5a0820fb73f6c42a7cfe01 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:17:31 +0200 Subject: [PATCH 4/5] bench(patterns): sparse sibling for nodal_balance; drop dev-script repro Replace dev-scripts/proto_deferred_bench.py with a suite entry: nodal_balance_sparse builds the identical balance via sum(sparse=True) + freeze=True under v1 (phases: build/matrices/to_lp). Paired with the existing nodal_balance severity sweep, the padding cost becomes CI-visible: pytest benchmarks/ -k 'nodal_balance and build' --benchmark-memory severity 0 50 100 dense (KiB) 951 5,460 9,909 sparse (KiB) 1,524 1,524 1,524 (and ~1.6x faster builds) Co-Authored-By: Claude Fable 5 --- benchmarks/patterns/nodal_balance.py | 48 ++++++++++++++++++++-- dev-scripts/proto_deferred_bench.py | 60 ---------------------------- 2 files changed, 44 insertions(+), 64 deletions(-) delete mode 100644 dev-scripts/proto_deferred_bench.py diff --git a/benchmarks/patterns/nodal_balance.py b/benchmarks/patterns/nodal_balance.py index 458df39a..f05d6989 100644 --- a/benchmarks/patterns/nodal_balance.py +++ b/benchmarks/patterns/nodal_balance.py @@ -7,6 +7,12 @@ result's ``_term`` axis blows up — most of it fill. ``severity`` dials that skew; the build's peak memory is expected to climb steeply with it on the current (dense) kernel. + +``nodal_balance_sparse`` builds the identical constraint through the +CSR-backed path (``sum(sparse=True)`` + ``freeze=True`` under v1): the +grouped sum never materializes the padded rectangle and is realized directly +as a CSRConstraint, so its peak memory should stay flat across the severity +sweep — the pair makes the padding cost visible. """ from __future__ import annotations @@ -16,7 +22,14 @@ import xarray as xr import linopy -from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern +from benchmarks.registry import ( + BUILD, + MATRICES, + SEVERITIES, + TO_LP, + BenchSpec, + register_pattern, +) N_GEN = 2000 N_BUS = 50 @@ -43,7 +56,7 @@ def _bus_of_gen(severity: int) -> np.ndarray: return bus -def build_nodal_balance(severity: int) -> linopy.Model: +def _build(severity: int, sparse: bool) -> linopy.Model: gens = pd.RangeIndex(N_GEN, name="gen") time = pd.RangeIndex(N_TIME, name="time") buses = pd.RangeIndex(N_BUS, name="bus") @@ -53,15 +66,29 @@ def build_nodal_balance(severity: int) -> linopy.Model: gen = m.add_variables(lower=0, coords=[gens, time], name="gen") bus_of_gen = pd.Series(_bus_of_gen(severity), index=gens, name="bus") - supply = (1 * gen).groupby(bus_of_gen).sum() + supply = (1 * gen).groupby(bus_of_gen).sum(sparse=sparse) demand = xr.DataArray( rng.uniform(10.0, 100.0, size=(N_BUS, N_TIME)), coords=[buses, time] ) - m.add_constraints(supply == demand, name="balance") + m.add_constraints(supply == demand, name="balance", freeze=sparse) m.add_objective(gen.sum()) return m +def build_nodal_balance(severity: int) -> linopy.Model: + return _build(severity, sparse=False) + + +def build_nodal_balance_sparse(severity: int) -> linopy.Model: + """The same balance via the sparse groupby + frozen CSR path (v1-only).""" + previous = linopy.options["semantics"] + linopy.options["semantics"] = "v1" + try: + return _build(severity, sparse=True) + finally: + linopy.options["semantics"] = previous + + SPEC = register_pattern( BenchSpec( name="nodal_balance", @@ -70,3 +97,16 @@ def build_nodal_balance(severity: int) -> linopy.Model: axis="severity", ) ) + +# narrower phases: the model holds a CSRConstraint from the start, which the +# build/matrix/LP paths consume natively; the netcdf/solver-API round-trips +# are the dense Constraint's territory and not what this pair measures +SPARSE_SPEC = register_pattern( + BenchSpec( + name="nodal_balance_sparse", + build=build_nodal_balance_sparse, + sweep=SEVERITIES, + axis="severity", + phases=frozenset({BUILD, MATRICES, TO_LP}), + ) +) diff --git a/dev-scripts/proto_deferred_bench.py b/dev-scripts/proto_deferred_bench.py deleted file mode 100644 index ce20b28d..00000000 --- a/dev-scripts/proto_deferred_bench.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -Benchmark: eager groupby-sum balance constraint vs sparse CSR realization. - -Scenario from issue #745 (hub-skewed nodal balance): 120 buses, one hub with -HUB generators, 100 on each other bus, ring of 120 lines, 24 snapshots. -Build-only, using the real API: ``groupby(g).sum(sparse=...)`` + -``add_constraints(lhs == load, freeze=...)``. - -Run under memray, one variant per process: - memray run -o eager.bin proto_deferred_bench.py eager [hub] - memray run -o sparse.bin proto_deferred_bench.py sparse [hub] -""" - -from __future__ import annotations - -import sys -import time - -sys.path.insert( - 0, str(__import__("pathlib").Path(__file__).resolve().parent.parent / "test") -) -from test_sparse_groupby import balance_lhs, base_model - -N_BUS = 120 -HUB = int(sys.argv[2]) if len(sys.argv) > 2 else 8000 -GENS_PER_BUS = tuple([HUB] + [100] * (N_BUS - 1)) -N_SNAP = 24 - - -def main() -> None: - import linopy - - linopy.options["semantics"] = "v1" # the sparse path is gated behind v1 - variant = sys.argv[1] - m, gen_p, flow, eff, gbus, bus0, bus1, load, buses = base_model( - gens_per_bus=GENS_PER_BUS, n_snap=N_SNAP - ) - start = time.perf_counter() - - if variant == "setup": - print(f"setup only: {time.perf_counter() - start:.2f}s") - return - - sparse = variant == "sparse" - lhs = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=sparse) - con = m.add_constraints(lhs == load, name="balance", freeze=sparse) - build_s = time.perf_counter() - start - - start = time.perf_counter() - n_rows = con.to_polars().height - export_s = time.perf_counter() - start - - print( - f"{variant}: build {build_s:.2f}s, to_polars {export_s:.2f}s, " - f"type={type(con).__name__}, polars term rows={n_rows}" - ) - - -if __name__ == "__main__": - main() From c207a6bcf3cba602a11a6ecac9be53328dda37d5 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:29:33 +0200 Subject: [PATCH 5/5] refactor: condense sparse-groupby code and docs DRY the payload module (shared rhs-alignment helper, comprehension-based assembly), fold inline comments into docstrings, shrink the module and kwarg docs. No behavior change: sparse suite, nodal_balance benchmarks smoke and core expression/constraint tests unchanged. Co-Authored-By: Claude Fable 5 --- benchmarks/patterns/nodal_balance.py | 3 - linopy/constraints.py | 2 +- linopy/csr.py | 235 ++++++++++++--------------- linopy/expressions.py | 21 +-- linopy/model.py | 2 - test/test_sparse_groupby.py | 11 +- 6 files changed, 109 insertions(+), 165 deletions(-) diff --git a/benchmarks/patterns/nodal_balance.py b/benchmarks/patterns/nodal_balance.py index f05d6989..4cc2e688 100644 --- a/benchmarks/patterns/nodal_balance.py +++ b/benchmarks/patterns/nodal_balance.py @@ -98,9 +98,6 @@ def build_nodal_balance_sparse(severity: int) -> linopy.Model: ) ) -# narrower phases: the model holds a CSRConstraint from the start, which the -# build/matrix/LP paths consume natively; the netcdf/solver-API round-trips -# are the dense Constraint's territory and not what this pair measures SPARSE_SPEC = register_pattern( BenchSpec( name="nodal_balance_sparse", diff --git a/linopy/constraints.py b/linopy/constraints.py index b73b872f..8a3c625b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1193,7 +1193,7 @@ def _from_pending( def data(self) -> Dataset: if self._data is None and self._pending is not None: lhs, sign, rhs = self._pending - lhs.data # noqa: B018 materialize; clears lhs's sparse payload + lhs.data # noqa: B018 self._data = lhs.to_constraint(sign, rhs).data self._assigned = "labels" in self._data self._pending = None diff --git a/linopy/csr.py b/linopy/csr.py index 474c9a6c..fecd7e39 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -1,39 +1,18 @@ """ CSR-backed expressions: a sparse internal representation for grouped sums. -``expr.groupby(g).sum(sparse=True)`` (or ``linopy.options["sparse_groupby"] -= True`` under v1) returns an ordinary -:class:`~linopy.expressions.LinearExpression` whose payload is a -:class:`CSRPayload` — the expression as a scipy CSR matrix (rows = flat -cells of the coordinate grid, columns = variable labels, values = -coefficients) plus a dense const vector — instead of the materialized -dense dataset. Modeled on dask-backed xarray objects: same public type, -different backing; any operation without a sparse branch transparently -converts to the dense form through the ``.data`` property. - -The CSR form is canonical: duplicate variables within a cell are summed -(``2x + 3x -> 5x``) and terms are ordered by variable label, so the -``_term`` axis is ragged by construction — the group-size padding of issue -#745 has no analog. Operations become sparse linear algebra: - -- grouping scatters members into group rows (conceptually ``G @ A`` with a - 0/1 grouping matrix), -- ``merge`` along ``_term`` — and therefore ``+``/``-`` — is sparse matrix - addition, -- unary minus and scalar multiplication scale the values, -- ``== rhs`` with a constant is carried as a pending payload on the - :class:`~linopy.constraints.Constraint`, and ``Model.add_constraints`` - with ``freeze=True`` staples sign and rhs on to produce a - :class:`~linopy.constraints.CSRConstraint` directly. - -Converting back to the dense rectangle yields the mathematically identical -expression in canonical form (not the eager kernel's exact term layout); -this is why the feature is gated behind v1 semantics, where the term -layout is explicitly non-contractual. - -Columns are raw variable labels (not positions), so payloads stay valid -when variables are added to the model later; realization maps labels to -dense positions via the model's label index. +``expr.groupby(g).sum(sparse=True)`` (or ``linopy.options["sparse_groupby"]`` +under v1) returns an ordinary :class:`~linopy.expressions.LinearExpression` +backed by a :class:`CSRPayload` instead of the dense dataset — same public +type, different backing, akin to dask-backed xarray objects. The CSR form is +canonical (duplicate variables summed, terms label-ordered) and ragged along +``_term``, so the group-size padding of issue #745 has no analog; grouping, +``merge``/``+``/``-`` and scaling become sparse linear algebra, and +``Model.add_constraints(..., freeze=True)`` staples sign and rhs on to form a +:class:`~linopy.constraints.CSRConstraint` directly. Anything without a +sparse branch expands through ``.data`` to the mathematically identical dense +rectangle in canonical term layout — the reason the feature is v1-gated, +where term layout is non-contractual. """ from __future__ import annotations @@ -59,8 +38,10 @@ class CSRPayload: """ An expression as ``A @ x + c`` over a fixed coordinate grid. - ``csr`` has one row per flat grid cell (C order over ``grid_dims``) - and one column per variable label; ``const`` is the per-cell constant. + ``csr`` has one row per flat grid cell (C order over ``grid_dims``) and + one column per raw variable label — label columns stay valid when + variables are added to the model later; realization maps them to dense + positions. ``const`` is the per-cell constant. """ csr: scipy.sparse.csr_array @@ -81,10 +62,15 @@ def n_cells(self) -> int: def from_grouper( cls, expr: LinearExpression, grouper: pd.Series, group_dim: str ) -> CSRPayload: - """Build the grouped sum directly in CSR form (no padded rectangle).""" + """ + Build the grouped sum directly in CSR form (no padded rectangle). + + The grouper is conformed to the expression's member index by label + (upstream alignment checks guarantee equal label sets) and group + labels are sorted, matching the dense kernel's output grid. + """ member_dim = str(grouper.index.name) if member_dim in expr.data.indexes: - # labels checked equal upstream; conform the order grouper = grouper.reindex(expr.data.indexes[member_dim]) elif len(grouper) != expr.data.sizes[member_dim]: raise ValueError(f"grouper length does not match dimension {member_dim!r}") @@ -109,7 +95,7 @@ def from_grouper( def from_expression( cls, expr: LinearExpression, template: CSRPayload ) -> CSRPayload | None: - """Convert a dense expression on the template's grid to CSR form.""" + """Convert a dense expression on the template's grid, else None.""" if set(expr.coord_dims) != set(template.grid_dims): return None for d in expr.coord_dims: @@ -131,18 +117,14 @@ def same_grid(self, other: CSRPayload) -> bool: def add(self, other: CSRPayload) -> CSRPayload: """Sparse matrix addition == merge along the term dimension.""" - a, b = self.csr, other.csr - width = max(a.shape[1], b.shape[1]) - a, b = _widen(a, width), _widen(b, width) - return replace(self, csr=a + b, const=self.const + other.const) + width = max(self.csr.shape[1], other.csr.shape[1]) + csr = _widen(self.csr, width) + _widen(other.csr, width) + return replace(self, csr=csr, const=self.const + other.const) def materialize(self) -> LinearExpression: """ - Expand to the dense rectangle in canonical form. - - Terms are variable-label-ordered and duplicates are summed; the - term axis is padded to the widest cell with the usual fill - (vars=-1, coeffs=NaN). + Expand to the dense rectangle in canonical form: terms label-ordered, + duplicates summed, padded to the widest cell with the usual fill. """ from linopy.expressions import LinearExpression @@ -152,24 +134,23 @@ def materialize(self) -> LinearExpression: lengths = np.diff(csr.indptr) nterm = max(int(lengths.max()) if len(lengths) else 0, 1) - labels_dtype = self.model._dtypes["labels"] - vars_flat = np.full((self.n_cells, nterm), -1, dtype=labels_dtype) + vars_flat = np.full( + (self.n_cells, nterm), -1, dtype=self.model._dtypes["labels"] + ) coeffs_flat = np.full((self.n_cells, nterm), np.nan) rows = np.repeat(np.arange(self.n_cells), lengths) pos = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], lengths) vars_flat[rows, pos] = csr.indices coeffs_flat[rows, pos] = csr.data - shape = self.shape dims = (*self.grid_dims, TERM_DIM) - coords = {d: self.indexes[d] for d in self.grid_dims} ds = Dataset( { - "coeffs": (dims, coeffs_flat.reshape(*shape, nterm)), - "vars": (dims, vars_flat.reshape(*shape, nterm)), - "const": (self.grid_dims, self.const.reshape(shape)), + "coeffs": (dims, coeffs_flat.reshape(*self.shape, nterm)), + "vars": (dims, vars_flat.reshape(*self.shape, nterm)), + "const": (self.grid_dims, self.const.reshape(self.shape)), }, - coords=coords, + coords={d: self.indexes[d] for d in self.grid_dims}, ) return LinearExpression(ds, self.model) @@ -191,28 +172,21 @@ def _assemble( codes: np.ndarray, ) -> CSRPayload: """ - Scatter an expression's terms into grid rows (conceptually G @ A). - - ``expr``'s ``member_dim`` scatters into the grid dim ``scatter_dim`` - at the row positions given by ``codes``; every other grid dim maps - one-to-one. + Scatter an expression's terms into grid rows (conceptually ``G @ A``): + ``member_dim`` lands in the grid dim ``scatter_dim`` at row positions + ``codes``, every other grid dim maps one-to-one, and the COO→CSR + conversion sums duplicates — which is the group sum. The constant is + reduced with the dense kernel's skipna semantics. """ ds = expr.data shape = tuple(len(indexes[d]) for d in grid_dims) - strides = np.array( - [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] - ) - - transposed: list[str] = [] - axis_positions: list[np.ndarray] = [] - for i, (d, stride) in enumerate(zip(grid_dims, strides)): - if d == scatter_dim: - transposed.append(member_dim) - axis_positions.append(codes * stride) - else: - transposed.append(d) - axis_positions.append(np.arange(shape[i]) * stride) + strides = [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] + transposed = [member_dim if d == scatter_dim else d for d in grid_dims] + axis_positions = [ + codes * stride if d == scatter_dim else np.arange(n) * stride + for d, n, stride in zip(grid_dims, shape, strides) + ] cell_rows = axis_positions[0] for pos in axis_positions[1:]: cell_rows = cell_rows[..., None] + pos @@ -220,34 +194,32 @@ def _assemble( coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) - nterm = ds.sizes[TERM_DIM] - rows = np.repeat(cell_rows, nterm) - + rows = np.repeat(cell_rows, ds.sizes[TERM_DIM]) keep = (vars_ != -1) & (coeffs != 0) & ~np.isnan(coeffs) + full_size = int(np.prod(shape, dtype=np.int64)) if shape else 1 - width = expr.model._xCounter coo = scipy.sparse.coo_array( - (coeffs[keep], (rows[keep], vars_[keep])), shape=(full_size, width) + (coeffs[keep], (rows[keep], vars_[keep])), + shape=(full_size, expr.model._xCounter), ) - csr = scipy.sparse.csr_array(coo) # sums duplicates == the group sum const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) const = np.zeros(full_size) np.add.at(const, cell_rows, np.where(np.isnan(const_vals), 0.0, const_vals)) - return CSRPayload(csr, const, grid_dims, indexes, expr.model) + return CSRPayload( + scipy.sparse.csr_array(coo), const, grid_dims, indexes, expr.model + ) def try_csr_merge( exprs: Any, dim: str, join: Any, kwargs: dict ) -> LinearExpression | None: """ - Sparse branch of :func:`linopy.expressions.merge`. - - Returns a CSR-backed combined expression when every input is a plain - LinearExpression on the same result grid (CSR-backed, or dense and - convertible) — where any join produces the identical result — and None - otherwise to fall through to the dense path. + Sparse branch of :func:`linopy.expressions.merge`: combine plain + LinearExpressions on one shared grid (CSR-backed or dense-convertible), + where any join produces the identical result. Returns None to fall + through to the dense path. """ from linopy.expressions import LinearExpression @@ -264,11 +236,9 @@ def try_csr_merge( combined: CSRPayload | None = None for e in exprs: - payload = e._csr + payload = e._csr or CSRPayload.from_expression(e, template) if payload is None: - payload = CSRPayload.from_expression(e, template) - if payload is None: - return None + return None combined = payload if combined is None else combined.add(payload) assert combined is not None return LinearExpression._from_csr(combined, exprs[0].model) @@ -288,19 +258,15 @@ def extract_pending( and sign is None and rhs is None ): - expr, sign, rhs = lhs._pending - if expr._csr is None: # already materialized elsewhere - return None - lhs = expr + lhs, sign, rhs = lhs._pending if not (isinstance(lhs, LinearExpression) and lhs._csr is not None): return None if not isinstance(sign, str) or rhs is None or not is_constant(rhs): return None - payload = lhs._csr rhs_da = _as_rhs_dataarray(rhs) - if rhs_da is None or not set(rhs_da.dims) <= set(payload.grid_dims): + if rhs_da is None or not set(rhs_da.dims) <= set(lhs._csr.grid_dims): return None - return payload, sign, rhs + return lhs._csr, sign, rhs def _as_rhs_dataarray(rhs: Any) -> DataArray | None: @@ -310,9 +276,36 @@ def _as_rhs_dataarray(rhs: Any) -> DataArray | None: da = as_dataarray(rhs) except (TypeError, ValueError): return None - if set(da.dims) & set(HELPER_DIMS): - return None - return da + return None if set(da.dims) & set(HELPER_DIMS) else da + + +def _rhs_grid_values(payload: CSRPayload, rhs: Any) -> np.ndarray: + """ + Align the rhs to the payload grid by label and flatten it, with v1 + parity: NaN in the rhs raises (§5) and a reordered or differing index + on a shared dim raises (§8), as on the dense path. + """ + from linopy.semantics import check_user_nan, is_v1 + + rhs_da = _as_rhs_dataarray(rhs) + assert rhs_da is not None + if is_v1(): + if bool(rhs_da.isnull().any()): + check_user_nan() + for d in rhs_da.dims: + if not rhs_da.get_index(d).equals(payload.indexes[str(d)]): + raise ValueError( + f"Coordinate mismatch on shared dimension {d!r} between " + "the rhs and the grouped result. Align the rhs with " + ".sel(...) / .reindex(...) before combining (§8)." + ) + rhs_da = rhs_da.reindex( + {d: payload.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan + ) + missing = {d: payload.indexes[d] for d in payload.grid_dims if d not in rhs_da.dims} + if missing: + rhs_da = rhs_da.expand_dims(missing) + return rhs_da.transpose(*payload.grid_dims).to_numpy().reshape(-1) def realize_csr_constraint( @@ -321,64 +314,36 @@ def realize_csr_constraint( """ Staple sign and rhs onto a CSR-backed lhs to form a CSRConstraint. - Equivalent to materializing the sparse sum, adding the constraint and - freezing it — but the padded dense ``_term`` rectangle never exists; - peak memory is proportional to the number of nonzero terms. + Label columns are mapped to dense variable positions, the payload's + constant moves to the rhs, labels are allocated as in + ``Model._allocate_constraint_labels``, and rows without terms or with a + NaN rhs are inactive — all as on the frozen dense path. """ from linopy.common import maybe_replace_sign from linopy.constraints import CSRConstraint - from linopy.semantics import check_user_nan, is_v1 sign = maybe_replace_sign(sign) full_size = payload.n_cells - # map label columns to dense positions in the active variable array label_index = model.variables.label_index - label_to_pos = label_index.label_to_pos coo = payload.csr.tocoo() csr = scipy.sparse.csr_array( scipy.sparse.coo_array( - (coo.data, (coo.coords[0], label_to_pos[coo.coords[1]])), + (coo.data, (coo.coords[0], label_index.label_to_pos[coo.coords[1]])), shape=(full_size, label_index.n_active_vars), ) ) csr.eliminate_zeros() - # rhs aligned by label to the grid; expression constants move to the rhs - rhs_da = _as_rhs_dataarray(rhs) - assert rhs_da is not None - if is_v1() and bool(rhs_da.isnull().any()): - check_user_nan() # §5: NaN in a user constant raises under v1 - if is_v1(): - # §8 parity with the dense path: a reordered or differing rhs index - # raises; the user aligns explicitly with .sel/.reindex - for d in rhs_da.dims: - if not rhs_da.get_index(d).equals(payload.indexes[str(d)]): - raise ValueError( - f"Coordinate mismatch on shared dimension {d!r} between " - "the rhs and the grouped result. Align the rhs with " - ".sel(...) / .reindex(...) before combining (§8)." - ) - rhs_da = rhs_da.reindex( - {d: payload.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan - ) - missing = {d: payload.indexes[d] for d in payload.grid_dims if d not in rhs_da.dims} - if missing: - rhs_da = rhs_da.expand_dims(missing) - rhs_flat = ( - rhs_da.transpose(*payload.grid_dims).to_numpy().reshape(-1) - payload.const - ) + rhs_flat = _rhs_grid_values(payload, rhs) - payload.const - # label allocation, mirroring Model._allocate_constraint_labels; rows - # without terms or with NaN rhs are inactive, as in the frozen dense path cindex = model._cCounter model._cCounter += full_size active = (np.diff(csr.indptr) > 0) & ~np.isnan(rhs_flat) - con_labels = np.arange(cindex, cindex + full_size)[active] return CSRConstraint( csr[active], - con_labels, + np.arange(cindex, cindex + full_size)[active], rhs_flat[active], sign, coords=[payload.indexes[d] for d in payload.grid_dims], diff --git a/linopy/expressions.py b/linopy/expressions.py index 4725dcf6..48408732 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -496,17 +496,12 @@ def sum( Defaults to False. sparse : bool, optional Build the grouped sum in CSR form behind the ordinary - LinearExpression type (akin to dask-backed xarray objects): one - row per result cell, one column per variable, duplicate - variables summed and terms label-ordered — the group-size - padding of the dense kernel has no analog. Operations without a - sparse branch transparently expand to the dense rectangle in - this canonical form; a still-sparse left-hand side reaching - ``Model.add_constraints`` with ``freeze=True`` becomes a - CSRConstraint directly. Requires v1 semantics. Defaults to - ``linopy.options["sparse_groupby"]`` (only honored under v1). - Only single-key groupers over an existing dimension are held - sparse. + LinearExpression type — no group-size padding; a still-sparse + lhs reaching ``Model.add_constraints`` with ``freeze=True`` + becomes a CSRConstraint directly, other operations expand to + the dense rectangle in canonical term layout. Single-key + groupers only; requires v1 semantics. Defaults to + ``linopy.options["sparse_groupby"]``. See :mod:`linopy.csr`. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -1498,7 +1493,6 @@ def type(self) -> str: @property def data(self) -> Dataset: if self._data is None and self._csr is not None: - # expand to the dense rectangle (canonical form); sparsity is spent self._data = self._csr.materialize().data self._csr = None return self._data @@ -1524,7 +1518,6 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: if self._data is None and self._csr is not None: - # answerable from the sparse payload without materializing return tuple(self._csr.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @@ -1715,8 +1708,6 @@ def to_constraint( to the right-hand side. """ if self._csr is not None and isinstance(sign, str) and is_constant(rhs): - # carry the sparse payload on the constraint; its .data property - # materializes through this method's dense path if needed return constraints.Constraint._from_pending(self, sign, rhs, self.model) rhs = as_constant(rhs) diff --git a/linopy/model.py b/linopy/model.py index 62fbcfbe..ac39a5f9 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1139,8 +1139,6 @@ def add_constraints( name = self._resolve_constraint_name(name) - # a still-sparse grouped lhs with freeze on becomes a CSRConstraint - # directly (sign and rhs stapled on), skipping the padded dense form resolved_freeze = self.freeze_constraints if freeze is None else freeze if resolved_freeze and mask is None and not self.chunk: from linopy.csr import extract_pending, realize_csr_constraint diff --git a/test/test_sparse_groupby.py b/test/test_sparse_groupby.py index 12e06ec0..beddef1d 100644 --- a/test/test_sparse_groupby.py +++ b/test/test_sparse_groupby.py @@ -26,7 +26,7 @@ def require_v1() -> None: def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): - """Model with gen_p (gen, snapshot) and flow (line, snapshot) on a ring.""" + """Model with gen_p and flow on a ring; load ordered like the sorted groups (v1).""" rng = np.random.default_rng(seed) n_bus = len(gens_per_bus) buses = pd.Index([f"bus{i}" for i in range(n_bus)], name="bus") @@ -42,7 +42,6 @@ def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): gbus = pd.Series(buses[gen_bus], index=gens, name="bus") bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") bus1 = pd.Series(buses[(np.arange(n_bus) + 1) % n_bus], index=lines, name="bus") - # v1 requires explicit alignment: order the load like the (sorted) groups load = xr.DataArray( rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" ).sortby("bus") @@ -74,7 +73,6 @@ def test_csr_requires_v1(): return with pytest.raises(ValueError, match="requires v1 semantics"): (eff * gen_p).groupby(gbus).sum(sparse=True) - # the option is not honored under legacy: result is eager and dense linopy.options["sparse_groupby"] = True try: res = (eff * gen_p).groupby(gbus).sum() @@ -89,7 +87,6 @@ def test_csr_is_plain_linear_expression_and_materializes_equivalently(): sparse = (eff * gen_p).groupby(gbus).sum(sparse=True) eager = (eff * gen_p).groupby(gbus).sum() assert type(sparse) is LinearExpression - # comparing data materializes; the result must be exactly today's assert_linequal(sparse, eager) @@ -108,9 +105,6 @@ def test_scalar_ops_stay_csr(): sparse = -2.0 * (eff * gen_p).groupby(gbus).sum(sparse=True) assert sparse._csr is not None # still unmaterialized after neg/scale eager = -2.0 * (eff * gen_p).groupby(gbus).sum() - # scaling after grouping normalizes padded-slot fills while scaling - # before grouping keeps fresh pads; masked-slot fill is not contractual, - # so compare with masked slots blanked on both sides from linopy.testing import _sort_by_vars_along_term a, b = _sort_by_vars_along_term(sparse), _sort_by_vars_along_term(eager) @@ -144,6 +138,7 @@ def test_freeze_realizes_csr_without_dense_rectangle(): def test_freeze_false_falls_back_to_identical_dense_constraint(): + """The fallback is canonical-form, so compare mathematically (strict=False).""" require_v1() m1, *rest = base_model() c1 = m1.add_constraints( @@ -155,8 +150,6 @@ def test_freeze_false_falls_back_to_identical_dense_constraint(): balance_lhs(*rest2[:6], sparse=True) == rest2[6], name="balance" ) assert isinstance(c2, Constraint) - # the sparse lhs materializes in canonical form (terms label-sorted, - # padding trailing), so compare mathematically, not layout-strictly assert_conequal(c1, c2, strict=False) assert np.array_equal(c1.labels.values, c2.labels.values)