perf: disaggregate unsegmented clusterings with a numpy gather - #116
Conversation
|
Warning Review limit reached
Next review available in: 49 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthrough
ChangesDisaggregation fast path
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AggregationResult
participant ClusteringResult
participant _disaggregate_gather
AggregationResult->>ClusteringResult: delegate disaggregate(data)
ClusteringResult->>ClusteringResult: build and validate per-slice clusterings
ClusteringResult->>_disaggregate_gather: gather compatible unsegmented assignments
_disaggregate_gather-->>ClusteringResult: return expanded time-series data
ClusteringResult-->>AggregationResult: return disaggregated result
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/tsam_xarray/_clustering.py (1)
713-741: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
Args:/Returns:sections to the new helper docstrings.
_is_gatherabledocuments behavior in prose only, and_same_time_indexhas no docstring. As per coding guidelines, use Google-style docstrings withArgs:,Returns:,Raises:,Attributes:sections.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tsam_xarray/_clustering.py` around lines 713 - 741, Add Google-style docstrings to `_is_gatherable` and `_same_time_index`, documenting each parameter under `Args:` and the boolean result under `Returns:`; preserve the existing behavior and prose where useful, and do not add unsupported `Raises:` or `Attributes:` sections.Source: Coding guidelines
test/test_aggregate.py (1)
1671-1723: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering the gather/fallback equivalence.
The new tests cover dim order, label gathering, dtype, and error paths, but nothing asserts that the vectorized sliced path produces the same values as the per-slice tsam fallback (e.g. slices with differing stored
time_index, which makes_is_gatherablereturnFalse). A single equivalence test would lock in the fast path's correctness.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/test_aggregate.py` around lines 1671 - 1723, Add a test alongside the existing disaggregate tests that constructs slices with differing stored time_index values so _is_gatherable returns False, then compares vectorized sliced disaggregation against the per-slice tsam fallback. Assert the resulting values are identical, while preserving the current coverage for dimension order, labels, dtype, and errors.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/tsam_xarray/_clustering.py`:
- Around line 713-741: Add Google-style docstrings to `_is_gatherable` and
`_same_time_index`, documenting each parameter under `Args:` and the boolean
result under `Returns:`; preserve the existing behavior and prose where useful,
and do not add unsupported `Raises:` or `Attributes:` sections.
In `@test/test_aggregate.py`:
- Around line 1671-1723: Add a test alongside the existing disaggregate tests
that constructs slices with differing stored time_index values so _is_gatherable
returns False, then compares vectorized sliced disaggregation against the
per-slice tsam fallback. Assert the resulting values are identical, while
preserving the current coverage for dimension order, labels, dtype, and errors.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c6f94bc2-8e18-49d2-960e-fc11d5eb9556
📒 Files selected for processing (4)
benchmarks/test_bench_aggregate.pysrc/tsam_xarray/_clustering.pysrc/tsam_xarray/_result.pytest/test_aggregate.py
Expanding an unsegmented clustering back to the full time axis is a gather along the cluster axis. The wrapper expressed it as a pandas round-trip — build a MultiIndexed DataFrame, hand it to tsam's disaggregate() (which does unstack -> .loc[assignments] -> stack), convert back — once per slice-dim combination. Callers expand every variable of a solution, so a 12-period model paid n_variables x 12 round-trips for what is one gather. The gather now runs in numpy across all slices at once. xarray guarantees the regular grid the fast path needs (dense cluster x timestep block, equal block sizes, one shared timestep order), so no structural guards are needed; the path declines only what it cannot express as one rectangular result: segmented clusterings, and slices whose stored time indices disagree. Those fall back to the per-slice tsam path unchanged. Two things the gather preserves, both easy to get wrong: - Dim order stays (*slice_dims, time, *other_dims), matching the per-slice concat. The naive vectorized form yields (time, *slice_dims), which compares equal under xarray ops but silently transposes .values. - Cluster assignments are labels, not positions — .loc[assignments] indexes by label, so labels are mapped to positions before indexing. Validated against the previous implementation on 19 shapes (slice-dim counts, segmentation, custom dim names, JSON round-trip, shuffled cluster and timestep coords, int/float32/bool/NaN payloads, extra trailing dims): values, dtype, shape, dim order and coords identical in every case. Error type and text for mismatched cluster IDs are unchanged. 365 days, 36 clusters, 100 variables: 481 ms -> 20 ms (24x) without slice dims, 6.2 s -> 146 ms (42x) across 12 slices. AggregationResult.disaggregate() now delegates to the identical loop on ClusteringResult instead of building a per-slice result view. Closes #115. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #116. _same_time_index had no docstring; both it and _is_gatherable now carry the project's Google-style Args/Returns, since which inputs make a clustering gatherable is the non-obvious part of the fast path. The time-index guard had no test. Slices whose stored indices disagree now have one: it asserts the gather is declined and that each slice still matches its own expansion through tsam. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
e2f20c3 to
e8e95f2
Compare
BenchmarksΔ% vs base, one table sorted by biggest change first Full tablebenchmarks/test_bench_aggregate.py
|
1 participant
Closes #115.
Expanding an unsegmented clustering back to the full time axis is a gather along the cluster axis. The wrapper expressed it as a pandas round-trip — build a MultiIndexed DataFrame, hand it to tsam's
disaggregate()(which doesunstack→.loc[assignments]→stack), convert back — once per slice-dim combination. Callers expand every variable of a solution, so a 12-period model paidn_variables × 12round-trips for what is one gather.The gather now runs in numpy across all slices at once.
Results
365 days, 36 clusters, 100 variables:
Why not just wait for the upstream tsam fix
FBumann/tsam#61 makes tsam's own
disaggregate()~22x faster. It is worth having, but it does not replace this. Keeping the per-slice wrapper loop and swapping in a shim that does exactly what that PR does:Without slice dims they land in the same ballpark. With slices, upstream alone leaves 511 ms on the table, because once tsam is fast the remaining cost is the wrapper: 12 DataFrame constructions, 12
cr.disaggregate()calls, 12DataArraybuilds, onexr.concat. No upstream change removes that loop.They also do not overlap. This PR bypasses tsam for unsegmented clusterings, so when #61 lands its win goes to the segmented path — which this PR deliberately leaves alone. Nothing to unwind later, and no tsam version gate.
Fallbacks
xarray guarantees the regular grid the fast path needs — a dense
cluster × timestepblock, equal block sizes, one shared timestep order — so the structural guards a pandas-level fix needs are unnecessary here. The path declines only what it cannot express as one rectangular result: segmented clusterings, and slices whose stored time indices disagree. Both fall back to the per-slice tsam path unchanged.Two things a patch here has to preserve
_concat_along_dimsleaves slice dims leading:(*slice_dims, time, *other_dims). The naive vectorized form yields(time, *slice_dims)— identical under xarray comparisons, but.valuessilently transposes from(12, 8761)to(8761, 12), breaking positional access downstream._expand_periodsindexes with.loc[list(cluster_assignments)], so assignments are cluster labels and must be mapped to positions rather than used as indices.Verification
DimNames, JSON round-trip, shuffled cluster and timestep coords, int/float32/bool/NaN payloads, extra trailing dims, missingtime_index. Values, dtype, shape, dim order and coords identical in every case.data has 12 timesteps per clusterrather than tsam'scluster 0 has 12 timesteps, since in xarray all clusters share the axis.ruffandmypyclean.AggregationResult.disaggregate()now delegates to the identical loop onClusteringResultinstead of building a per-slice result view.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests