tsam-xarray 0.6.0__tar.gz → 0.6.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (44) hide show
  1. tsam_xarray-0.6.1/.release-please-manifest.json +3 -0
  2. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/CHANGELOG.md +8 -0
  3. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/PKG-INFO +1 -1
  4. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/examples/multi-dim.ipynb +39 -1
  5. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_core.py +264 -26
  6. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_version.py +2 -2
  7. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/test/test_aggregate.py +163 -0
  8. tsam_xarray-0.6.0/.release-please-manifest.json +0 -3
  9. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/dependabot.yml +0 -0
  10. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/workflows/ci.yaml +0 -0
  11. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/workflows/dependabot-auto-merge.yaml +0 -0
  12. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/workflows/pr-title.yaml +0 -0
  13. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/workflows/publish.yaml +0 -0
  14. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.github/workflows/release.yaml +0 -0
  15. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.gitignore +0 -0
  16. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.pre-commit-config.yaml +0 -0
  17. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.readthedocs.yaml +0 -0
  18. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/.release-please-config.json +0 -0
  19. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/CLAUDE.md +0 -0
  20. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/LICENSE +0 -0
  21. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/README.md +0 -0
  22. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/codecov.yml +0 -0
  23. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/assets/multi-dim-input.png +0 -0
  24. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/assets/multi-dim-metrics.png +0 -0
  25. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/changelog.md +0 -0
  26. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/data-model.md +0 -0
  27. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/examples/clustering-io.ipynb +0 -0
  28. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/examples/getting-started.ipynb +0 -0
  29. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/examples/segmentation.ipynb +0 -0
  30. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/examples/tuning.ipynb +0 -0
  31. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/gen_ref_pages.py +0 -0
  32. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/index.md +0 -0
  33. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/docs/stylesheets/extra.css +0 -0
  34. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/mkdocs.yml +0 -0
  35. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/pyproject.toml +0 -0
  36. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/scripts/generate_readme_images.py +0 -0
  37. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/__init__.py +0 -0
  38. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_clustering.py +0 -0
  39. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_result.py +0 -0
  40. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_sample_data.py +0 -0
  41. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/src/tsam_xarray/_tuning.py +0 -0
  42. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/test/conftest.py +0 -0
  43. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/test/test_parametrized.py +0 -0
  44. {tsam_xarray-0.6.0 → tsam_xarray-0.6.1}/test/test_tuning.py +0 -0
