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. Stacked on #717 (feat/arithmetic-convention) — draft until that lands.

🤖 Generated with Claude Code