From cdae729d15252bcc56f3244ce45363df01ce78dd Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 23 Jul 2026 11:56:25 +0200 Subject: [PATCH] feat(extremes): force count-stable clustering and allow transferable replace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consume the tsam fix in FZJ-IEK3-VSA/tsam#414 (closes upstream #410) so extreme periods compose with the wrapper's fan-out and transfer paths. - Always force ExtremeConfig(preserve_n_clusters=True) for the additive methods (append/new_cluster) so every independently-aggregated slice yields exactly n_clusters and results stack rectangularly. Extremes now count against the cluster budget instead of adding a data-dependent number of extra clusters. - Lift the cluster_on + method="replace" rejection when the installed tsam can transfer the hybrid representative (stores extreme_replacements). - Both behaviours are feature-detected: no-ops on released tsam (which lacks the flag), so existing behaviour and CI are unchanged. They activate automatically once tsam_xarray runs on a tsam that ships the fix — which requires the v4 pipeline architecture the fix is built on. Adds test/test_extremes_composability.py: a 30-case fuzz over additive extremes on sliced data (asserts the fallback count contract today, tightens to exact-count once the flag is present) plus flag-gated count-stability and replace-transfer tests. Co-Authored-By: Claude Opus 4.8 --- src/tsam_xarray/_core.py | 125 ++++++++++++++---- test/test_aggregate.py | 11 +- test/test_extremes_composability.py | 196 ++++++++++++++++++++++++++++ 3 files changed, 302 insertions(+), 30 deletions(-) create mode 100644 test/test_extremes_composability.py diff --git a/src/tsam_xarray/_core.py b/src/tsam_xarray/_core.py index f6f25aa..600fb8c 100644 --- a/src/tsam_xarray/_core.py +++ b/src/tsam_xarray/_core.py @@ -2,6 +2,7 @@ from __future__ import annotations +import dataclasses import itertools import warnings from collections.abc import Hashable, Sequence @@ -33,6 +34,54 @@ def _cluster_counts(tsam_result: Any) -> dict[int, float]: return cast("dict[int, float]", counts) +def _extreme_config_supports_preserve() -> bool: + """True when the installed tsam's ``ExtremeConfig`` has ``preserve_n_clusters``. + + Added in the tsam release that fixes FZJ-IEK3-VSA/tsam#410. On older tsam the + field is absent and the wrapper falls back to today's behaviour, where extreme + counts may vary across slices. + """ + return any( + f.name == "preserve_n_clusters" for f in dataclasses.fields(tsam.ExtremeConfig) + ) + + +def _clustering_supports_replace_transfer() -> bool: + """True when the installed tsam stores replace injections for transfer. + + Added alongside ``preserve_n_clusters`` (FZJ-IEK3-VSA/tsam#410). When present, + ``ClusteringResult.apply()`` reproduces the hybrid 'replace' representative, so + the ``cluster_on`` transfer path is safe with ``method="replace"``. + """ + return any( + f.name == "extreme_replacements" + for f in dataclasses.fields(tsam.ClusteringResult) + ) + + +def _force_preserve_n_clusters(tsam_kwargs: dict[str, Any]) -> None: + """Pin the extreme-period count so every slice yields exactly ``n_clusters``. + + tsam adds a data-dependent number of extreme clusters — criteria that resolve + to the same period, or an extreme already coinciding with a cluster center, get + deduplicated — so independently-aggregated slices can end up with different + cluster counts and fail to stack. tsam's ``preserve_n_clusters=True`` clusters + into ``n_clusters - D`` and adds the ``D`` extremes back, making the total + exactly ``n_clusters`` regardless of the data. The wrapper always turns it on + for the additive methods so its rectangular output holds; ``method="replace"`` + is already count-stable (and the flag would only warn there), so it is left + untouched. No-op on tsam versions without the flag. + """ + extremes = tsam_kwargs.get("extremes") + if extremes is None or not _extreme_config_supports_preserve(): + return + if getattr(extremes, "method", "append") == "replace": + return + if getattr(extremes, "preserve_n_clusters", False): + return + tsam_kwargs["extremes"] = dataclasses.replace(extremes, preserve_n_clusters=True) + + def aggregate( da: xr.DataArray, *, @@ -95,15 +144,14 @@ def aggregate( to be clustered on. At least one column must remain selected. - Not compatible with ``ExtremeConfig(method= - "replace")`` — the carried columns are filled by - transferring the clustering, which cannot - reproduce the hybrid 'replace' representative. Use - a low ``weights`` value to de-emphasise a column - instead of excluding it. An ``ExtremeConfig`` also - may not reference an excluded coordinate, since - extreme periods are identified only on the - clustered-on columns. + ``ExtremeConfig(method="replace")`` needs a tsam + that can transfer the hybrid 'replace' + representative (FZJ-IEK3-VSA/tsam#410); on older + tsam it is rejected — use a low ``weights`` value + to de-emphasise a column instead of excluding it. + An ``ExtremeConfig`` may not reference an excluded + coordinate, since extreme periods are identified + only on the clustered-on columns. dim_names: Names for the structural output dimensions (``cluster``, ``timestep``, ``period``, ``segment``). @@ -113,6 +161,16 @@ def aggregate( **tsam_kwargs: Additional keyword arguments passed to ``tsam.aggregate()``. + + When ``extremes`` uses an additive method + (``"append"`` / ``"new_cluster"``), the wrapper + forces ``preserve_n_clusters=True`` so every slice + yields exactly ``n_clusters`` representatives and the + output stays rectangular. Extremes count against the + cluster budget rather than adding to it. Requires a + tsam with the flag (FZJ-IEK3-VSA/tsam#410); on older + tsam the count is data-dependent and mismatched + slices raise. """ resolved_dim_names = dim_names if dim_names is not None else DimNames() _validate_time_dim(da, time_dim) @@ -125,6 +183,7 @@ def aggregate( per_dim_weights = _normalize_weights(weights, da, col_dims) active_coords = _normalize_cluster_on(cluster_on, da, col_dims) _validate_extremes_with_cluster_on(tsam_kwargs, active_coords, da) + _force_preserve_n_clusters(tsam_kwargs) if not slice_dims: return _aggregate_single( @@ -220,15 +279,17 @@ def _validate_extremes_with_cluster_on( """Reject extremes settings incompatible with the cluster_on transfer. cluster_on clusters on the active subset and transfers the result to the - carried columns via ``ClusteringResult.apply()``. Two extremes settings - don't survive that transfer: + carried columns via ``ClusteringResult.apply()``. Two extremes settings can + fail to survive that transfer: - ``method="replace"`` builds a hybrid representative (some columns from the - medoid, some from the extreme period) that apply() cannot reproduce — tsam - silently falls back to the medoid. Rather than degrade quietly, error and - point to ``weights``, which de-emphasises a column without excluding it - (tsam clamps a low weight to a minimal one, so the column stays in the - single aggregation pass and 'replace' remains reproducible). + medoid, some from the extreme period). On tsam versions that store the + injection (FZJ-IEK3-VSA/tsam#410), apply() replays it and the transfer is + exact, so 'replace' is allowed. On older tsam, apply() silently falls back + to the medoid — rather than degrade quietly, error and point to ``weights``, + which de-emphasises a column without excluding it (tsam clamps a low weight + to a minimal one, so the column stays in the single aggregation pass and + 'replace' remains reproducible). - An ExtremeConfig that references a coordinate cluster_on excluded: those columns never enter the clustering, so tsam can't identify extremes on them (it would raise a misleading "not found in data"). Surface the real @@ -240,15 +301,20 @@ def _validate_extremes_with_cluster_on( if extremes is None: return - if getattr(extremes, "method", None) == "replace": + if ( + getattr(extremes, "method", None) == "replace" + and not _clustering_supports_replace_transfer() + ): msg = ( "extremes method 'replace' is not supported together with " - "cluster_on. cluster_on fully excludes the carried columns and " - "transfers the clustering to them, which cannot reproduce the " - "hybrid 'replace' representative. To reduce a column's influence " - "without excluding it, give it a low 'weights' value instead of " - "using cluster_on (weights never removes a column entirely, so " - "'replace' stays reproducible)." + "cluster_on on this tsam version. cluster_on fully excludes the " + "carried columns and transfers the clustering to them, which this " + "tsam cannot reproduce for the hybrid 'replace' representative. " + "Upgrade tsam to a version with transferable 'replace' " + "(FZJ-IEK3-VSA/tsam#410), or reduce a column's influence without " + "excluding it by giving it a low 'weights' value instead of using " + "cluster_on (weights never removes a column entirely, so 'replace' " + "stays reproducible)." ) raise ValueError(msg) @@ -371,12 +437,13 @@ def _validate_consistent_cluster_counts( if len(unique) > 1: msg = ( "Slices produced different cluster counts: " - f"{counts}. This can happen with " - "ExtremeConfig(method='append'). Use " - "method='replace', aggregate slices separately, " - "or open an issue at " - "github.com/FBumann/tsam_xarray/issues " - "to discuss expected behaviour." + f"{counts}. This happens with additive extremes " + "(ExtremeConfig(method='append' / 'new_cluster')) on " + "a tsam without preserve_n_clusters. Upgrade tsam to " + "a version with count-stable extremes " + "(FZJ-IEK3-VSA/tsam#410), which the wrapper enables " + "automatically, or use method='replace' or aggregate " + "slices separately." ) raise ValueError(msg) diff --git a/test/test_aggregate.py b/test/test_aggregate.py index bf912b4..bdedc3f 100644 --- a/test/test_aggregate.py +++ b/test/test_aggregate.py @@ -10,6 +10,7 @@ import xarray as xr import tsam_xarray +from tsam_xarray._core import _clustering_supports_replace_transfer def _make_da( @@ -410,9 +411,17 @@ def test_composes_with_weights(self, da_3var: xr.DataArray): assert result.n_clusters == 4 assert "price" in result.cluster_representatives.coords["variable"].values + @pytest.mark.skipif( + _clustering_supports_replace_transfer(), + reason="tsam transfers replace (#410) — cluster_on+replace is allowed", + ) def test_replace_extremes_rejected(self, da_3var: xr.DataArray): """extremes='replace' can't survive the cluster_on transfer — error, - and steer users to weights instead of full exclusion.""" + and steer users to weights instead of full exclusion. + + Only on tsam that cannot transfer the hybrid 'replace' representative; + newer tsam reproduces it, so the rejection is lifted (covered by + ``test_extremes_composability.TestReplaceTransfers``).""" import tsam with pytest.raises(ValueError, match=r"replace.*not supported.*cluster_on"): diff --git a/test/test_extremes_composability.py b/test/test_extremes_composability.py new file mode 100644 index 0000000..d844057 --- /dev/null +++ b/test/test_extremes_composability.py @@ -0,0 +1,196 @@ +"""Composability + fuzz tests for extreme-period handling across slices. + +Two guarantees are exercised here, both riding on the tsam fix in +FZJ-IEK3-VSA/tsam#410: + +* **Count stability** — with additive extremes (``append`` / ``new_cluster``), + every independently-aggregated slice must yield exactly ``n_clusters`` so the + results stack into a rectangular array. The wrapper forces + ``preserve_n_clusters=True`` to achieve this. Where the running tsam lacks the + flag, the tests instead assert the documented fallback contract: aggregate() + either succeeds with *uniform* counts or raises the specific count-mismatch + error — never anything else. +* **Transferable replace** — ``method="replace"`` composes with ``cluster_on`` + and round-trips through ``apply()``. Gated on the tsam that stores the + injection. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import tsam +import xarray as xr + +import tsam_xarray +from tsam_xarray._core import ( + _clustering_supports_replace_transfer, + _extreme_config_supports_preserve, +) + +SUPPORTS_PRESERVE = _extreme_config_supports_preserve() +SUPPORTS_REPLACE_TRANSFER = _clustering_supports_replace_transfer() + +VARIABLES = ["solar", "wind", "price"] + + +def _sliced_da(seed: int, n_scenarios: int = 3, n_days: int = 20) -> xr.DataArray: + """A (time, variable, scenario) array — scenario is an independent slice.""" + rng = np.random.default_rng(seed) + time = pd.date_range("2020-01-01", periods=n_days * 24, freq="h") + return xr.DataArray( + rng.random((len(time), len(VARIABLES), n_scenarios)), + dims=["time", "variable", "scenario"], + coords={ + "time": time, + "variable": VARIABLES, + "scenario": [f"s{i}" for i in range(n_scenarios)], + }, + ) + + +@pytest.mark.parametrize("seed", range(30)) +def test_additive_extremes_across_slices_fuzz(seed: int) -> None: + """Fuzz additive extremes over sliced data; assert the count contract. + + Criteria are capped so the distinct-extreme count D stays below + ``n_clusters`` (tsam rejects ``n_clusters <= D`` under preserve). + """ + rng = np.random.default_rng(10_000 + seed) + method = str(rng.choice(["append", "new_cluster"])) + + def pick() -> list[str]: + k = int(rng.integers(0, 3)) # 0-2 columns per criterion + if k == 0: + return [] + return list(rng.choice(VARIABLES, size=k, replace=False)) + + max_value, min_value = pick(), pick() + if not (max_value or min_value): # ensure at least one active criterion + max_value = ["solar"] + + extremes = tsam.ExtremeConfig( + method=method, max_value=max_value, min_value=min_value + ) + n_clusters = int(rng.integers(8, 12)) + da = _sliced_da(seed) + + try: + res = tsam_xarray.aggregate( + da, + time_dim="time", + cluster_dim="variable", + n_clusters=n_clusters, + extremes=extremes, + ) + except ValueError as exc: + # The only tolerated failure is the documented count mismatch, and only + # on a tsam that cannot pin the count. + assert not SUPPORTS_PRESERVE, f"unexpected failure under preserve: {exc}" + assert "different cluster counts" in str(exc) + return + + # Success path: counts are uniform across slices (concat succeeded). + if SUPPORTS_PRESERVE: + assert res.clustering.n_clusters == n_clusters + assert res.cluster_representatives.sizes["cluster"] == n_clusters + else: + # Fallback tsam: slices happened to agree; the shared count is + # n_clusters plus a uniform number of extras. + assert res.clustering.n_clusters >= n_clusters + + +@pytest.mark.skipif( + not SUPPORTS_PRESERVE, reason="tsam lacks preserve_n_clusters (#410)" +) +class TestCountStability: + """Strict count guarantees, only meaningful when the flag is available.""" + + def test_sliced_append_pins_count(self) -> None: + extremes = tsam.ExtremeConfig( + method="append", + max_value=["solar", "wind", "price"], + min_value=["solar", "wind", "price"], + ) + res = tsam_xarray.aggregate( + _sliced_da(0, n_scenarios=4), + time_dim="time", + cluster_dim="variable", + n_clusters=8, + extremes=extremes, + ) + assert res.clustering.n_clusters == 8 + + def test_forced_even_for_single_aggregation(self) -> None: + """The flag is forced whether or not the data is sliced.""" + rng = np.random.default_rng(0) + time = pd.date_range("2020-01-01", periods=20 * 24, freq="h") + da = xr.DataArray( + rng.random((len(time), 3)), + dims=["time", "variable"], + coords={"time": time, "variable": VARIABLES}, + ) + extremes = tsam.ExtremeConfig( + method="append", max_value=["solar"], min_value=["wind"] + ) + res = tsam_xarray.aggregate( + da, time_dim="time", cluster_dim="variable", n_clusters=6, extremes=extremes + ) + assert res.clustering.n_clusters == 6 + + def test_user_preserve_false_is_overridden(self) -> None: + """The wrapper always forces the flag, even against an explicit False.""" + extremes = tsam.ExtremeConfig( + method="append", + max_value=["solar", "wind", "price"], + min_value=["solar", "wind", "price"], + preserve_n_clusters=False, + ) + res = tsam_xarray.aggregate( + _sliced_da(1, n_scenarios=3), + time_dim="time", + cluster_dim="variable", + n_clusters=8, + extremes=extremes, + ) + assert res.clustering.n_clusters == 8 + + +@pytest.mark.skipif( + not SUPPORTS_REPLACE_TRANSFER, + reason="tsam lacks transferable replace (#410)", +) +class TestReplaceTransfers: + """`replace` now survives the transfer paths it used to be barred from.""" + + def test_cluster_on_replace_allowed(self) -> None: + da = _sliced_da(0).isel(scenario=0, drop=True) + res = tsam_xarray.aggregate( + da, + time_dim="time", + cluster_dim="variable", + n_clusters=6, + cluster_on=["solar", "wind"], + extremes=tsam.ExtremeConfig( + method="replace", max_value=["solar"], min_value=["wind"] + ), + ) + # passive column carried through the transfer + assert "price" in res.cluster_representatives.coords["variable"].values + + def test_replace_roundtrips_through_apply(self) -> None: + da = _sliced_da(2).isel(scenario=0, drop=True) + extremes = tsam.ExtremeConfig( + method="replace", max_value=["solar", "wind"], min_value=["solar", "wind"] + ) + direct = tsam_xarray.aggregate( + da, time_dim="time", cluster_dim="variable", n_clusters=8, extremes=extremes + ) + applied = direct.clustering.apply(da) + np.testing.assert_allclose( + direct.cluster_representatives.values, + applied.cluster_representatives.transpose( + *direct.cluster_representatives.dims + ).values, + )