@@ -0,0 +1,3 @@
1
+ {
2
+ ".": "0.6.1"
3
+ }
@@ -1,5 +1,13 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.6.1](https://github.com/FBumann/tsam_xarray/compare/v0.6.0...v0.6.1) (2026-07-09)
4
+
5
+
6
+ ### Features
7
+
8
+ * add cluster_on to select which coordinates drive clustering ([#89](https://github.com/FBumann/tsam_xarray/issues/89)) ([ddff73b](https://github.com/FBumann/tsam_xarray/commit/ddff73b24243e53d5a182b9e8c794e2b57a8d1b7))
9
+ * validate extremes options against cluster_on ([#91](https://github.com/FBumann/tsam_xarray/issues/91)) ([2fa5a91](https://github.com/FBumann/tsam_xarray/commit/2fa5a91c5222f8b6dd6932d9572dec62e90f595b))
10
+
3
11
  ## [0.6.0](https://github.com/FBumann/tsam_xarray/compare/v0.5.2...v0.6.0) (2026-05-27)
4
12
 
5
13
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tsam_xarray
3
- Version: 0.6.0
3
+ Version: 0.6.1
4
4
  Summary: Lightweight xarray wrapper for tsam time series aggregation
5
5
  License-Expression: MIT
6
6
  License-File: LICENSE
@@ -11,7 +11,7 @@
11
11
  "- **`cluster_dim`** — dimensions clustered together (shared clustering)\n",
12
12
  "- **Auto-slicing** — remaining dimensions get independent clusterings\n",
13
13
  "\n",
14
- "This notebook covers stacking, slicing, and weights."
14
+ "This notebook covers stacking, slicing, weights, and selecting what to cluster on."
15
15
  ]
16
16
  },
17
17
  {
@@ -212,6 +212,44 @@
212
212
  ")\n",
213
213
  "result_w2.accuracy.rmse.to_dataframe(\"RMSE\")"
214
214
  ]
215
+ },
216
+ {
217
+ "cell_type": "markdown",
218
+ "metadata": {},
219
+ "source": [
220
+ "## Selecting what to cluster on\n",
221
+ "\n",
222
+ "Sometimes a variable should be *carried through* the aggregation without\n",
223
+ "influencing how the clusters are formed — e.g. a price or demand series you\n",
224
+ "want on the representative grid, but that shouldn't drive period selection.\n",
225
+ "\n",
226
+ "Use `cluster_on` to restrict the clustering to specific coordinates. Everything\n",
227
+ "else is still aggregated and reconstructed from the resulting clusters, but has\n",
228
+ "no effect on the cluster distances. This differs from a small `weights` value:\n",
229
+ "`cluster_on` fully removes the coordinate from the distance metric.\n",
230
+ "\n",
231
+ "Here we cluster on `solar` and `wind` only, while `demand` is carried along."
232
+ ]
233
+ },
234
+ {
235
+ "cell_type": "code",
236
+ "execution_count": null,
237
+ "metadata": {},
238
+ "outputs": [],
239
+ "source": [
240
+ "# Cluster on solar + wind only; demand is carried along but doesn't\n",
241
+ "# influence cluster selection.\n",
242
+ "result_on = tsam_xarray.aggregate(\n",
243
+ " da_single,\n",
244
+ " time_dim=\"time\",\n",
245
+ " cluster_dim=[\"variable\", \"region\"],\n",
246
+ " n_clusters=4,\n",
247
+ " cluster_on={\"variable\": [\"solar\", \"wind\"]},\n",
248
+ ")\n",
249
+ "# demand is still present in the output, aggregated from the same clusters\n",
250
+ "print(\"Variables:\", list(result_on.cluster_representatives.coords[\"variable\"].values))\n",
251
+ "result_on.accuracy.rmse.to_dataframe(\"RMSE\")"
252
+ ]
215
253
  }
216
254
  ],
217
255
  "metadata": {
@@ -14,6 +14,7 @@ import xarray as xr
14
14
  from tsam_xarray._result import AccuracyMetrics, AggregationResult
15
15
 
16
16
  Weights = dict[str, float] | dict[str, dict[str, float]] | None
17
+ ClusterOn = str | Sequence[str] | dict[str, Sequence[str]] | None
17
18
 
18
19
 
19
20
  def aggregate(
@@ -23,6 +24,7 @@ def aggregate(
23
24
  cluster_dim: Sequence[str] | str,
24
25
  n_clusters: int,
25
26
  weights: Weights = None,
27
+ cluster_on: ClusterOn = None,
26
28
  **tsam_kwargs: Any,
27
29
  ) -> AggregationResult:
28
30
  """Aggregate an xarray DataArray using tsam.
@@ -55,6 +57,37 @@ def aggregate(
55
57
  e.g. ``("solar", "north")`` gets weight
56
58
  ``2.0 * 1.5 = 3.0``.
57
59
 
60
+ cluster_on: Restrict which coordinates drive the
61
+ clustering. Coordinates not selected are still
62
+ aggregated and reconstructed from the resulting
63
+ clusters, but have no influence on how the clusters
64
+ are formed. ``None`` (default) clusters on
65
+ everything. Mirrors the ``weights`` formats:
66
+
67
+ - **List** (single ``cluster_dim``)::
68
+
69
+ cluster_on=["solar", "wind"]
70
+
71
+ - **Dict** (multiple ``cluster_dim``)::
72
+
73
+ cluster_on={"variable": ["solar", "wind"]}
74
+
75
+ A dim omitted from the dict is unrestricted (all
76
+ its coordinates are eligible). When several dims
77
+ are listed, a column must match on *all* of them
78
+ to be clustered on. At least one column must
79
+ remain selected.
80
+
81
+ Not compatible with ``ExtremeConfig(method=
82
+ "replace")`` — the carried columns are filled by
83
+ transferring the clustering, which cannot
84
+ reproduce the hybrid 'replace' representative. Use
85
+ a low ``weights`` value to de-emphasise a column
86
+ instead of excluding it. An ``ExtremeConfig`` also
87
+ may not reference an excluded coordinate, since
88
+ extreme periods are identified only on the
89
+ clustered-on columns.
90
+
58
91
  **tsam_kwargs: Additional keyword arguments passed to
59
92
  ``tsam.aggregate()``.
60
93
  """
@@ -65,10 +98,18 @@ def aggregate(
65
98
  da = _validate_data(da, time_dim, col_dims, slice_dims)
66
99
  _validate_no_cluster_config_weights(tsam_kwargs)
67
100
  per_dim_weights = _normalize_weights(weights, da, col_dims)
101
+ active_coords = _normalize_cluster_on(cluster_on, da, col_dims)
102
+ _validate_extremes_with_cluster_on(tsam_kwargs, active_coords, da)
68
103
 
69
104
  if not slice_dims:
70
105
  return _aggregate_single(
71
- da, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
106
+ da,
107
+ n_clusters,
108
+ time_dim,
109
+ col_dims,
110
+ per_dim_weights,
111
+ active_coords,
112
+ tsam_kwargs,
72
113
  )
73
114
 
74
115
  slice_coords = {d: da.coords[d].values for d in slice_dims}
@@ -80,7 +121,13 @@ def aggregate(
80
121
  sel = dict(zip(slice_dims, key, strict=True))
81
122
  da_slice = da.sel(sel)
82
123
  r = _aggregate_single(
83
- da_slice, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
124
+ da_slice,
125
+ n_clusters,
126
+ time_dim,
127
+ col_dims,
128
+ per_dim_weights,
129
+ active_coords,
130
+ tsam_kwargs,
84
131
  )
85
132
  results.append(r)
86
133
 
@@ -129,6 +176,68 @@ def _validate_no_cluster_config_weights(
129
176
  raise ValueError(msg)
130
177
 
131
178
 
179
+ def _validate_extremes_with_cluster_on(
180
+ tsam_kwargs: dict[str, Any],
181
+ active_coords: dict[str, set[str]] | None,
182
+ da: xr.DataArray,
183
+ ) -> None:
184
+ """Reject extremes settings incompatible with the cluster_on transfer.
185
+
186
+ cluster_on clusters on the active subset and transfers the result to the
187
+ carried columns via ``ClusteringResult.apply()``. Two extremes settings
188
+ don't survive that transfer:
189
+
190
+ - ``method="replace"`` builds a hybrid representative (some columns from the
191
+ medoid, some from the extreme period) that apply() cannot reproduce — tsam
192
+ silently falls back to the medoid. Rather than degrade quietly, error and
193
+ point to ``weights``, which de-emphasises a column without excluding it
194
+ (tsam clamps a low weight to a minimal one, so the column stays in the
195
+ single aggregation pass and 'replace' remains reproducible).
196
+ - An ExtremeConfig that references a coordinate cluster_on excluded: those
197
+ columns never enter the clustering, so tsam can't identify extremes on
198
+ them (it would raise a misleading "not found in data"). Surface the real
199
+ cause instead.
200
+ """
201
+ if active_coords is None:
202
+ return
203
+ extremes = tsam_kwargs.get("extremes")
204
+ if extremes is None:
205
+ return
206
+
207
+ if getattr(extremes, "method", None) == "replace":
208
+ msg = (
209
+ "extremes method 'replace' is not supported together with "
210
+ "cluster_on. cluster_on fully excludes the carried columns and "
211
+ "transfers the clustering to them, which cannot reproduce the "
212
+ "hybrid 'replace' representative. To reduce a column's influence "
213
+ "without excluding it, give it a low 'weights' value instead of "
214
+ "using cluster_on (weights never removes a column entirely, so "
215
+ "'replace' stays reproducible)."
216
+ )
217
+ raise ValueError(msg)
218
+
219
+ passive_coords: set[str] = set()
220
+ for dim, active_set in active_coords.items():
221
+ all_coords = {str(c) for c in da.coords[dim].values}
222
+ passive_coords |= all_coords - active_set
223
+
224
+ referenced = (
225
+ set(extremes.max_value)
226
+ | set(extremes.min_value)
227
+ | set(extremes.max_period)
228
+ | set(extremes.min_period)
229
+ )
230
+ excluded = referenced & passive_coords
231
+ if excluded:
232
+ msg = (
233
+ f"extremes references coordinates {excluded} that cluster_on "
234
+ "excluded from the clustering. Extreme periods can only be "
235
+ "identified on clustered-on columns — add these coordinates to "
236
+ "cluster_on, or remove them from the ExtremeConfig."
237
+ )
238
+ raise ValueError(msg)
239
+
240
+
132
241
  def _validate(
133
242
  da: xr.DataArray,
134
243
  time_dim: str,
@@ -265,7 +374,7 @@ def _representatives_to_da(
265
374
  df = df.copy()
266
375
  # With segmentation, index has 3 levels: (cluster, segment_step, segment_duration)
267
376
  # Without: 2 levels: (cluster, timestep)
268
- if df.index.nlevels == 3:
377
+ if isinstance(df.index, pd.MultiIndex) and df.index.nlevels == 3:
269
378
  df.index = df.index.droplevel(2) # drop segment_duration
270
379
  df.index.names = ["cluster", "timestep"]
271
380
 
@@ -339,6 +448,36 @@ def _metric_to_da(
339
448
  return xr.DataArray(series.to_xarray())
340
449
 
341
450
 
451
+ def _validate_dim_coord_mapping(
452
+ mapping: dict[str, Any],
453
+ da: xr.DataArray,
454
+ col_dims: list[str],
455
+ label: str,
456
+ ) -> None:
457
+ """Validate a ``{dim: coords}`` mapping references real dims and coords.
458
+
459
+ ``coords`` may be any iterable of coordinate names (a list, or the keys
460
+ of a ``{coord: weight}`` dict). Shared by ``weights`` and ``cluster_on``.
461
+ """
462
+ extra_dims = set(mapping.keys()) - set(col_dims)
463
+ if extra_dims:
464
+ msg = (
465
+ f"{label} has unknown dims {extra_dims}, "
466
+ f"must be subset of cluster_dim {col_dims}"
467
+ )
468
+ raise ValueError(msg)
469
+
470
+ for dim_name, coords in mapping.items():
471
+ valid_coords = set(str(c) for c in da.coords[dim_name].values)
472
+ unknown = set(str(c) for c in coords) - valid_coords
473
+ if unknown:
474
+ msg = (
475
+ f"{label} has unknown coords {unknown} for dim {dim_name!r}, "
476
+ f"valid coords: {sorted(valid_coords)}"
477
+ )
478
+ raise ValueError(msg)
479
+
480
+
342
481
  def _normalize_weights(
343
482
  weights: dict[str, float] | dict[str, dict[str, float]] | None,
344
483
  da: xr.DataArray,
@@ -370,27 +509,51 @@ def _normalize_weights(
370
509
  raise ValueError(msg)
371
510
  per_dim_weights = {col_dims[0]: weights} # type: ignore[dict-item]
372
511
 
373
- # Validate dim names exist in cluster_dim
374
- extra_dims = set(per_dim_weights.keys()) - set(col_dims)
375
- if extra_dims:
376
- msg = (
377
- f"weights has unknown dims {extra_dims}, "
378
- f"must be subset of cluster_dim {col_dims}"
379
- )
380
- raise ValueError(msg)
512
+ _validate_dim_coord_mapping(per_dim_weights, da, col_dims, "weights")
513
+ return per_dim_weights
381
514
 
382
- # Validate coord values exist in the DataArray
383
- for dim_name, coord_weights in per_dim_weights.items():
384
- valid_coords = set(str(c) for c in da.coords[dim_name].values)
385
- unknown = set(coord_weights.keys()) - valid_coords
386
- if unknown:
515
+
516
+ def _normalize_cluster_on(
517
+ cluster_on: ClusterOn,
518
+ da: xr.DataArray,
519
+ col_dims: list[str],
520
+ ) -> dict[str, set[str]] | None:
521
+ """Normalize cluster_on to a per-dim set of eligible coords.
522
+
523
+ A list/str form requires a single cluster_dim. Returns ``None``
524
+ when nothing is restricted (cluster on everything).
525
+ """
526
+ if cluster_on is None:
527
+ return None
528
+
529
+ if isinstance(cluster_on, dict):
530
+ per_dim: dict[str, list[str]] = {
531
+ k: [str(c) for c in v] for k, v in cluster_on.items()
532
+ }
533
+ else:
534
+ if len(col_dims) != 1:
387
535
  msg = (
388
- f"weights has unknown coords {unknown} for dim {dim_name!r}, "
389
- f"valid coords: {sorted(valid_coords)}"
536
+ "List-form cluster_on requires a single cluster_dim. "
537
+ "For multiple cluster_dim, use a dict: "
538
+ '{"dim_name": ["coord", ...]}.'
390
539
  )
391
540
  raise ValueError(msg)
541
+ coords = [cluster_on] if isinstance(cluster_on, str) else list(cluster_on)
542
+ per_dim = {col_dims[0]: [str(c) for c in coords]}
392
543
 
393
- return per_dim_weights
544
+ _validate_dim_coord_mapping(per_dim, da, col_dims, "cluster_on")
545
+
546
+ # An empty coord list for any restricted dim leaves no column to cluster
547
+ # on (the stacked columns are the full coordinate grid). Fail here, before
548
+ # any aggregation runs, rather than deep in the per-slice loop.
549
+ if any(not coords for coords in per_dim.values()):
550
+ msg = (
551
+ "cluster_on selected no columns — nothing to cluster on. "
552
+ "Select at least one coordinate present in the data."
553
+ )
554
+ raise ValueError(msg)
555
+
556
+ return {d: set(cs) for d, cs in per_dim.items()}
394
557
 
395
558
 
396
559
  def _translate_weights(
@@ -414,28 +577,103 @@ def _translate_weights(
414
577
  return flat
415
578
 
416
579
 
580
+ # tsam.aggregate() kwargs that ClusteringResult.apply() also accepts. Only
581
+ # these are forwarded to apply() when re-aggregating passive columns onto a
582
+ # clustering built from the active subset; the rest (period_duration, cluster,
583
+ # segments, extremes, ...) are already baked into the stored clustering.
584
+ _APPLY_KWARGS = frozenset(
585
+ {"temporal_resolution", "round_decimals", "numerical_tolerance"}
586
+ )
587
+
588
+
589
+ def _active_columns(
590
+ df: pd.DataFrame,
591
+ active_coords: dict[str, set[str]] | None,
592
+ col_dims: list[str],
593
+ ) -> list[Hashable]:
594
+ """Columns that drive the clustering, given a cluster_on selection.
595
+
596
+ A column is active only if, for every restricted dim, its coordinate is
597
+ in that dim's eligible set; unrestricted dims impose no constraint. With
598
+ no restriction (``active_coords is None``) every column is active. The
599
+ remaining columns are *passive*: still aggregated and reconstructed, but
600
+ excluded from the cluster distances.
601
+ """
602
+ if active_coords is None:
603
+ return list(df.columns)
604
+
605
+ def is_active(col: Hashable) -> bool:
606
+ pairs = (
607
+ zip(col_dims, col, strict=True)
608
+ if isinstance(col, tuple)
609
+ else [(col_dims[0], col)]
610
+ )
611
+ return all(
612
+ str(coord) in active_coords[dim]
613
+ for dim, coord in pairs
614
+ if dim in active_coords
615
+ )
616
+
617
+ return [col for col in df.columns if is_active(col)]
618
+
619
+
417
620
  def _aggregate_single(
418
621
  da: xr.DataArray,
419
622
  n_clusters: int,
420
623
  time_dim: str,
421
624
  col_dims: list[str],
422
625
  weights: dict[str, dict[str, float]] | None,
626
+ active_coords: dict[str, set[str]] | None,
423
627
  tsam_kwargs: dict[str, Any],
424
628
  ) -> AggregationResult:
425
- """Run a single tsam aggregation on a DataArray."""
629
+ """Run a single tsam aggregation on a DataArray.
630
+
631
+ When ``cluster_on`` restricts the clustering, the clustering runs on
632
+ the active subset only, then is applied back to the full frame so the
633
+ passive columns still get representatives and reconstruction.
634
+ """
426
635
  df = _to_dataframe(da, time_dim, col_dims)
427
636
 
428
637
  tsam_weights: dict[Hashable, float] | None = None
429
638
  if weights is not None:
430
639
  tsam_weights = _translate_weights(weights, df, col_dims)
431
640
 
432
- tsam_result = tsam.aggregate(
433
- df,
434
- n_clusters,
435
- weights=tsam_weights, # type: ignore[arg-type]
436
- **tsam_kwargs,
437
- )
641
+ active_cols = _active_columns(df, active_coords, col_dims)
438
642
 
643
+ if len(active_cols) == len(df.columns):
644
+ # Nothing carried passively — cluster on the full frame directly.
645
+ tsam_result = tsam.aggregate(
646
+ df,
647
+ n_clusters,
648
+ weights=tsam_weights, # type: ignore[arg-type]
649
+ **tsam_kwargs,
650
+ )
651
+ else:
652
+ # Cluster on the active subset, then apply that clustering back to the
653
+ # full frame so passive columns get representatives and reconstruction.
654
+ active_weights = (
655
+ None if tsam_weights is None else {c: tsam_weights[c] for c in active_cols}
656
+ )
657
+ clustering = tsam.aggregate(
658
+ df[active_cols],
659
+ n_clusters,
660
+ weights=active_weights, # type: ignore[arg-type]
661
+ **tsam_kwargs,
662
+ ).clustering
663
+ apply_kwargs = {k: v for k, v in tsam_kwargs.items() if k in _APPLY_KWARGS}
664
+ tsam_result = clustering.apply(df, **apply_kwargs)
665
+
666
+ return _result_from_tsam(tsam_result, da, df, time_dim, col_dims)
667
+
668
+
669
+ def _result_from_tsam(
670
+ tsam_result: Any,
671
+ da: xr.DataArray,
672
+ df: pd.DataFrame,
673
+ time_dim: str,
674
+ col_dims: list[str],
675
+ ) -> AggregationResult:
676
+ """Build an AggregationResult from a tsam aggregation result."""
439
677
  typical = _representatives_to_da(tsam_result.cluster_representatives, col_dims)
440
678
  reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims)
441
679
 
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '0.6.0'
22
- __version_tuple__ = version_tuple = (0, 6, 0)
21
+ __version__ = version = '0.6.1'
22
+ __version_tuple__ = version_tuple = (0, 6, 1)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -306,6 +306,169 @@ class TestWeights:
306
306
  )
307
307
 
308
308
 
309
+ @pytest.fixture
310
+ def da_3var() -> xr.DataArray:
311
+ """Single cluster_dim: solar/wind/price over time (region dropped)."""
312
+ return _make_da(variables=["solar", "wind", "price"]).isel(region=0, drop=True)
313
+
314
+
315
+ @pytest.fixture
316
+ def da_3var_region() -> xr.DataArray:
317
+ """Multi cluster_dim: solar/wind/price over time and region."""
318
+ return _make_da(variables=["solar", "wind", "price"])
319
+
320
+
321
+ class TestClusterOn:
322
+ """cluster_on restricts which coords drive clustering, keeping the rest."""
323
+
324
+ @pytest.mark.parametrize(
325
+ ("cluster_dim", "cluster_on", "subset"),
326
+ [
327
+ ("variable", ["solar", "wind"], ["solar", "wind"]),
328
+ ("variable", "solar", ["solar"]),
329
+ (
330
+ ["variable", "region"],
331
+ {"variable": ["solar", "wind"]},
332
+ ["solar", "wind"],
333
+ ),
334
+ ],
335
+ ids=["list", "string", "dict-multi"],
336
+ )
337
+ def test_matches_subset_clustering(
338
+ self,
339
+ da_3var: xr.DataArray,
340
+ da_3var_region: xr.DataArray,
341
+ cluster_dim: list[str] | str,
342
+ cluster_on: object,
343
+ subset: list[str],
344
+ ):
345
+ """cluster_on gives the same clustering as aggregating the subset alone,
346
+ while carried coords still appear in the output."""
347
+ da = da_3var if cluster_dim == "variable" else da_3var_region
348
+ restricted = tsam_xarray.aggregate(
349
+ da,
350
+ time_dim="time",
351
+ cluster_dim=cluster_dim,
352
+ n_clusters=4,
353
+ cluster_on=cluster_on,
354
+ )
355
+ reference = tsam_xarray.aggregate(
356
+ da.sel(variable=subset),
357
+ time_dim="time",
358
+ cluster_dim=cluster_dim,
359
+ n_clusters=4,
360
+ )
361
+ np.testing.assert_array_equal(
362
+ restricted.clustering.cluster_assignments.values,
363
+ reference.clustering.cluster_assignments.values,
364
+ )
365
+ reps = restricted.cluster_representatives
366
+ assert "price" in reps.coords["variable"].values
367
+ assert "price" in restricted.reconstructed.coords["variable"].values
368
+
369
+ def test_including_var_would_change_clustering(self, da_3var: xr.DataArray):
370
+ """Sanity: clustering on everything gives different assignments."""
371
+ restricted = tsam_xarray.aggregate(
372
+ da_3var,
373
+ time_dim="time",
374
+ cluster_dim="variable",
375
+ n_clusters=4,
376
+ cluster_on=["solar", "wind"],
377
+ )
378
+ full = tsam_xarray.aggregate(
379
+ da_3var, time_dim="time", cluster_dim="variable", n_clusters=4
380
+ )
381
+ assert (
382
+ restricted.clustering.cluster_assignments.values
383
+ != full.clustering.cluster_assignments.values
384
+ ).any()
385
+
386
+ def test_composes_with_weights(self, da_3var: xr.DataArray):
387
+ """cluster_on and weights coexist; weights apply to the active subset."""
388
+ result = tsam_xarray.aggregate(
389
+ da_3var,
390
+ time_dim="time",
391
+ cluster_dim="variable",
392
+ n_clusters=4,
393
+ cluster_on=["solar", "wind"],
394
+ weights={"solar": 5.0, "wind": 1.0},
395
+ )
396
+ assert result.n_clusters == 4
397
+ assert "price" in result.cluster_representatives.coords["variable"].values
398
+
399
+ def test_replace_extremes_rejected(self, da_3var: xr.DataArray):
400
+ """extremes='replace' can't survive the cluster_on transfer — error,
401
+ and steer users to weights instead of full exclusion."""
402
+ import tsam
403
+
404
+ with pytest.raises(ValueError, match=r"replace.*not supported.*cluster_on"):
405
+ tsam_xarray.aggregate(
406
+ da_3var,
407
+ time_dim="time",
408
+ cluster_dim="variable",
409
+ n_clusters=4,
410
+ cluster_on=["solar", "wind"],
411
+ extremes=tsam.ExtremeConfig(method="replace", max_value=["solar"]),
412
+ )
413
+
414
+ def test_extremes_on_excluded_column_rejected(self, da_3var: xr.DataArray):
415
+ """An ExtremeConfig referencing a cluster_on-excluded coord errors with
416
+ a clear cause, not tsam's misleading 'not found in data'."""
417
+ import tsam
418
+
419
+ with pytest.raises(ValueError, match=r"cluster_on excluded"):
420
+ tsam_xarray.aggregate(
421
+ da_3var,
422
+ time_dim="time",
423
+ cluster_dim="variable",
424
+ n_clusters=4,
425
+ cluster_on=["solar", "wind"],
426
+ extremes=tsam.ExtremeConfig(method="append", max_value=["price"]),
427
+ )
428
+
429
+ def test_append_extremes_allowed(self, da_3var: xr.DataArray):
430
+ """extremes='append' on an active column composes with cluster_on."""
431
+ import tsam
432
+
433
+ result = tsam_xarray.aggregate(
434
+ da_3var,
435
+ time_dim="time",
436
+ cluster_dim="variable",
437
+ n_clusters=4,
438
+ cluster_on=["solar", "wind"],
439
+ extremes=tsam.ExtremeConfig(method="append", max_value=["solar"]),
440
+ )
441
+ assert "price" in result.cluster_representatives.coords["variable"].values
442
+
443
+ @pytest.mark.parametrize(
444
+ ("cluster_dim", "cluster_on", "match"),
445
+ [
446
+ (["variable", "region"], ["solar"], "single cluster_dim"),
447
+ (["variable", "region"], {"scenario": ["low"]}, "unknown dims"),
448
+ ("variable", ["nuclear"], "unknown coords"),
449
+ ("variable", {"variable": []}, "nothing to cluster on"),
450
+ ],
451
+ ids=["list-multi", "unknown-dim", "unknown-coord", "empty"],
452
+ )
453
+ def test_invalid_cluster_on_raises(
454
+ self,
455
+ da_3var: xr.DataArray,
456
+ da_3var_region: xr.DataArray,
457
+ cluster_dim: list[str] | str,
458
+ cluster_on: object,
459
+ match: str,
460
+ ):
461
+ da = da_3var if cluster_dim == "variable" else da_3var_region
462
+ with pytest.raises(ValueError, match=match):
463
+ tsam_xarray.aggregate(
464
+ da,
465
+ time_dim="time",
466
+ cluster_dim=cluster_dim,
467
+ n_clusters=4,
468
+ cluster_on=cluster_on,
469
+ )
470
+
471
+
309
472
  class TestWeightTranslation:
310
473
  """Verify weight values are correctly translated."""
311
474
 
@@ -1,3 +0,0 @@
1
- {
2
- ".": "0.6.0"
3
- }
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes