tsam-xarray 0.6.0__tar.gz → 0.6.2__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.2/.release-please-manifest.json +3 -0
  2. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/CHANGELOG.md +20 -0
  3. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/PKG-INFO +1 -1
  4. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/examples/getting-started.ipynb +43 -9
  5. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/examples/multi-dim.ipynb +39 -1
  6. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_core.py +271 -28
  7. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_result.py +60 -2
  8. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_version.py +2 -2
  9. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/test/test_aggregate.py +287 -0
  10. tsam_xarray-0.6.0/.release-please-manifest.json +0 -3
  11. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/dependabot.yml +0 -0
  12. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/workflows/ci.yaml +0 -0
  13. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/workflows/dependabot-auto-merge.yaml +0 -0
  14. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/workflows/pr-title.yaml +0 -0
  15. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/workflows/publish.yaml +0 -0
  16. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.github/workflows/release.yaml +0 -0
  17. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.gitignore +0 -0
  18. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.pre-commit-config.yaml +0 -0
  19. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.readthedocs.yaml +0 -0
  20. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/.release-please-config.json +0 -0
  21. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/CLAUDE.md +0 -0
  22. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/LICENSE +0 -0
  23. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/README.md +0 -0
  24. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/codecov.yml +0 -0
  25. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/assets/multi-dim-input.png +0 -0
  26. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/assets/multi-dim-metrics.png +0 -0
  27. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/changelog.md +0 -0
  28. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/data-model.md +0 -0
  29. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/examples/clustering-io.ipynb +0 -0
  30. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/examples/segmentation.ipynb +0 -0
  31. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/examples/tuning.ipynb +0 -0
  32. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/gen_ref_pages.py +0 -0
  33. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/index.md +0 -0
  34. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/docs/stylesheets/extra.css +0 -0
  35. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/mkdocs.yml +0 -0
  36. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/pyproject.toml +0 -0
  37. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/scripts/generate_readme_images.py +0 -0
  38. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/__init__.py +0 -0
  39. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_clustering.py +0 -0
  40. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_sample_data.py +0 -0
  41. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/src/tsam_xarray/_tuning.py +0 -0
  42. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/test/conftest.py +0 -0
  43. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/test/test_parametrized.py +0 -0
  44. {tsam_xarray-0.6.0 → tsam_xarray-0.6.2}/test/test_tuning.py +0 -0
@@ -0,0 +1,3 @@
1
+ {
2
+ ".": "0.6.2"
3
+ }
@@ -1,5 +1,25 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.6.2](https://github.com/FBumann/tsam_xarray/compare/v0.6.1...v0.6.2) (2026-07-21)
4
+
5
+
6
+ ### Features
7
+
8
+ * first-class original-vs-reconstructed comparison on AggregationResult ([#93](https://github.com/FBumann/tsam_xarray/issues/93)) ([b3fd174](https://github.com/FBumann/tsam_xarray/commit/b3fd174cdbe379eb1de586aba5bfb75c0e4d21bb))
9
+
10
+
11
+ ### Bug Fixes
12
+
13
+ * align reconstructed coordinate order with input (+ comparison recipe) ([#97](https://github.com/FBumann/tsam_xarray/issues/97)) ([afd6818](https://github.com/FBumann/tsam_xarray/commit/afd6818715704822bacfad6edfa5c6595e9f8fc4))
14
+
15
+ ## [0.6.1](https://github.com/FBumann/tsam_xarray/compare/v0.6.0...v0.6.1) (2026-07-09)
16
+
17
+
18
+ ### Features
19
+
20
+ * 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))
21
+ * validate extremes options against cluster_on ([#91](https://github.com/FBumann/tsam_xarray/issues/91)) ([2fa5a91](https://github.com/FBumann/tsam_xarray/commit/2fa5a91c5222f8b6dd6932d9572dec62e90f595b))
22
+
3
23
  ## [0.6.0](https://github.com/FBumann/tsam_xarray/compare/v0.5.2...v0.6.0) (2026-05-27)
4
24
 
5
25
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tsam_xarray
3
- Version: 0.6.0
3
+ Version: 0.6.2
4
4
  Summary: Lightweight xarray wrapper for tsam time series aggregation
5
5
  License-Expression: MIT
6
6
  License-File: LICENSE
@@ -124,7 +124,13 @@
124
124
  "cell_type": "markdown",
125
125
  "metadata": {},
126
126
  "source": [
127
- "## Reconstructed vs original"
127
+ "## Reconstructed vs original\n",
128
+ "\n",
129
+ "`result.compare()` stacks `original` and `reconstructed` along a `variant`\n",
130
+ "dimension (on the original time axis), so the comparison plots directly with\n",
131
+ "`color=\"variant\"` — no manual `concat`/`melt`. Pass a coordinate selection\n",
132
+ "(e.g. `variable=\"solar\"`) to focus on one column. `result.original` and\n",
133
+ "`result.reconstructed` share the same dim order, so they always line up."
128
134
  ]
129
135
  },
130
136
  {
@@ -133,14 +139,8 @@
133
139
  "metadata": {},
134
140
  "outputs": [],
135
141
  "source": [
136
- "import xarray as xr\n",
137
- "\n",
138
- "comparison = xr.concat(\n",
139
- " [da_simple.sel(variable=\"solar\"), result.reconstructed.sel(variable=\"solar\")],\n",
140
- " dim=\"source\",\n",
141
- ").assign_coords(source=[\"original\", \"reconstructed\"])\n",
142
- "comparison.plotly.line(\n",
143
- " x=\"time\", color=\"source\", title=\"Original vs reconstructed (solar)\"\n",
142
+ "result.compare(variable=\"solar\").plotly.line(\n",
143
+ " x=\"time\", color=\"variant\", title=\"Original vs reconstructed (solar)\"\n",
144
144
  ")"
145
145
  ]
146
146
  },
@@ -153,6 +153,38 @@
153
153
  "result.residuals.plotly.line(x=\"time\", facet_col=\"variable\", title=\"Residuals\")"
154
154
  ]
155
155
  },
156
+ {
157
+ "cell_type": "markdown",
158
+ "metadata": {},
159
+ "source": [
160
+ "Because `to_dataframe()` returns tidy long-form data, richer comparisons are a\n",
161
+ "few lines in the plotting library of your choice — no bespoke helper needed.\n",
162
+ "For example, load-duration curves (each series sorted descending, Original\n",
163
+ "dotted / Reconstructed solid):"
164
+ ]
165
+ },
166
+ {
167
+ "cell_type": "code",
168
+ "execution_count": null,
169
+ "metadata": {},
170
+ "outputs": [],
171
+ "source": [
172
+ "tidy = result.to_dataframe()\n",
173
+ "tidy[\"rank\"] = (\n",
174
+ " tidy.groupby([\"variant\", \"variable\"])[\"energy\"]\n",
175
+ " .rank(method=\"first\", ascending=False)\n",
176
+ " .astype(int)\n",
177
+ ")\n",
178
+ "duration = tidy.set_index([\"rank\", \"variant\", \"variable\"])[\"energy\"].to_xarray()\n",
179
+ "duration.plotly.line(\n",
180
+ " x=\"rank\",\n",
181
+ " color=\"variable\",\n",
182
+ " line_dash=\"variant\",\n",
183
+ " line_dash_map={\"original\": \"dot\", \"reconstructed\": \"solid\"},\n",
184
+ " title=\"Load-duration curves: original vs reconstructed\",\n",
185
+ ")"
186
+ ]
187
+ },
156
188
  {
157
189
  "cell_type": "markdown",
158
190
  "metadata": {},
@@ -195,6 +227,8 @@
195
227
  "metadata": {},
196
228
  "outputs": [],
197
229
  "source": [
230
+ "import xarray as xr\n",
231
+ "\n",
198
232
  "disaggregated = result.disaggregate(result.cluster_representatives)\n",
199
233
  "\n",
200
234
  "comparison = xr.concat(\n",
@@ -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": {
@@ -3,6 +3,7 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  import itertools
6
+ import warnings
6
7
  from collections.abc import Hashable, Sequence
7
8
  from typing import Any
8
9
 
@@ -14,6 +15,7 @@ import xarray as xr
14
15
  from tsam_xarray._result import AccuracyMetrics, AggregationResult
15
16
 
16
17
  Weights = dict[str, float] | dict[str, dict[str, float]] | None
18
+ ClusterOn = str | Sequence[str] | dict[str, Sequence[str]] | None
17
19
 
18
20
 
19
21
  def aggregate(
@@ -23,6 +25,7 @@ def aggregate(
23
25
  cluster_dim: Sequence[str] | str,
24
26
  n_clusters: int,
25
27
  weights: Weights = None,
28
+ cluster_on: ClusterOn = None,
26
29
  **tsam_kwargs: Any,
27
30
  ) -> AggregationResult:
28
31
  """Aggregate an xarray DataArray using tsam.
@@ -55,6 +58,37 @@ def aggregate(
55
58
  e.g. ``("solar", "north")`` gets weight
56
59
  ``2.0 * 1.5 = 3.0``.
57
60
 
61
+ cluster_on: Restrict which coordinates drive the
62
+ clustering. Coordinates not selected are still
63
+ aggregated and reconstructed from the resulting
64
+ clusters, but have no influence on how the clusters
65
+ are formed. ``None`` (default) clusters on
66
+ everything. Mirrors the ``weights`` formats:
67
+
68
+ - **List** (single ``cluster_dim``)::
69
+
70
+ cluster_on=["solar", "wind"]
71
+
72
+ - **Dict** (multiple ``cluster_dim``)::
73
+
74
+ cluster_on={"variable": ["solar", "wind"]}
75
+
76
+ A dim omitted from the dict is unrestricted (all
77
+ its coordinates are eligible). When several dims
78
+ are listed, a column must match on *all* of them
79
+ to be clustered on. At least one column must
80
+ remain selected.
81
+
82
+ Not compatible with ``ExtremeConfig(method=
83
+ "replace")`` — the carried columns are filled by
84
+ transferring the clustering, which cannot
85
+ reproduce the hybrid 'replace' representative. Use
86
+ a low ``weights`` value to de-emphasise a column
87
+ instead of excluding it. An ``ExtremeConfig`` also
88
+ may not reference an excluded coordinate, since
89
+ extreme periods are identified only on the
90
+ clustered-on columns.
91
+
58
92
  **tsam_kwargs: Additional keyword arguments passed to
59
93
  ``tsam.aggregate()``.
60
94
  """
@@ -65,10 +99,18 @@ def aggregate(
65
99
  da = _validate_data(da, time_dim, col_dims, slice_dims)
66
100
  _validate_no_cluster_config_weights(tsam_kwargs)
67
101
  per_dim_weights = _normalize_weights(weights, da, col_dims)
102
+ active_coords = _normalize_cluster_on(cluster_on, da, col_dims)
103
+ _validate_extremes_with_cluster_on(tsam_kwargs, active_coords, da)
68
104
 
69
105
  if not slice_dims:
70
106
  return _aggregate_single(
71
- da, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
107
+ da,
108
+ n_clusters,
109
+ time_dim,
110
+ col_dims,
111
+ per_dim_weights,
112
+ active_coords,
113
+ tsam_kwargs,
72
114
  )
73
115
 
74
116
  slice_coords = {d: da.coords[d].values for d in slice_dims}
@@ -80,7 +122,13 @@ def aggregate(
80
122
  sel = dict(zip(slice_dims, key, strict=True))
81
123
  da_slice = da.sel(sel)
82
124
  r = _aggregate_single(
83
- da_slice, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
125
+ da_slice,
126
+ n_clusters,
127
+ time_dim,
128
+ col_dims,
129
+ per_dim_weights,
130
+ active_coords,
131
+ tsam_kwargs,
84
132
  )
85
133
  results.append(r)
86
134
 
@@ -129,6 +177,68 @@ def _validate_no_cluster_config_weights(
129
177
  raise ValueError(msg)
130
178
 
131
179
 
180
+ def _validate_extremes_with_cluster_on(
181
+ tsam_kwargs: dict[str, Any],
182
+ active_coords: dict[str, set[str]] | None,
183
+ da: xr.DataArray,
184
+ ) -> None:
185
+ """Reject extremes settings incompatible with the cluster_on transfer.
186
+
187
+ cluster_on clusters on the active subset and transfers the result to the
188
+ carried columns via ``ClusteringResult.apply()``. Two extremes settings
189
+ don't survive that transfer:
190
+
191
+ - ``method="replace"`` builds a hybrid representative (some columns from the
192
+ medoid, some from the extreme period) that apply() cannot reproduce — tsam
193
+ silently falls back to the medoid. Rather than degrade quietly, error and
194
+ point to ``weights``, which de-emphasises a column without excluding it
195
+ (tsam clamps a low weight to a minimal one, so the column stays in the
196
+ single aggregation pass and 'replace' remains reproducible).
197
+ - An ExtremeConfig that references a coordinate cluster_on excluded: those
198
+ columns never enter the clustering, so tsam can't identify extremes on
199
+ them (it would raise a misleading "not found in data"). Surface the real
200
+ cause instead.
201
+ """
202
+ if active_coords is None:
203
+ return
204
+ extremes = tsam_kwargs.get("extremes")
205
+ if extremes is None:
206
+ return
207
+
208
+ if getattr(extremes, "method", None) == "replace":
209
+ msg = (
210
+ "extremes method 'replace' is not supported together with "
211
+ "cluster_on. cluster_on fully excludes the carried columns and "
212
+ "transfers the clustering to them, which cannot reproduce the "
213
+ "hybrid 'replace' representative. To reduce a column's influence "
214
+ "without excluding it, give it a low 'weights' value instead of "
215
+ "using cluster_on (weights never removes a column entirely, so "
216
+ "'replace' stays reproducible)."
217
+ )
218
+ raise ValueError(msg)
219
+
220
+ passive_coords: set[str] = set()
221
+ for dim, active_set in active_coords.items():
222
+ all_coords = {str(c) for c in da.coords[dim].values}
223
+ passive_coords |= all_coords - active_set
224
+
225
+ referenced = (
226
+ set(extremes.max_value)
227
+ | set(extremes.min_value)
228
+ | set(extremes.max_period)
229
+ | set(extremes.min_period)
230
+ )
231
+ excluded = referenced & passive_coords
232
+ if excluded:
233
+ msg = (
234
+ f"extremes references coordinates {excluded} that cluster_on "
235
+ "excluded from the clustering. Extreme periods can only be "
236
+ "identified on clustered-on columns — add these coordinates to "
237
+ "cluster_on, or remove them from the ExtremeConfig."
238
+ )
239
+ raise ValueError(msg)
240
+
241
+
132
242
  def _validate(
133
243
  da: xr.DataArray,
134
244
  time_dim: str,
@@ -171,8 +281,6 @@ def _validate_data(
171
281
 
172
282
  # Dask arrays — compute before other checks
173
283
  if hasattr(da.data, "dask"):
174
- import warnings
175
-
176
284
  warnings.warn(
177
285
  "DataArray is backed by dask. Computing eagerly for tsam.",
178
286
  stacklevel=3,
@@ -265,7 +373,7 @@ def _representatives_to_da(
265
373
  df = df.copy()
266
374
  # With segmentation, index has 3 levels: (cluster, segment_step, segment_duration)
267
375
  # Without: 2 levels: (cluster, timestep)
268
- if df.index.nlevels == 3:
376
+ if isinstance(df.index, pd.MultiIndex) and df.index.nlevels == 3:
269
377
  df.index = df.index.droplevel(2) # drop segment_duration
270
378
  df.index.names = ["cluster", "timestep"]
271
379
 
@@ -339,6 +447,36 @@ def _metric_to_da(
339
447
  return xr.DataArray(series.to_xarray())
340
448
 
341
449
 
450
+ def _validate_dim_coord_mapping(
451
+ mapping: dict[str, Any],
452
+ da: xr.DataArray,
453
+ col_dims: list[str],
454
+ label: str,
455
+ ) -> None:
456
+ """Validate a ``{dim: coords}`` mapping references real dims and coords.
457
+
458
+ ``coords`` may be any iterable of coordinate names (a list, or the keys
459
+ of a ``{coord: weight}`` dict). Shared by ``weights`` and ``cluster_on``.
460
+ """
461
+ extra_dims = set(mapping.keys()) - set(col_dims)
462
+ if extra_dims:
463
+ msg = (
464
+ f"{label} has unknown dims {extra_dims}, "
465
+ f"must be subset of cluster_dim {col_dims}"
466
+ )
467
+ raise ValueError(msg)
468
+
469
+ for dim_name, coords in mapping.items():
470
+ valid_coords = set(str(c) for c in da.coords[dim_name].values)
471
+ unknown = set(str(c) for c in coords) - valid_coords
472
+ if unknown:
473
+ msg = (
474
+ f"{label} has unknown coords {unknown} for dim {dim_name!r}, "
475
+ f"valid coords: {sorted(valid_coords)}"
476
+ )
477
+ raise ValueError(msg)
478
+
479
+
342
480
  def _normalize_weights(
343
481
  weights: dict[str, float] | dict[str, dict[str, float]] | None,
344
482
  da: xr.DataArray,
@@ -370,27 +508,51 @@ def _normalize_weights(
370
508
  raise ValueError(msg)
371
509
  per_dim_weights = {col_dims[0]: weights} # type: ignore[dict-item]
372
510
 
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)
511
+ _validate_dim_coord_mapping(per_dim_weights, da, col_dims, "weights")
512
+ return per_dim_weights
381
513
 
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:
514
+
515
+ def _normalize_cluster_on(
516
+ cluster_on: ClusterOn,
517
+ da: xr.DataArray,
518
+ col_dims: list[str],
519
+ ) -> dict[str, set[str]] | None:
520
+ """Normalize cluster_on to a per-dim set of eligible coords.
521
+
522
+ A list/str form requires a single cluster_dim. Returns ``None``
523
+ when nothing is restricted (cluster on everything).
524
+ """
525
+ if cluster_on is None:
526
+ return None
527
+
528
+ if isinstance(cluster_on, dict):
529
+ per_dim: dict[str, list[str]] = {
530
+ k: [str(c) for c in v] for k, v in cluster_on.items()
531
+ }
532
+ else:
533
+ if len(col_dims) != 1:
387
534
  msg = (
388
- f"weights has unknown coords {unknown} for dim {dim_name!r}, "
389
- f"valid coords: {sorted(valid_coords)}"
535
+ "List-form cluster_on requires a single cluster_dim. "
536
+ "For multiple cluster_dim, use a dict: "
537
+ '{"dim_name": ["coord", ...]}.'
390
538
  )
391
539
  raise ValueError(msg)
540
+ coords = [cluster_on] if isinstance(cluster_on, str) else list(cluster_on)
541
+ per_dim = {col_dims[0]: [str(c) for c in coords]}
392
542
 
393
- return per_dim_weights
543
+ _validate_dim_coord_mapping(per_dim, da, col_dims, "cluster_on")
544
+
545
+ # An empty coord list for any restricted dim leaves no column to cluster
546
+ # on (the stacked columns are the full coordinate grid). Fail here, before
547
+ # any aggregation runs, rather than deep in the per-slice loop.
548
+ if any(not coords for coords in per_dim.values()):
549
+ msg = (
550
+ "cluster_on selected no columns — nothing to cluster on. "
551
+ "Select at least one coordinate present in the data."
552
+ )
553
+ raise ValueError(msg)
554
+
555
+ return {d: set(cs) for d, cs in per_dim.items()}
394
556
 
395
557
 
396
558
  def _translate_weights(
@@ -414,30 +576,111 @@ def _translate_weights(
414
576
  return flat
415
577
 
416
578
 
579
+ # tsam.aggregate() kwargs that ClusteringResult.apply() also accepts. Only
580
+ # these are forwarded to apply() when re-aggregating passive columns onto a
581
+ # clustering built from the active subset; the rest (period_duration, cluster,
582
+ # segments, extremes, ...) are already baked into the stored clustering.
583
+ _APPLY_KWARGS = frozenset(
584
+ {"temporal_resolution", "round_decimals", "numerical_tolerance"}
585
+ )
586
+
587
+
588
+ def _active_columns(
589
+ df: pd.DataFrame,
590
+ active_coords: dict[str, set[str]] | None,
591
+ col_dims: list[str],
592
+ ) -> list[Hashable]:
593
+ """Columns that drive the clustering, given a cluster_on selection.
594
+
595
+ A column is active only if, for every restricted dim, its coordinate is
596
+ in that dim's eligible set; unrestricted dims impose no constraint. With
597
+ no restriction (``active_coords is None``) every column is active. The
598
+ remaining columns are *passive*: still aggregated and reconstructed, but
599
+ excluded from the cluster distances.
600
+ """
601
+ if active_coords is None:
602
+ return list(df.columns)
603
+
604
+ def is_active(col: Hashable) -> bool:
605
+ pairs = (
606
+ zip(col_dims, col, strict=True)
607
+ if isinstance(col, tuple)
608
+ else [(col_dims[0], col)]
609
+ )
610
+ return all(
611
+ str(coord) in active_coords[dim]
612
+ for dim, coord in pairs
613
+ if dim in active_coords
614
+ )
615
+
616
+ return [col for col in df.columns if is_active(col)]
617
+
618
+
417
619
  def _aggregate_single(
418
620
  da: xr.DataArray,
419
621
  n_clusters: int,
420
622
  time_dim: str,
421
623
  col_dims: list[str],
422
624
  weights: dict[str, dict[str, float]] | None,
625
+ active_coords: dict[str, set[str]] | None,
423
626
  tsam_kwargs: dict[str, Any],
424
627
  ) -> AggregationResult:
425
- """Run a single tsam aggregation on a DataArray."""
628
+ """Run a single tsam aggregation on a DataArray.
629
+
630
+ When ``cluster_on`` restricts the clustering, the clustering runs on
631
+ the active subset only, then is applied back to the full frame so the
632
+ passive columns still get representatives and reconstruction.
633
+ """
426
634
  df = _to_dataframe(da, time_dim, col_dims)
427
635
 
428
636
  tsam_weights: dict[Hashable, float] | None = None
429
637
  if weights is not None:
430
638
  tsam_weights = _translate_weights(weights, df, col_dims)
431
639
 
432
- tsam_result = tsam.aggregate(
433
- df,
434
- n_clusters,
435
- weights=tsam_weights, # type: ignore[arg-type]
436
- **tsam_kwargs,
437
- )
640
+ active_cols = _active_columns(df, active_coords, col_dims)
438
641
 
642
+ with warnings.catch_warnings():
643
+ warnings.filterwarnings(
644
+ "ignore",
645
+ message=".*sorted alphabetically.*",
646
+ category=FutureWarning,
647
+ )
648
+ if len(active_cols) == len(df.columns):
649
+ tsam_result = tsam.aggregate(
650
+ df,
651
+ n_clusters,
652
+ weights=tsam_weights, # type: ignore[arg-type]
653
+ **tsam_kwargs,
654
+ )
655
+ else:
656
+ active_weights = (
657
+ None
658
+ if tsam_weights is None
659
+ else {c: tsam_weights[c] for c in active_cols}
660
+ )
661
+ clustering = tsam.aggregate(
662
+ df[active_cols],
663
+ n_clusters,
664
+ weights=active_weights, # type: ignore[arg-type]
665
+ **tsam_kwargs,
666
+ ).clustering
667
+ apply_kwargs = {k: v for k, v in tsam_kwargs.items() if k in _APPLY_KWARGS}
668
+ tsam_result = clustering.apply(df, **apply_kwargs)
669
+
670
+ return _result_from_tsam(tsam_result, da, df, time_dim, col_dims)
671
+
672
+
673
+ def _result_from_tsam(
674
+ tsam_result: Any,
675
+ da: xr.DataArray,
676
+ df: pd.DataFrame,
677
+ time_dim: str,
678
+ col_dims: list[str],
679
+ ) -> AggregationResult:
680
+ """Build an AggregationResult from a tsam aggregation result."""
439
681
  typical = _representatives_to_da(tsam_result.cluster_representatives, col_dims)
440
682
  reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims)
683
+ reconstructed = reconstructed.transpose(*da.dims).reindex_like(da)
441
684
 
442
685
  cw = tsam_result.cluster_weights
443
686
  cluster_ids = np.array(sorted(cw.keys()))
@@ -5,6 +5,7 @@ from __future__ import annotations
5
5
  from dataclasses import dataclass
6
6
  from typing import TYPE_CHECKING
7
7
 
8
+ import pandas as pd
8
9
  import xarray as xr
9
10
 
10
11
  if TYPE_CHECKING:
@@ -71,7 +72,7 @@ class AggregationResult:
71
72
  *slice_dims)``.
72
73
  accuracy: Per-column and weighted accuracy metrics.
73
74
  reconstructed: Reconstructed time series
74
- (same shape as input).
75
+ (same shape and dim order as ``original``).
75
76
  original: The input data.
76
77
  clustering: Reusable clustering metadata.
77
78
  See `ClusteringResult`.
@@ -121,9 +122,66 @@ class AggregationResult:
121
122
 
122
123
  @property
123
124
  def residuals(self) -> xr.DataArray:
124
- """Difference between original and reconstructed data."""
125
+ """Difference between original and reconstructed data.
126
+
127
+ Shares the dim order of ``original`` and ``reconstructed``.
128
+ """
125
129
  return self.original - self.reconstructed
126
130
 
131
+ def compare(self, **sel: object) -> xr.DataArray:
132
+ """Stack ``original`` and ``reconstructed`` along a ``variant`` dim.
133
+
134
+ Returns a single DataArray on the original time axis with a new
135
+ ``variant`` coordinate ``["original", "reconstructed"]``, ready to
136
+ plot directly with a ``color=``/``hue="variant"`` grouping — no
137
+ ``melt`` step. This is the canonical way to eyeball aggregation
138
+ quality per column and per slice dim.
139
+
140
+ Args:
141
+ **sel: Optional label-based selection applied to both arrays
142
+ before stacking, e.g. ``compare(variable="solar")`` to
143
+ compare a single column.
144
+
145
+ Returns:
146
+ DataArray with dims ``("variant", *original.dims)``.
147
+
148
+ Examples:
149
+ >>> agg.compare(variable="solar").plotly.line(
150
+ ... x="time", color="variant"
151
+ ... )
152
+ """
153
+ original = self.original
154
+ reconstructed = self.reconstructed
155
+ if sel:
156
+ original = original.sel(sel)
157
+ reconstructed = reconstructed.sel(sel)
158
+ variant = pd.Index(["original", "reconstructed"], name="variant")
159
+ combined = xr.concat([original, reconstructed], dim=variant)
160
+ combined.name = self.original.name
161
+ return combined
162
+
163
+ def to_dataframe(self, **sel: object) -> pd.DataFrame:
164
+ """Tidy/long-form ``original`` vs ``reconstructed`` DataFrame.
165
+
166
+ A flat DataFrame with a ``variant`` column
167
+ (``"original"``/``"reconstructed"``), the ``time`` axis, every
168
+ cluster and slice dim, and a value column — ready to hand
169
+ straight to a plotting library.
170
+
171
+ Args:
172
+ **sel: Optional label-based selection forwarded to
173
+ `compare` (e.g. ``variable="solar"``).
174
+
175
+ Returns:
176
+ Long-form DataFrame with a ``variant`` column and one value
177
+ column (named after the input DataArray, or ``"value"``).
178
+ """
179
+ combined = self.compare(**sel)
180
+ name = combined.name
181
+ if name is None or str(name) == "variant":
182
+ name = "value"
183
+ return combined.to_dataframe(name=str(name)).reset_index()
184
+
127
185
  def disaggregate(self, data: xr.DataArray) -> xr.DataArray:
128
186
  """Map data on ``(cluster, timestep)`` back to original time.
129
187
 
@@ -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.2'
22
+ __version_tuple__ = version_tuple = (0, 6, 2)
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
 
@@ -1366,3 +1529,127 @@ class TestSliceEdgeCases:
1366
1529
  _validate_consistent_cluster_counts(
1367
1530
  [r1, r2], [("low",), ("high",)]
1368
1531
  ) # should not raise
1532
+
1533
+
1534
+ def _make_time_last_da() -> xr.DataArray:
1535
+ """DataArray with time as the LAST dim (the #92 papercut case)."""
1536
+ time = pd.date_range("2020-01-01", periods=10 * 24, freq="h")
1537
+ rng = np.random.default_rng(0)
1538
+ return xr.DataArray(
1539
+ rng.random((2, 2, len(time))),
1540
+ dims=["scenario", "variable", "time"],
1541
+ coords={
1542
+ "scenario": ["a", "b"],
1543
+ "variable": ["solar", "wind"],
1544
+ "time": time,
1545
+ },
1546
+ name="power",
1547
+ )
1548
+
1549
+
1550
+ class TestDimOrderConsistency:
1551
+ """original / reconstructed / residuals must share dim order (#92)."""
1552
+
1553
+ def test_time_last_input_aligned(self):
1554
+ da = _make_time_last_da()
1555
+ result = tsam_xarray.aggregate(
1556
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1557
+ )
1558
+ assert result.original.dims == da.dims
1559
+ assert result.reconstructed.dims == result.original.dims
1560
+ assert result.residuals.dims == result.original.dims
1561
+
1562
+ def test_time_first_input_aligned(self):
1563
+ da = _make_da(scenarios=["low", "high"])
1564
+ result = tsam_xarray.aggregate(
1565
+ da,
1566
+ time_dim="time",
1567
+ cluster_dim=["variable", "region"],
1568
+ n_clusters=4,
1569
+ )
1570
+ assert result.reconstructed.dims == result.original.dims
1571
+ assert result.residuals.dims == result.original.dims
1572
+
1573
+ def test_single_dim_no_slice(self):
1574
+ da = _make_time_last_da().isel(scenario=0).drop_vars("scenario")
1575
+ result = tsam_xarray.aggregate(
1576
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1577
+ )
1578
+ assert result.reconstructed.dims == result.original.dims == da.dims
1579
+
1580
+ def test_non_alphabetical_coord_order_preserved(self):
1581
+ """reconstructed keeps the input's coord order, not tsam's sorted one."""
1582
+ time = pd.date_range("2020-01-01", periods=5 * 24, freq="h")
1583
+ rng = np.random.default_rng(0)
1584
+ da = xr.DataArray(
1585
+ rng.random((len(time), 3)),
1586
+ dims=["time", "variable"],
1587
+ coords={"time": time, "variable": ["wind", "solar", "demand"]},
1588
+ name="p",
1589
+ )
1590
+ result = tsam_xarray.aggregate(
1591
+ da, time_dim="time", cluster_dim="variable", n_clusters=3
1592
+ )
1593
+ assert list(result.reconstructed["variable"].values) == [
1594
+ "wind",
1595
+ "solar",
1596
+ "demand",
1597
+ ]
1598
+ assert not bool(result.reconstructed.isnull().any())
1599
+
1600
+
1601
+ class TestCompare:
1602
+ """AggregationResult.compare() and .to_dataframe() (#92)."""
1603
+
1604
+ def test_compare_adds_variant_dim(self):
1605
+ da = _make_time_last_da()
1606
+ result = tsam_xarray.aggregate(
1607
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1608
+ )
1609
+ cmp = result.compare()
1610
+ assert cmp.dims == ("variant", *result.original.dims)
1611
+ assert list(cmp["variant"].values) == ["original", "reconstructed"]
1612
+ assert cmp.name == "power"
1613
+ xr.testing.assert_equal(cmp.sel(variant="original", drop=True), result.original)
1614
+ xr.testing.assert_equal(
1615
+ cmp.sel(variant="reconstructed", drop=True), result.reconstructed
1616
+ )
1617
+
1618
+ def test_compare_with_selection(self):
1619
+ da = _make_time_last_da()
1620
+ result = tsam_xarray.aggregate(
1621
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1622
+ )
1623
+ cmp = result.compare(variable="solar")
1624
+ assert "variable" not in cmp.dims
1625
+ assert cmp.sizes["variant"] == 2
1626
+
1627
+ def test_to_dataframe_is_long_form(self):
1628
+ da = _make_time_last_da()
1629
+ result = tsam_xarray.aggregate(
1630
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1631
+ )
1632
+ df = result.to_dataframe(variable="solar")
1633
+ assert "variant" in df.columns
1634
+ assert "power" in df.columns
1635
+ assert set(df["variant"].unique()) == {"original", "reconstructed"}
1636
+ n_time = da.sizes["time"]
1637
+ assert len(df) == 2 * da.sizes["scenario"] * n_time
1638
+
1639
+ def test_to_dataframe_unnamed_uses_value(self):
1640
+ da = _make_time_last_da().rename(None)
1641
+ result = tsam_xarray.aggregate(
1642
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1643
+ )
1644
+ df = result.to_dataframe()
1645
+ assert "value" in df.columns
1646
+
1647
+ def test_to_dataframe_name_collision_with_variant(self):
1648
+ """An input named 'variant' falls back to 'value' (no insert clash)."""
1649
+ da = _make_time_last_da().rename("variant")
1650
+ result = tsam_xarray.aggregate(
1651
+ da, time_dim="time", cluster_dim="variable", n_clusters=4
1652
+ )
1653
+ df = result.to_dataframe(variable="solar")
1654
+ assert "value" in df.columns
1655
+ assert list(df.columns).count("variant") == 1
@@ -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