feat: add cluster_on to select which coordinates drive clustering#89
Conversation
`cluster_on` restricts the clustering to a subset of coordinates along `cluster_dim`. Coordinates left out are still aggregated and reconstructed from the resulting clusters, but have no influence on how the clusters are formed — useful for carrying prices/demand onto the representative grid without letting them drive period selection. Mirrors the `weights` formats: a list/bare string for a single cluster_dim, or a dict keyed by dimension for multiple. Internally clusters on the active subset then applies the clustering back to the full frame, so `weights` stays purely about relative importance (no magic zero-weight discontinuity). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR adds Changescluster_on feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant aggregate
participant _normalize_cluster_on
participant _aggregate_single
participant tsam
aggregate->>_normalize_cluster_on: cluster_on
_normalize_cluster_on-->>aggregate: active_coords
aggregate->>_aggregate_single: active_coords
alt passive columns exist
_aggregate_single->>tsam: aggregate(active subset)
tsam-->>_aggregate_single: clustering
_aggregate_single->>tsam: clustering.apply(full DataFrame)
tsam-->>_aggregate_single: reconstructed result
else all columns active
_aggregate_single->>tsam: aggregate(full DataFrame)
tsam-->>_aggregate_single: reconstructed result
end
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@src/tsam_xarray/_core.py`:
- Around line 445-448: Treat dict-form cluster_on string values as a single
coordinate instead of iterating over characters. Update the normalization logic
in the cluster_on handling branch of _core.py so the dict path matches the
list-form behavior: when a per-dimension value is a string, wrap it as one item
before converting to strings. Keep the existing validation flow in the
cluster_on processing code, but ensure {"variable": "solar"} becomes ["solar"],
not a character list.
- Around line 585-601: The passive-column path in the TSAM aggregation flow is
passing aggregate-only options into ClusteringResult.apply(), which only accepts
the data frame in tsam 3.4.0. Update the logic around tsam.aggregate() and
ClusteringResult.apply() in _core.py so that apply() is called without
temporal_resolution, round_decimals, numerical_tolerance, or any other
aggregate-only kwargs; either remove apply_kwargs entirely or filter it strictly
to parameters supported by apply(). Ensure the active-column branch still uses
the clustering.apply() result correctly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: aafe8147-4644-4d34-96de-1fdd7da59b36
📒 Files selected for processing (3)
docs/examples/multi-dim.ipynbsrc/tsam_xarray/_core.pytest/test_aggregate.py
| if isinstance(cluster_on, dict): | ||
| per_dim: dict[str, list[str]] = { | ||
| k: [str(c) for c in v] for k, v in cluster_on.items() | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and inspect the relevant region.
wc -l src/tsam_xarray/_core.py
ast-grep outline src/tsam_xarray/_core.py --view expanded | sed -n '1,220p'
sed -n '420,490p' src/tsam_xarray/_core.pyRepository: FBumann/tsam_xarray
Length of output: 3571
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from collections.abc import Sequence
v = "solar"
print("isinstance(v, Sequence):", isinstance(v, Sequence))
print("[str(c) for c in v]:", [str(c) for c in v])
cluster_on = {"variable": "solar"}
per_dim = {k: [str(c) for c in v] for k, v in cluster_on.items()}
print("per_dim:", per_dim)
PYRepository: FBumann/tsam_xarray
Length of output: 283
Treat dict-form cluster_on string values as one coord
{"variable": "solar"} is expanded into ['s', 'o', 'l', 'a', 'r'], which then fails validation with an unknown-coords error. Normalize per-dim string values the same way the list-form branch already does.
Proposed fix
if isinstance(cluster_on, dict):
per_dim: dict[str, list[str]] = {
- k: [str(c) for c in v] for k, v in cluster_on.items()
+ k: [v] if isinstance(v, str) else [str(c) for c in v]
+ for k, v in cluster_on.items()
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if isinstance(cluster_on, dict): | |
| per_dim: dict[str, list[str]] = { | |
| k: [str(c) for c in v] for k, v in cluster_on.items() | |
| } | |
| if isinstance(cluster_on, dict): | |
| per_dim: dict[str, list[str]] = { | |
| k: [v] if isinstance(v, str) else [str(c) for c in v] | |
| for k, v in cluster_on.items() | |
| } |
🤖 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/_core.py` around lines 445 - 448, Treat dict-form cluster_on
string values as a single coordinate instead of iterating over characters.
Update the normalization logic in the cluster_on handling branch of _core.py so
the dict path matches the list-form behavior: when a per-dimension value is a
string, wrap it as one item before converting to strings. Keep the existing
validation flow in the cluster_on processing code, but ensure {"variable":
"solar"} becomes ["solar"], not a character list.
Clean up the cluster_on implementation without changing the public API
or any test:
- Factor the duplicated dim/coord validation shared by weights and
cluster_on into _validate_dim_coord_mapping().
- Collapse _col_is_active + _partition_by_cluster_on into a single
_active_columns() pass; "has passive columns" is len(active) != total.
- Move the empty-selection ("nothing to cluster on") check up-front into
_normalize_cluster_on, so bad input fails before any aggregation runs
instead of deep in the per-slice loop.
517 passed, mypy + ruff clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@src/tsam_xarray/_core.py`:
- Around line 590-591: The passive-column path is incorrectly forwarding
filtered kwargs into ClusteringResult.apply(), but that method only accepts the
DataFrame and will raise a TypeError. Update the apply logic in the clustering
flow around clustering.apply so it is called with df only, and either stop
constructing apply_kwargs here or keep _APPLY_KWARGS empty so
temporal_resolution, round_decimals, and numerical_tolerance are not passed
through.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| apply_kwargs = {k: v for k, v in tsam_kwargs.items() if k in _APPLY_KWARGS} | ||
| tsam_result = clustering.apply(df, **apply_kwargs) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '505,516p' src/tsam_xarray/_core.py
rg -n "_APPLY_KWARGS" src/tsam_xarray/_core.pyRepository: FBumann/tsam_xarray
Length of output: 702
🌐 Web query:
tsam 3.4.0 ClusteringResult.apply signature temporal_resolution round_decimals numerical_tolerance source
💡 Result:
In tsam 3.4.0, the ClusteringResult.apply() method is used to transfer a previously computed clustering configuration to new input data [1][2][3]. The signature of this method is: ClusteringResult.apply(data: pandas.DataFrame) -> AggregationResult The method takes a single argument, data, which is the pandas DataFrame to which the existing clustering configuration will be applied [4][3]. Regarding the parameters temporal_resolution, round_decimals, and numerical_tolerance: These parameters are generally associated with the initial tsam.aggregate() function (used to create the clustering in the first place) rather than the ClusteringResult.apply() method [5][6]. When performing the initial aggregation, they serve the following purposes [5][6]: - temporal_resolution: Defines the time resolution of the input data. It accepts a float (representing hours) or a pandas Timedelta string (e.g., '1h', '15min'). If not provided, it is inferred from the datetime index. - round_decimals: An integer specifying the number of decimal places to which the results should be rounded. - numerical_tolerance: A float (default is 1e-13) used for floating-point comparisons during the aggregation process. Because ClusteringResult.apply() uses a pre-calculated clustering (such as predefined cluster centers or indices), these tuning parameters are typically set during the initial aggregate() call and are encapsulated within the ClusteringResult object itself [4][3]. Applying the result to new data simply maps the new data to the existing cluster representatives [4][2].
Citations:
- 1: https://github.com/FZJ-IEK3-VSA/tsam/releases
- 2: https://tsam.readthedocs.io/en/latest/newsDoc.html
- 3: https://tsam.readthedocs.io/en/latest/migration-guide/
- 4: https://tsam.readthedocs.io/en/latest/notebooks/clustering_transfer/
- 5: https://tsam.readthedocs.io/en/latest/api/tsam/tuning/
- 6: https://tsam.readthedocs.io/en/v3.1.1/timeseriesaggregationDoc.html
Don't forward kwargs into ClusteringResult.apply()
ClusteringResult.apply() only accepts the DataFrame, so passing temporal_resolution, round_decimals, or numerical_tolerance here will raise TypeError on the passive-column path. Call clustering.apply(df) directly, or leave _APPLY_KWARGS empty.
🤖 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/_core.py` around lines 590 - 591, The passive-column path is
incorrectly forwarding filtered kwargs into ClusteringResult.apply(), but that
method only accepts the DataFrame and will raise a TypeError. Update the apply
logic in the clustering flow around clustering.apply so it is called with df
only, and either stop constructing apply_kwargs here or keep _APPLY_KWARGS empty
so temporal_resolution, round_decimals, and numerical_tolerance are not passed
through.
1 participant
What
Adds a
cluster_onparameter toaggregate()that restricts which coordinates alongcluster_dimdrive the clustering. Coordinates left out are still aggregated and reconstructed from the resulting clusters, but have no influence on how the clusters are formed.Why
Common ask: carry a price/demand series onto the representative grid without letting it drive period selection.
The obvious alternative —
weights={"price": 0}— is a footgun: tsam silently clamps a 0 weight to a "minimal tolerable weighting", so the variable still weakly influences the clustering. Encoding a hard structural decision (which columns enter the distance metric) as a magic value of the softweightsknob also creates a0.0-vs-1e-9discontinuity.cluster_onkeepsweightspurely about relative importance and expresses exclusion explicitly.How
Internally clusters on the active subset via
tsam.aggregate(), then applies that clustering back to the full frame (ClusteringResult.apply()) so passive columns get representatives. Composes withweights(which apply to the active subset). Mirrors theweightsAPI shapes.Validation covers: unknown dims/coords, list-form with multiple
cluster_dim, and an empty selection ("nothing to cluster on").Tests
New parametrized
TestClusterOn(fixtures +parametrize) proving a carried variable yields identical cluster assignments to clustering the subset alone, plus error-path cases. Full suite: 517 passed, mypy + ruff clean.Docs
New "Selecting what to cluster on" section in
docs/examples/multi-dim.ipynb; executes end-to-end.🤖 Generated with Claude Code
Summary by CodeRabbit
cluster_ontoaggregate()to control which coordinates influence clustering, while still reconstructing and returning all coordinates/variables.weights.cluster_onuse, interaction withweights, differing outcomes vs full-variable clustering, and error cases.