perf(groupby): unify scatter kernel over numpy and dask via apply_ufunc - #802
Conversation
The fast path of LinearExpression.groupby(...).sum() used ds.unstack(group_dim, fill_value=...) followed by a stack, which materializes 2-3 intermediate copies of the padded result (n_groups x max_group_size x nterm) and goes through pandas MultiIndex machinery sized by the number of elements. Instead, factorize the groups and scatter coeffs/vars directly into the preallocated padded result arrays; constants are group-summed with np.add.at. Peak memory drops to input + result (the minimum for the padded layout) and the grouping itself gets considerably faster. The result is unchanged: same dims, coords, term ordering and padding. The unstack-based implementation is kept as _sum_by_unstack and still used for chunked (dask-backed) data, which cannot be scattered into numpy arrays. NaN group labels now raise an informative ValueError instead of failing inside unstack. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a test for grouping over an empty group dimension, which the scatter fast path handles cleanly but the unstack fallback cannot. Trim comments that duplicated the helper docstrings.
Relax the groupby-sum scatter gate to a pure numpy/dask check: auxiliary coordinates on the grouped dimension no longer force the slow unstack path. Summing over groups collapses that dimension, so both kernels drop every coordinate tied to it — the scatter result is identical, just cheaper. The unstack kernel now serves only chunked (dask) data, and a debug log records when that fallback is taken. Inline the now-trivial predicate into the dispatch and consolidate the kernel tests into a TestGroupbySumScatterKernel class: a one-line case table over a shared fixture, with added coverage for combined structures, auxiliary coords, and a MultiIndex grouped dimension. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Merging this PR will improve performance by ×2.1
Performance Changes
Tip Curious why this is faster? Comment Comparing 
Footnotes
|
06ab9af to
2abf407
Compare
Replace the previous numpy-scatter / dask-unstack split with a single kernel (`_grouped_sum`) wrapped in `xarray.apply_ufunc`. It scatters terms into the padded result arrays for numpy-backed data and runs the same scatter lazily on chunked (dask) data via `dask="parallelized"`, after gathering the grouped and term dimensions (the scatter's core dims) into single chunks. This removes the last `pd.MultiIndex`/`unstack` usage in groupby-sum, drops the numpy-vs-dask branch in `sum()`, and keeps peak memory at input + result on both backends. Multi-key / DataFrame grouping and its `MultiIndex` result are unaffected — that logic sits above the kernel. Tests verify the kernel from first principles (each group's terms and constant must match its members) across every case shape on both numpy and dask, plus explicit anchors pinning the exact padded layout — member order, fill position, term interleaving and the factor axis — for the linear, multidim and quadratic cases. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2abf407 to
ac8ec47
Compare
|
I ran the benchmark locally with pytets-benchmem to also check timing. Small improvement there too! Note The following content was generated by AI. Local verification of the performance claim on the
Both peak memory (~2.4–2.5× lower) and build time (~5–13% faster) improve, and the gain grows with group skew — the pathological case the scatter targets. Consistent with the CodSpeed report (×2–3 memory on build/to_lp/to_solver). MethodEach version's A |
|
I'll put in on my stack, but no promises this week. I also think this is not time critical. The actual use of dask, is non-existent in my expectation. |
coroa
left a comment
coroa
left a comment
There was a problem hiding this comment.
Cool, the core stuff is very good. I find the tests a bit bloated and have a few small comments.
Great work
| @pytest.mark.parametrize("backend", ["numpy", "dask"]) | ||
| @pytest.mark.parametrize( | ||
| "build", | ||
| GROUPBY_SUM_CASES.values(), | ||
| ids=GROUPBY_SUM_CASES.keys(), | ||
| ) | ||
| def test_grouped_sum_correct( | ||
| self, | ||
| build: Callable[[SimpleNamespace], tuple[LinearExpression, pd.Series]], | ||
| groupby_ctx: SimpleNamespace, | ||
| backend: str, | ||
| ) -> None: | ||
| """ | ||
| Each group's terms and constant must match its members, from first | ||
| principles, on both numpy and dask backends. See ``GROUPBY_SUM_CASES`` | ||
| for the structures covered. | ||
| """ | ||
| if backend == "dask": | ||
| pytest.importorskip("dask") | ||
| expr, groups = build(groupby_ctx) | ||
| _assert_grouped_sum_correct(expr, groups, chunked=backend == "dask") |
There was a problem hiding this comment.
Wouldn't the simpler pattern to compare fallback to scatter implementation, like the skewed one above, be the better comparison. Didn't check in detail what assert_grouped_sum_correct actually does, but it looks too complicated.
There was a problem hiding this comment.
The fallback errors on quadratics (no _factor handling) and diverges on masked and auxiliary-coordinate cases, so it cannot validate the scatter kernel in general.
But we should split the cases up i think!
There was a problem hiding this comment.
Wouldn't it be easier then to do a simpler test example for quadratics, and fix the divergences by re-ordering or whatever it would take?
I don't understand the test and where it would fail and find this concerning. I agree with large test coverage is good. With claude i often can't tell whether the tests actually test what we would want them to
… path, and the ones that fail there and need manual validation
…thods, and improving variable names
|
@coroa Im done here. About the bloated tests. I agree, but thought its better to over cover for now. If you agree, im open to removing some tests.
|
|
@coroa I would merge this tomorrow morning if you have no objections to my latest changes. |
| return self.map(func, **kwargs, shortcut=True) | ||
| return self.map(func) |
There was a problem hiding this comment.
@FabianHofmann Can you clarify what kwargs were meant to be passed here? Shortcuts was dead code, so fine to remove.
There was a problem hiding this comment.
yes, it was a no-op that was mainly there for mirroring the signature of the internal xarray groupby.map function. it was never explicitly used as I (guess I) wanted to tackle this later which never happened. so fine to remove. the kwargs seems to be dead anyway and setting it would have let to errors. so this is fixing it
| @pytest.mark.parametrize("backend", ["numpy", "dask"]) | ||
| @pytest.mark.parametrize( | ||
| "build", | ||
| GROUPBY_SUM_CASES.values(), | ||
| ids=GROUPBY_SUM_CASES.keys(), | ||
| ) | ||
| def test_grouped_sum_correct( | ||
| self, | ||
| build: Callable[[SimpleNamespace], tuple[LinearExpression, pd.Series]], | ||
| groupby_ctx: SimpleNamespace, | ||
| backend: str, | ||
| ) -> None: | ||
| """ | ||
| Each group's terms and constant must match its members, from first | ||
| principles, on both numpy and dask backends. See ``GROUPBY_SUM_CASES`` | ||
| for the structures covered. | ||
| """ | ||
| if backend == "dask": | ||
| pytest.importorskip("dask") | ||
| expr, groups = build(groupby_ctx) | ||
| _assert_grouped_sum_correct(expr, groups, chunked=backend == "dask") |
There was a problem hiding this comment.
Wouldn't it be easier then to do a simpler test example for quadratics, and fix the divergences by re-ordering or whatever it would take?
I don't understand the test and where it would fail and find this concerning. I agree with large test coverage is good. With claude i often can't tell whether the tests actually test what we would want them to
Drop the first-principles multiset oracle (_term_multisets, _assert_grouped_sum_correct) and its all-cases sweep: it was hard to review and gave no way to see, by hand, what it actually checked. Cover the cases the old fallback engine cannot reproduce (masked multi-dim vars, coords on the grouped dim, mixed quadratic + linear + const) with concrete, hand-pinned layout tests instead, run on both the numpy and the lazy dask kernel via a small _grouped_sum_on_backend helper. test_grouped_sum_matches_fallback stays as the equivalence guard for the cases the old engine does reproduce. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Read results through .vars/.coeffs/.const/.coords/.sizes instead of the internal .data Dataset, and chunk with expr.chunk(...) rather than reconstructing LinearExpression(expr.data.chunk(...)). Lazy dask results realise when their values are read, so the manual .compute() round-trip is gone too. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
@coroa I changed the testing from the generic but hard to read oracle to a few hand rolled cases. I dropped the oracle. We could also keep the oracle if you value the test surface, but I think its fine like this. |
|
@FBumann I had a another view now. I think it's fine. let's pull this in |
|
Sorry for not keeping up with review speed. Yes, thanks, the hand-rolled cases are clearer to me. Both: But, I would appreciate if you waited for me once i am reviewing things until i give the ok. |
* test: pin groupby-sum to preserve grouped-dim position (xfail) groupby-sum moves the group dim to the trailing position instead of keeping the grouped dim's slot, diverging from xarray's native groupby-reduce and linopy <= 0.8.0 (regressed in #802). Existing tests use assert_linequal, which aligns dims and misses the reorder. Strict xfail: remove the marker once grouping preserves dim position. * fix(groupby): preserve grouped-dim position in groupby-sum groupby-sum moved the group dim to the trailing position; both the apply_ufunc scatter kernel and the fallback append it there, and the LinearExpression constructor's broadcast then locks that order. Restore the grouped dim's original slot via _restore_group_dim_position: transpose the terms and reorder the dimension coordinates so the constructor's canonical dim order matches. MultiIndex group coords are left untouched (multikey/observed results stay order-agnostic). Drops the strict xfail; adds a release note. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor(groupby): fold insertion flag into new_dims sentinel * perf(groupby): make reordered terms C-contiguous to keep export cheap The dim-position transpose left coeffs/vars as non-contiguous views, so LP/solver export had to materialise a contiguous copy on ravel (CodSpeed flagged ~4x export memory on the nodal_balance padding benchmark). Force numpy-backed terms C-contiguous after the reorder so downstream flattening ravels a view. ascontiguousarray is a no-op when no reorder moved data; dask arrays are left lazy (chunks preserved). * perf(groupby): scatter into target axis order in one contiguous alloc Make _grouped_sum emit the grouped dimension in place directly: pass all surviving dims as core dims so apply_ufunc's output order is fully chosen, and build the padded array with the group axis in its target slot. The result is contiguous in a single allocation, so the downstream reorder and contiguity coercion become no-ops (no copy). Removes the build-memory regression from the previous force-contiguous approach: build and export peak now match master (~10MB / ~3MB) on nodal_balance instead of 19MB / 3MB. Trade-off: every dim is now a core dim, so dask inputs collapse to a single chunk (no per-surviving-dim parallelism). Kept lazy. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * perf(groupby): sum const before term scatters to restore peak-memory parity Inlining group_sum(data.const) into the Dataset literal made it evaluate after the coeffs/vars scatters, stacking its nan-masking temporaries on top of the retained padded arrays (+31% peak in expr_groupby_sum). Reducing const first frees them before the padded arrays exist, restoring the exact peak of master (verified with memray: 171 KB vs 275 KB). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(groupby): keep dim position when the grouper reuses the grouped dim's name Restoring the group dim's position inferred the consumed dims as those absent from the result, so a grouper named like the dim it groups (e.g. relabelling 'name' to a coarser 'name') shadowed the consumed dim, the restore early-returned and the dim drifted to the trailing position on both paths. Pass the grouper's own dims (DataArray/IndexVariable dims, Series/DataFrame index name) through _grouper_dims instead; unknown grouper types (grouper mappings, multi-key lists under use_fallback) keep the inference. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: FBumann <117816358+FBumann@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* docs: harmonize PR/issue refs in release notes Convert bare GitHub URLs and mixed RST link forms to a single anonymous linked short-ref form (`#NNN <url>`__) across the Upcoming and 0.8.0 sections, linkify plain #580 mentions, and fix a stray Markdown link in the 0.1.4 notes. * docs: add PR refs to remaining upcoming release-note entries Trace each previously unreferenced entry in the Upcoming Version section to its introducing PR and append the linked short-ref, so every bullet is traceable (#718, #566, #780, #783, #790, #801, #802, #824, #860).
> [!NOTE] > This PR was prepared by AI (Claude Code) at the maintainer's request. Manual equivalent of the monthly Dependabot bump for `benchmark/uv.lock`: `linopy` 0.8.0 → 0.9.0. #269 pinned `linopy==0.9.0` at the repo root, but the benchmark env is a separate project resolved from its own committed lock with `uv sync --frozen` — which never re-reads the root requirement. CodSpeed therefore ran **0.8.0 on both sides** of that PR's comparison, which is why it reported "will not alter performance". This bump is what actually puts 0.9.0 in front of the instruments. linopy 0.9.0 carries several changes we should feel — we use `groupby().sum()` in `src/fluxopt/constraints/sparse.py` and `src/fluxopt/model.py`: - PyPSA/linopy#802 — `perf(groupby)`: unified scatter kernel over numpy and dask via `apply_ufunc` - PyPSA/linopy#566 — int32 label arrays (~25% memory reduction) - PyPSA/linopy#815, PyPSA/linopy#816 — drop/prune explicit zeros from the constraint matrix Labelled `trigger:benchmark` so the walltime job runs alongside memory — this is a dependency step change on the dashboard baseline, so it is worth having both instruments record it on the same commit. <details> <summary>Local verification</summary> ``` $ cd benchmark && uv lock --upgrade-package linopy Resolved 59 packages in 574ms Updated linopy v0.8.0 -> v0.9.0 $ uv sync --frozen && uv run --no-sync python -c "import linopy; print(linopy.__version__)" 0.9.0 $ uv run --no-sync pytest . --benchmark-disable -q 41 passed, 185 warnings in 22.51s ``` </details> 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
3 participants
I use apply_ufunc to make this dask capable. As we dont have the reference unstack implementation anymore, i introduced quite a heavy testing part (fast though), as I find the apply_ufunc version harder to understand personally. Happy to strip it down.
Note
The following content was generated by AI.
Stacked on #793. Until #793 merges, this PR's diff includes its commit too — review only the top commit
perf(groupby): unify scatter kernel ....What this does
#793 split groupby-sum into a numpy kernel and a dask
unstackfallback. Thiscollapses them into a single kernel (
_grouped_sum) wrapped inxarray.apply_ufunc:dask="parallelized",after gathering the grouped dimension into a single chunk (which unstacking
required too).
This removes the last
pd.MultiIndex/unstackusage in groupby-sum, drops thenumpy-vs-dask branch in
sum(), and keeps peak memory at input + result on bothbackends. Multi-key / DataFrame grouping and its
MultiIndexresult areunaffected — that logic sits above the kernel (existing tests cover it).
Tests
The kernel is verified from first principles — for every group and every
slice over the non-grouped dims, the result's live terms must equal the multiset
of its members' terms and the constant their NaN-skipping sum — across every
case shape on both numpy and dask backends. Three explicit anchors pin the
exact padded layout (member order, fill position,
(nterm, max_size)interleaving, and the
_factoraxis) for the linear, multidim and quadraticcases.
Benchmark (300k elem × 8 dim × 1000 groups, numpy)
_sum_by_scatter(#793)_sum_by_unstack(#793 dask path)The unified kernel matches the scatter kernel's memory and time; the old dask
path cost 2.2× peak.
Notes
linopy/expressions.pyand the groupby kernel tests; fulltest_linear_expression+test_quadratic_expressionpass (366), broadersuite green.
group_diminto one chunk is unavoidable for a scatter (a group'smembers can sit in any chunk) and is exactly what the old unstack path forced.