fix(groupby): preserve grouped-dim position in groupby-sum#860
Merged
Conversation
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.
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.
for more information, see https://pre-commit.ci
Merging this PR will not alter performance
Comparing 
Footnotes
|
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).
Collaborator
|
@FabianHofmann Good catch, but it seems like this needs to be fixed differently to keep the memory low |
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.
for more information, see https://pre-commit.ci
Collaborator
|
Ill fix the memory issue! |
…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>
…im'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>
Collaborator
|
@FabianHofmann fixed the memory issue in c5d08f4 |
Collaborator
Author
|
@FBumann great, thanks! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
The following content was generated by AI.
Changes proposed in this Pull Request
LinearExpression.groupby(...).sum()moved the grouped dimension to the trailingposition instead of replacing it in place, diverging from xarray's native
groupby-reduce and from linopy <= 0.8.0 (regressed in #802).
assert_linequalaligns dimension order, so existing tests missed the reorder.
Root cause: both the
apply_ufuncscatter kernel (fast path) and thegroupby.mapfallback append the new group dimension at the end. The
LinearExpressionconstructor's
xr.broadcastthen reorders every variable to the dataset'scanonical dim order — which is driven by coordinate-insertion order — so a plain
transposealone doesn't stick.The fix adds
_restore_group_dim_position, applied to both paths:dimension(s) were, surviving dims keeping their relative order.
MultiIndexgroup coordinates untouched, so multikey /observed=Trueresults (which are order-agnostic) are unaffected and no
FutureWarningisemitted.
The strict
xfailmarker ontest_linear_expression_groupby_preserves_dim_positionis removed; the test now passes for both
use_fallback=FalseandTrue.Checklist
AGENTS.md).doc.doc/release_notes.rstof the upcoming release is included.