feat(v1): CSR-backed sparse groupby-sum - skew-independent build memory (6-9x on the #745 hub case)

[FBumann]

This is a drive by PR done with CLAUDE FABLE. Its here to explore how we can better ahdnle sparsity.

Not sure where this goes yet.

Note

The following content was generated by AI (Claude Code), prompted and reviewed by @FBumann.

pytest-benchmem, the #745 hub scenario (nodal balance, 120 buses × 24 snapshots, one hub bus with hub generators, 100 on each other bus; per-benchmark allocator peak of the build):

hub generators	dense groupby build peak	sparse=True + freeze=True	median build time
8,000	277.1 MiB	43.5 MiB (6.4×)	19.4 ms → 12.9 ms
16,000	546.1 MiB	61.0 MiB (9.0×)	38.5 ms → 16.5 ms

Doubling the skew doubles the dense peak but only grows the sparse one by its actual new terms: the sparse path is O(nnz) and independent of the group-size distribution, so the padding cost that PyPSA's meshed-bus banding exists to contain (see #745) is gone rather than capped. Constraint export also gets cheaper, because the result already is the export representation.

What this adds

Addresses the inherent half of #745 (and #757) on the v1 line, following the umbrella plan in #756: groupby(g).sum() must pad every group to the largest group's term count because a LinearExpression stores terms in a dense cells × _term rectangle. The padded cells are pure intermediate waste for constraint-bound results — the LP/matrix export drops them again.

expr.groupby(g).sum(sparse=True) (or linopy.options["sparse_groupby"] = True) builds the grouped sum in CSR form behind the unchanged LinearExpression type — modeled on dask-backed xarray objects: same public class, different backing, no new public primitive. The payload (linopy/csr.py, CSRPayload) stores the expression as A @ x + c: one CSR row per result cell, one column per variable label, plus a dense per-cell constant. The _term axis is ragged by construction, so group-size padding has no analog, and the operations between a groupby and its constraint 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 / scalar multiplication scale values,
• == rhs is carried as a pending payload on the Constraint,
• Model.add_constraints(..., freeze=True) (or Model.freeze_constraints) staples sign and rhs on and registers a CSRConstraint directly — the existing frozen backend (perf: matrix accessor rewrite #630), so the LP writer and matrix export work unchanged.

Anything without a sparse branch transparently expands to the dense rectangle through the .data property and proceeds exactly as today, so compatibility is a fallback, not a constraint.

The contract, and why it is v1-gated

The CSR form is canonical: duplicate variables within a cell are summed (2x + 3x → 5x) and terms are ordered by variable label. Expanding back to the dense form therefore yields the mathematically identical expression in canonical term layout, not the eager kernel's exact positional layout. Term layout is non-contractual under v1 (linopy already treats it as such at export: zero-coeff filtering, maybe_group_terms_polars, densify_terms), which is why the feature requires options["semantics"] = "v1" — under legacy, sparse=True raises and the option is ignored. Equivalence is asserted at the level that matters: identical polars constraint rows and identical LP files (the test suite includes an end-to-end LP diff, term order within a row canonicalized).

v1 parity is kept in the direct realization: NaN in the rhs raises (§5), a reordered/differing rhs index raises with the standard alignment message (§8), and rows whose constraint is absent realize as masked (§12). Labels and the constraint grid bit-match the dense path.

Reproduce

The suite gains a sparse sibling of the existing nodal_balance pattern (the #745 severity sweep), so the padding cost is CI-visible on CodSpeed and one command locally:

pytest benchmarks/ -k "nodal_balance and build" --benchmark-memory

build peak (KiB)	severity 0	50	100
nodal_balance (dense)	951	5,460	9,909
nodal_balance_sparse	1,524	1,524	1,524

Standalone repro (single pytest file — the headline table above is its output)

# repro745.py — the #745 hub scenario, dense vs sparse (CSR) groupby-sum.
# run:  pytest repro745.py --benchmark-memory
import numpy as np
import pandas as pd
import pytest
import xarray as xr

import linopy


def build_balance(sparse: bool, hub: int) -> linopy.Model:
    linopy.options["semantics"] = "v1"
    n_bus, n_snap = 120, 24
    buses = pd.RangeIndex(n_bus, name="bus")
    gen_of_bus = np.repeat(np.arange(n_bus), [hub] + [100] * (n_bus - 1))
    gens = pd.RangeIndex(len(gen_of_bus), name="gen")
    snaps = pd.RangeIndex(n_snap, name="snapshot")

    m = linopy.Model()
    gen_p = m.add_variables(coords=[gens, snaps], name="gen_p")
    load = xr.DataArray(np.ones((n_bus, n_snap)), coords=[buses, snaps])

    grouper = pd.Series(gen_of_bus, index=gens, name="bus")
    supply = (1 * gen_p).groupby(grouper).sum(sparse=sparse)
    m.add_constraints(supply == load, name="balance", freeze=sparse)
    return m


@pytest.mark.parametrize("hub", [8000, 16000])
@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"])
def test_hub_balance_build(benchmark, sparse, hub):
    benchmark(build_balance, sparse, hub)

Needs pytest-benchmark and pytest-benchmem (both in the benchmarks extra). Measured on this branch, macOS arm64, python 3.13. Both paths produce identical constraints — the test suite asserts identical polars term rows and identical LP files for this construction.

Scope and current limitations

• Single-key groupers over an existing dimension; multi-key / multidim groupers, use_fallback and observed take the eager path.
• The sparse branches cover the constraint-building chain (neg, scalar mul, merge/+/- on the same grid, comparison with a constant rhs); everything else materializes canonically.
• freeze=False falls back to a dense Constraint (mathematically equal, canonical layout).
• Quadratic expressions are untouched.

Natural follow-ups on this seam (not in this PR): dot as a payload op (the v1-side counterpart of #867, which stays the master-era #748 fix) and ragged/KVL-style merge (#749).

History

Two commits kept deliberately: the first implements the same public behavior with a deferred-recipe payload (ungrouped parts + groupers, bit-identical fallback via replaying the eager kernel); the second swaps the payload to CSR. The recipe variant measured within ~5 % of the CSR numbers but accumulates unbounded part lists, holds all input expressions alive until realization, and every future op would need its own replay logic — CSR is the representation the rest of #756 (dot, merge) composes on. The swap trades the recipe's bit-identical fallback for the canonical-form contract above.

Full suite: 6336 passed, 557 skipped — unchanged from the base branch. #717 (feat/arithmetic-convention) has landed on master, so this branch is now rebased directly onto master and is no longer stacked.

🤖 Generated with Claude Code

[codspeed-hq]

Merging this PR will degrade performance by 9.1%

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
❌ 2 regressed benchmarks
✅ 172 untouched benchmarks
🆕 6 new benchmarks
⏩ 175 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
❌	test_to_lp[qp-n=1000]	2 MB	2.6 MB	-23.35%
❌	test_to_lp[nodal_balance-severity=0]	3.3 MB	3.7 MB	-11.45%
⚡	test_to_lp[expression_arithmetic-n=250]	45.3 MB	40.9 MB	+10.66%
🆕	test_build[nodal_balance_sparse-severity=0]	N/A	1.5 MB	N/A
🆕	test_build[nodal_balance_sparse-severity=100]	N/A	1.5 MB	N/A
🆕	test_build[nodal_balance_sparse-severity=50]	N/A	1.5 MB	N/A
🆕	test_to_lp[nodal_balance_sparse-severity=0]	N/A	3.9 MB	N/A
🆕	test_to_lp[nodal_balance_sparse-severity=100]	N/A	3.7 MB	N/A
🆕	test_to_lp[nodal_balance_sparse-severity=50]	N/A	3.7 MB	N/A

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing feat/v1-sparse-groupby (13f6681) with master (6b133d5)}

Footnotes

1. 175 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[github-actions]

Build cost — v1 vs legacy

v1 build peak & time relative to legacy, on this commit — not a comparison against master (that is CodSpeed).

peak — v1 / legacy	time — v1 / legacy

Full table (time + peak, mean)

benchmarks/drivers/test_build.py::test_build[basic-n=10]
                  time (s)         peak (KiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.07746 (1.08)   │   15.03 (1.00) 
 (v1)        0.07148 (1.0)   │    15.00 (1.0) 

benchmarks/drivers/test_build.py::test_build[basic-n=250]
                  time (s)         peak (MiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.08456 (1.09)   │   12.04 (1.00) 
 (v1)        0.07727 (1.0)   │    12.04 (1.0) 

benchmarks/drivers/test_build.py::test_build[cumsum-severity=0]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.03362 (1.08)   │   15.20 (1.0) 
 (v1)         0.0311 (1.0)   │   15.20 (1.0) 

benchmarks/drivers/test_build.py::test_build[cumsum-severity=100]
                  time (s)        peak (MiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.04906 (1.09)   │   44.93 (1.0) 
 (v1)        0.04486 (1.0)   │   44.93 (1.0) 

benchmarks/drivers/test_build.py::test_build[cumsum-severity=50]
                  time (s)        peak (MiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.03657 (1.06)   │   11.51 (1.0) 
 (v1)        0.03455 (1.0)   │   11.51 (1.0) 

benchmarks/drivers/test_build.py::test_build[expression_arithmetic-n=10]
                  time (s)         peak (KiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.08662 (1.07)   │   24.34 (1.06) 
 (v1)         0.0812 (1.0)   │    23.05 (1.0) 

benchmarks/drivers/test_build.py::test_build[expression_arithmetic-n=250]
                  time (s)         peak (MiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.09494 (1.06)   │   16.12 (1.00) 
 (v1)        0.08986 (1.0)   │    16.12 (1.0) 

benchmarks/drivers/test_build.py::test_build[knapsack-n=10000]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.02067 (1.06)   │   752.18 (1.10) 
 (v1)        0.01944 (1.0)   │    685.15 (1.0) 

benchmarks/drivers/test_build.py::test_build[knapsack-n=100]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.02025 (1.08)   │   3.12 (1.33) 
 (v1)        0.01881 (1.0)   │    2.34 (1.0) 

benchmarks/drivers/test_build.py::test_build[kvl_cycles-severity=0]
                  time (s)          peak (MiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.05748 (1.22)   │   126.16 (1.44) 
 (v1)        0.04698 (1.0)   │     87.71 (1.0) 

benchmarks/drivers/test_build.py::test_build[kvl_cycles-severity=100]
                  time (s)          peak (MiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.05636 (1.23)   │   126.16 (1.44) 
 (v1)        0.04601 (1.0)   │     87.71 (1.0) 

benchmarks/drivers/test_build.py::test_build[kvl_cycles-severity=50]
                  time (s)          peak (MiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.05705 (1.23)   │   126.16 (1.44) 
 (v1)        0.04636 (1.0)   │     87.71 (1.0) 

benchmarks/drivers/test_build.py::test_build[masked-n=100]
                 time (s)          peak (KiB) 
 name                mean   │            mean 
──────────────────────────────────────────────
 (legacy)    0.047 (1.01)   │    715.14 (1.0) 
 (v1)       0.04647 (1.0)   │   787.74 (1.10) 

benchmarks/drivers/test_build.py::test_build[masked-n=10]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.04622 (1.11)   │   4.55 (1.27) 
 (v1)        0.04181 (1.0)   │    3.59 (1.0) 

benchmarks/drivers/test_build.py::test_build[merge_balance-severity=0]
                 time (s)          peak (KiB) 
 name                mean   │            mean 
──────────────────────────────────────────────
 (legacy)   0.3242 (1.04)   │   704.12 (1.09) 
 (v1)        0.3123 (1.0)   │    647.38 (1.0) 

benchmarks/drivers/test_build.py::test_build[merge_balance-severity=100]
                 time (s)        peak (MiB) 
 name                mean   │          mean 
────────────────────────────────────────────
 (legacy)   0.3408 (1.05)   │   18.34 (1.0) 
 (v1)        0.3255 (1.0)   │   18.34 (1.0) 

benchmarks/drivers/test_build.py::test_build[merge_balance-severity=50]
                 time (s)       peak (MiB) 
 name                mean   │         mean 
───────────────────────────────────────────
 (legacy)   0.3376 (1.04)   │   9.54 (1.0) 
 (v1)        0.3249 (1.0)   │   9.54 (1.0) 

benchmarks/drivers/test_build.py::test_build[milp-n=10]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.06496 (1.08)   │   3.77 (1.12) 
 (v1)        0.06023 (1.0)   │    3.37 (1.0) 

benchmarks/drivers/test_build.py::test_build[milp-n=50]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.06562 (1.09)   │   216.59 (1.10) 
 (v1)        0.06016 (1.0)   │    196.23 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance-severity=0]
                  time (s)         peak (KiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.03339 (1.09)   │   938.49 (1.0) 
 (v1)        0.03072 (1.0)   │   938.49 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance-severity=100]
                  time (s)       peak (MiB) 
 name                 mean   │         mean 
────────────────────────────────────────────
 (legacy)   0.03463 (1.09)   │   9.66 (1.0) 
 (v1)         0.0319 (1.0)   │   9.66 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance-severity=50]
                  time (s)       peak (MiB) 
 name                 mean   │         mean 
────────────────────────────────────────────
 (legacy)   0.03412 (1.09)   │   5.32 (1.0) 
 (v1)        0.03126 (1.0)   │   5.32 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance_sparse-severity=0]
                  time (s)       peak (MiB) 
 name                 mean   │         mean 
────────────────────────────────────────────
 (legacy)   0.01841 (1.02)   │   1.41 (1.0) 
 (v1)        0.01806 (1.0)   │   1.41 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance_sparse-severity=100]
                  time (s)       peak (MiB) 
 name                 mean   │         mean 
────────────────────────────────────────────
 (legacy)   0.01835 (1.01)   │   1.41 (1.0) 
 (v1)        0.01818 (1.0)   │   1.41 (1.0) 

benchmarks/drivers/test_build.py::test_build[nodal_balance_sparse-severity=50]
                  time (s)       peak (MiB) 
 name                 mean   │         mean 
────────────────────────────────────────────
 (legacy)   0.01829 (1.01)   │   1.41 (1.0) 
 (v1)        0.01812 (1.0)   │   1.41 (1.0) 

benchmarks/drivers/test_build.py::test_build[piecewise-n=1000]
                 time (s)          peak (KiB) 
 name                mean   │            mean 
──────────────────────────────────────────────
 (legacy)   0.1626 (1.04)   │   946.85 (1.06) 
 (v1)         0.157 (1.0)   │    891.54 (1.0) 

benchmarks/drivers/test_build.py::test_build[piecewise-n=10]
                 time (s)        peak (KiB) 
 name                mean   │          mean 
────────────────────────────────────────────
 (legacy)   0.1594 (1.03)   │   12.01 (1.0) 
 (v1)        0.1545 (1.0)   │   12.01 (1.0) 

benchmarks/drivers/test_build.py::test_build[qp-n=1000]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.04145 (1.06)   │   147.70 (1.06) 
 (v1)        0.03902 (1.0)   │    139.87 (1.0) 

benchmarks/drivers/test_build.py::test_build[qp-n=10]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.04123 (1.07)   │   2.60 (1.09) 
 (v1)        0.03856 (1.0)   │    2.38 (1.0) 

benchmarks/drivers/test_build.py::test_build[rolling-severity=0]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.03418 (1.07)   │   696.75 (1.03) 
 (v1)        0.03188 (1.0)   │    673.70 (1.0) 

benchmarks/drivers/test_build.py::test_build[rolling-severity=100]
                  time (s)         peak (MiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.07293 (1.00)   │   137.97 (1.0) 
 (v1)        0.07263 (1.0)   │   137.97 (1.0) 

benchmarks/drivers/test_build.py::test_build[rolling-severity=50]
                  time (s)        peak (MiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.05332 (1.08)   │   69.22 (1.0) 
 (v1)        0.04936 (1.0)   │   69.22 (1.0) 

benchmarks/drivers/test_build.py::test_build[sos-n=1000]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)   0.03964 (1.09)   │   402.33 (1.00) 
 (v1)        0.03621 (1.0)   │    402.30 (1.0) 

benchmarks/drivers/test_build.py::test_build[sos-n=10]
                  time (s)        peak (KiB) 
 name                 mean   │          mean 
─────────────────────────────────────────────
 (legacy)   0.03883 (1.09)   │   3.19 (1.19) 
 (v1)        0.03573 (1.0)   │    2.69 (1.0) 

benchmarks/drivers/test_build.py::test_build[sparse_network-n=10]
                  time (s)         peak (KiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.04294 (1.06)   │   29.00 (1.54) 
 (v1)        0.04041 (1.0)   │    18.84 (1.0) 

benchmarks/drivers/test_build.py::test_build[sparse_network-n=250]
                  time (s)         peak (MiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)   0.05312 (1.12)   │   37.95 (1.43) 
 (v1)        0.04727 (1.0)   │    26.51 (1.0) 

benchmarks/drivers/test_build.py::test_build[storage-n=10]
                  time (s)          peak (KiB) 
 name                 mean   │            mean 
───────────────────────────────────────────────
 (legacy)    0.08254 (1.0)   │    410.93 (1.0) 
 (v1)       0.08417 (1.02)   │   427.84 (1.04) 

benchmarks/drivers/test_build.py::test_build[storage-n=250]
                  time (s)         peak (MiB) 
 name                 mean   │           mean 
──────────────────────────────────────────────
 (legacy)    0.08814 (1.0)   │     9.94 (1.0) 
 (v1)       0.09119 (1.03)   │   10.22 (1.03) 

📊 Interactive plots + CSV: download the semantics-report-v1-vs-legacy artifact from this run.

_{Report-only · not a gate · refreshed on every push · obsolete once legacy is dropped.}

[FabianHofmann]

Note

The following content was generated by AI (Claude Code), reviewed by @FabianHofmann.

Follow-up once #926 (label columns for CSRConstraint) lands

Parts of this PR are a stopgap for the fact that CSRConstraint stores dense variable positions. With label columns, the following is deleted or simplified, nothing else moves:

• Delete Constraint._from_pending, the lazy branch in Constraint.data, and extract_pending (linopy/constraints.py). LinearExpression.to_constraint on a payload-backed expression returns an unassigned CSRConstraint (cindex=None) directly, so no pending (lhs, sign, rhs) state is needed.
• Simplify CSRConstraint.from_payload: drop the label-to-position mapping and the label-index dependency. It becomes a thin constructor from payload, sign and rhs; the has_terms / eliminate_zeros ordering stays.
• Reduce the dispatch in Model.add_constraints to "if the incoming constraint is an unassigned CSRConstraint, assign labels", which is the step the dense path already performs.

Unchanged: linopy/sparse_expression.py (CSRPayload, try_csr_merge), the _payload hooks on LinearExpression, from_payload as the sparse sibling of from_mutable, and _rhs_grid_values.

Note that frozen constraints produced here inherit the bug in #926: adding or removing variables after a sparse freeze=True constraint fails at matrix assembly, exactly as for the pre-existing dense freeze=True path.

[FabianHofmann]

highly needed for the math-spec integration in linopy (for pypsa networks with loads of connections). gave it a push. merging now