fix(groupby): group by non-dimension coordinate names; fast multi-key grouping by names (#750, #753)

[FBumann]

What

Brings LinearExpression.groupby to parity with xarray.Dataset.groupby for grouping by coordinate names — on the fast path, with no breaking change. Two related pieces:

1. Group by an attached non-dimension coordinate (closes #750).

expr = (1 * x).assign_coords(period=period)   # 'period' rides on the 'snapshot' dim
expr.groupby("period").sum()   # before: ValueError: period already exists as coordinate
expr.groupby(period).sum()     # before: KeyError: 'period'

Both now work and take the fast path, mirroring xarray. (Grouping by a dimension name always worked — only attached non-dim coordinates were broken.) The old workaround was to detach the coord first (expr.drop_vars("period").groupby(period).sum()).

2. Multi-key grouping by names now takes the fast path (closes #753).

expr.groupby(["period", "season"]).sum()   # before: silently dropped to the slow xarray fallback

It now rides the reindex fast path and returns the same output — one dimension per key, byte-identical to the fallback, sparse fill cells included. The pd.DataFrame grouper is untouched and keeps its compact stacked-MultiIndex output, so this is non-breaking.

How

In LinearExpressionGroupby:

• _resolve_group normalizes a key: unwrap a single-element list (groupby(["period"]) → scalar), and resolve a string coord name to its coordinate so it takes the fast path.
• sum() drops every coordinate aligned to the grouped dimension before reshaping, so an attached aux coord (including the one being grouped by) no longer collides on the final rename.
• The groupby property detaches a free (non-indexed) coordinate before handing it to xarray — fixing the use_fallback=True path — but never a MultiIndex level.
• A list of coord names (1-D, sharing one dim) is gathered into a value frame to ride the fast path, then unstacked back into one dimension per key.

Memory

One dimension per key is a dense cartesian grid, so a sparse key crossing materialises mostly-fill cells (measured ~100× vs the compact DataFrame grouper for a diagonal crossing — see #740). A UserWarning nudges sparse/high-cardinality users to the DataFrame grouper; it reads the collapsed MultiIndex levels, so it is O(observed), not O(N), and fires before unstack allocates the grid. Getting separate dims and compact storage would need a sparse/long-format kernel — tracked in #757 (the groupby-densification follow-up) under the sparse-kernel umbrella #756.

Tests

TestGroupbyByAttachedCoordinate asserts grouped vars/coeffs against hard-coded results on a deterministic model: single-key (name/DataArray × use_fallback), 2-D variable, dimension-coordinate-by-name, single-element list, MultiIndex level, and a pytest.raises row pinning that a list of DataArrays is unsupported. TestMultiKeyFastPath covers the multi-key fast path: fast == fallback (list/tuple, including a sparse crossing), separate-dims-not-stacked, sparse-combination-filled, DataFrame-grouper-stays-compact, and the blow-up warning (fires when sparse, silent when dense). Full suite green; ruff + mypy clean.

Context: relation to PyPSA usage and #744

This is a building block for the flat-dimension + auxiliary-level-coord direction discussed in #744. If n.snapshots becomes a flat dimension carrying period / scenario level coordinates (instead of a stacked pd.MultiIndex), aggregating an expression over a level becomes expr.groupby("period").sum() — exactly what this PR makes work, now also for multiple levels at once (groupby(["period", "scenario"])).

PyPSA still carries per-period workarounds citing a broken MultiIndex groupby (pydata/xarray#6836); that upstream bug is now fixed (verified on xarray 2025.9.0), but those comments are about MultiIndex grouping and per-period rolling, orthogonal to this PR.

Closes #750, #753. Drafted with an agent (Claude Code).

🤖 Generated with Claude Code

[FabianHofmann]

@FBumann

Note

The following content was generated by AI (Claude Code), posted on behalf of @FabianHofmann.

A follow-up commit added an observed parameter to LinearExpressionGroupby.sum (pandas-style, cf. the observed=True default flip in pandas 3.0) addressing the memory concern with multi-key grouping:

expr.groupby(["period", "season"]).sum(observed=True)

keeps the result stacked over only the observed key combinations (a MultiIndex group dimension) instead of unstacking into one dimension per key, which materialises the dense cartesian grid. The default observed=False keeps xarray parity (fast path == fallback, byte-identical); the dense-grid UserWarning now points to observed=True instead of the pd.DataFrame grouper workaround. observed=True with use_fallback=True raises, since the xarray fallback cannot honor it.

The multikey gather and unstack epilogue were also extracted into helpers (_multikey_value_frame, _unstack_multikey), and the release notes gained a Features entry.

[FBumann]

Amazing!