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

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/.github/workflows/dependabot-auto-merge.yaml

@@ -13,7 +13,7 @@ jobs:
     if: github.actor == 'dependabot[bot]'
     runs-on: ubuntu-24.04
     steps:
-      - uses: dependabot/fetch-metadata@v2
+      - uses: dependabot/fetch-metadata@v3
         id: metadata
 
       - name: Generate app token

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/.github/workflows/release.yaml

@@ -20,7 +20,7 @@ jobs:
           app-id: ${{ secrets.APP_ID }}
           private-key: ${{ secrets.APP_PRIVATE_KEY }}
 
-      - uses: googleapis/release-please-action@v4
+      - uses: googleapis/release-please-action@v5
         id: release
         with:
           token: ${{ steps.app-token.outputs.token }}

tsam_xarray-0.6.1/.release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+  ".": "0.6.1"
+}

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/CHANGELOG.md

@@ -1,5 +1,32 @@
 # Changelog
 
+## [0.6.1](https://github.com/FBumann/tsam_xarray/compare/v0.6.0...v0.6.1) (2026-07-09)
+
+
+### Features
+
+* 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))
+* validate extremes options against cluster_on ([#91](https://github.com/FBumann/tsam_xarray/issues/91)) ([2fa5a91](https://github.com/FBumann/tsam_xarray/commit/2fa5a91c5222f8b6dd6932d9572dec62e90f595b))
+
+## [0.6.0](https://github.com/FBumann/tsam_xarray/compare/v0.5.2...v0.6.0) (2026-05-27)
+
+
+### ⚠ BREAKING CHANGES
+
+* **ClusteringResult:** `time_coords` attribute and the `_time_coords_to_dict` / `_time_coords_from_dict` helpers have been removed. The time index now lives inside the tsam `ClusteringResult` payload (`time_index`) and flows through `aggregate` / `disaggregate` natively. Pre-0.6 JSONs are still loadable — the legacy outer `time_coords` field is forwarded to the inner `time_index` with a `DeprecationWarning`. Re-save to silence the warning. ([#83](https://github.com/FBumann/tsam_xarray/pull/83))
+
+
+### Refactors
+
+* Reuse tsam 3.4's `DatetimeIndex` round-trip in `disaggregate`, dropping the parallel `time_coords` field, the compact serialization helpers, and the manual `MultiIndex` truncation in `_disaggregate_single`. ([#83](https://github.com/FBumann/tsam_xarray/pull/83))
+
+
+### Dependencies
+
+* Bump minimum `tsam` to `>=3.4.0` (required for the `time_index` round-trip above).
+* Bump `googleapis/release-please-action` from 4 to 5. ([#82](https://github.com/FBumann/tsam_xarray/pull/82))
+* Bump `dependabot/fetch-metadata` from 2 to 3. ([#81](https://github.com/FBumann/tsam_xarray/pull/81))
+
 ## [0.5.2](https://github.com/FBumann/tsam_xarray/compare/v0.5.1...v0.5.2) (2026-04-01)
 
 

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tsam_xarray
-Version: 0.5.2
+Version: 0.6.1
 Summary: Lightweight xarray wrapper for tsam time series aggregation
 License-Expression: MIT
 License-File: LICENSE
@@ -13,7 +13,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.11
 Requires-Dist: bottleneck>=1.4
-Requires-Dist: tsam>=3.3.0
+Requires-Dist: tsam>=3.4.0
 Requires-Dist: xarray>=2024.1
 Provides-Extra: plot
 Requires-Dist: plotly>=5; extra == 'plot'

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/docs/examples/multi-dim.ipynb

@@ -11,7 +11,7 @@
     "- **`cluster_dim`** — dimensions clustered together (shared clustering)\n",
     "- **Auto-slicing** — remaining dimensions get independent clusterings\n",
     "\n",
-    "This notebook covers stacking, slicing, and weights."
+    "This notebook covers stacking, slicing, weights, and selecting what to cluster on."
    ]
   },
   {
@@ -212,6 +212,44 @@
     ")\n",
     "result_w2.accuracy.rmse.to_dataframe(\"RMSE\")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Selecting what to cluster on\n",
+    "\n",
+    "Sometimes a variable should be *carried through* the aggregation without\n",
+    "influencing how the clusters are formed — e.g. a price or demand series you\n",
+    "want on the representative grid, but that shouldn't drive period selection.\n",
+    "\n",
+    "Use `cluster_on` to restrict the clustering to specific coordinates. Everything\n",
+    "else is still aggregated and reconstructed from the resulting clusters, but has\n",
+    "no effect on the cluster distances. This differs from a small `weights` value:\n",
+    "`cluster_on` fully removes the coordinate from the distance metric.\n",
+    "\n",
+    "Here we cluster on `solar` and `wind` only, while `demand` is carried along."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cluster on solar + wind only; demand is carried along but doesn't\n",
+    "# influence cluster selection.\n",
+    "result_on = tsam_xarray.aggregate(\n",
+    "    da_single,\n",
+    "    time_dim=\"time\",\n",
+    "    cluster_dim=[\"variable\", \"region\"],\n",
+    "    n_clusters=4,\n",
+    "    cluster_on={\"variable\": [\"solar\", \"wind\"]},\n",
+    ")\n",
+    "# demand is still present in the output, aggregated from the same clusters\n",
+    "print(\"Variables:\", list(result_on.cluster_representatives.coords[\"variable\"].values))\n",
+    "result_on.accuracy.rmse.to_dataframe(\"RMSE\")"
+   ]
   }
  ],
  "metadata": {

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
 ]
 dependencies = [
     "bottleneck>=1.4",
-    "tsam>=3.3.0",
+    "tsam>=3.4.0",
     "xarray>=2024.1",
 ]
 

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/src/tsam_xarray/_clustering.py

@@ -21,25 +21,6 @@ from tsam_xarray._core import (
 )
 
 
-def _time_coords_to_dict(tc: pd.DatetimeIndex) -> dict[str, Any] | list[str]:
-    """Serialize a DatetimeIndex compactly when possible.
-
-    Regular indices are stored as ``{start, periods, freq}`` (~3 values).
-    Irregular indices fall back to a full ISO string list.
-    """
-    freq = pd.infer_freq(tc)
-    if freq is not None:
-        return {"start": tc[0].isoformat(), "periods": len(tc), "freq": freq}
-    return [t.isoformat() for t in tc]
-
-
-def _time_coords_from_dict(raw: dict[str, Any] | list[str]) -> pd.DatetimeIndex:
-    """Deserialize a DatetimeIndex from either compact or list format."""
-    if isinstance(raw, dict):
-        return pd.date_range(raw["start"], periods=raw["periods"], freq=raw["freq"])
-    return pd.DatetimeIndex(raw)
-
-
 @dataclass(frozen=True, repr=False)
 class ClusteringResult:
     """Reusable clustering result with xarray dimension metadata.
@@ -54,8 +35,6 @@ class ClusteringResult:
         slice_dims: Dimension(s) aggregated independently.
         clusterings: Per-slice tsam clustering.
             Single entry ``{(): result}`` when no slicing.
-        time_coords: Original time coordinates.
-            Needed for ``disaggregate()``.
         n_clusters: Number of clusters.
         n_original_periods: Number of original periods.
         n_timesteps_per_period: Timesteps per period.
@@ -80,7 +59,6 @@ class ClusteringResult:
     cluster_dim: list[str]
     slice_dims: list[str]
     clusterings: dict[tuple[Hashable, ...], tsam.ClusteringResult]
-    time_coords: pd.DatetimeIndex | None = field(default=None, repr=False)
     _cache: dict[str, Any] = field(
         default_factory=dict, repr=False, init=False, compare=False
     )
@@ -427,29 +405,10 @@ class ClusteringResult:
         Returns:
             Data with ``cluster`` and ``timestep`` replaced by
             the original ``time`` dimension.
-
-        Raises:
-            ValueError: If time coordinates are not available
-                (e.g., loaded from an old JSON that predates
-                this feature).
         """
-        if self.time_coords is None:
-            msg = (
-                "No time coordinates available. "
-                "This ClusteringResult was loaded from a JSON file "
-                "that does not contain time coordinate data. "
-                "Re-run aggregate() or save from a newer version."
-            )
-            raise ValueError(msg)
-
         slice_dims = self.slice_dims
         if not slice_dims:
-            cr = self.clusterings[()]
-            return _disaggregate_single(
-                self.time_coords,
-                cr,
-                data,
-            )
+            return _disaggregate_single(self.clusterings[()], data)
 
         import itertools
 
@@ -460,7 +419,7 @@ class ClusteringResult:
             sel = dict(zip(slice_dims, key, strict=True))
             data_slice = data.sel(sel)
             cr = _lookup_clustering(self.clusterings, key)
-            results.append(_disaggregate_single(self.time_coords, cr, data_slice))
+            results.append(_disaggregate_single(cr, data_slice))
 
         return _concat_along_dims(results, slice_dims, slice_coords)
 
@@ -479,15 +438,12 @@ class ClusteringResult:
                     "clustering": cr.to_dict(),
                 }
             )
-        data: dict[str, Any] = {
+        return {
             "time_dim": self.time_dim,
             "cluster_dim": self.cluster_dim,
             "slice_dims": self.slice_dims,
             "clusterings": entries,
         }
-        if self.time_coords is not None:
-            data["time_coords"] = _time_coords_to_dict(self.time_coords)
-        return data
 
     def to_json(self, path: str | Path, **json_kwargs: Any) -> None:
         """Save clustering to JSON file.
@@ -510,21 +466,31 @@ class ClusteringResult:
         Returns:
             The loaded ``ClusteringResult``.
         """
+        # Backcompat: pre-0.6 wrappers stored the time index as an outer
+        # ``time_coords`` key while the inner tsam blob (written by tsam<3.4)
+        # had no ``time_index``. Forward it so disaggregate keeps datetimes.
+        if "time_coords" in data:
+            import warnings
+
+            warnings.warn(
+                "Loading a legacy tsam_xarray JSON with an outer 'time_coords' "
+                "field; re-save with to_json() to silence this warning.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            for entry in data["clusterings"]:
+                entry["clustering"].setdefault("time_index", data["time_coords"])
+
         clusterings: dict[tuple[Hashable, ...], tsam.ClusteringResult] = {}
         for entry in data["clusterings"]:
             key = tuple(entry["key"])
             clusterings[key] = tsam.ClusteringResult.from_dict(entry["clustering"])
 
-        time_coords: pd.DatetimeIndex | None = None
-        if "time_coords" in data:
-            time_coords = _time_coords_from_dict(data["time_coords"])
-
         return cls(
             time_dim=data["time_dim"],
             cluster_dim=data["cluster_dim"],
             slice_dims=data.get("slice_dims", []),
             clusterings=clusterings,
-            time_coords=time_coords,
         )
 
     @classmethod
@@ -653,7 +619,6 @@ def _apply_single(
         cluster_dim=col_dims,
         slice_dims=[],
         clusterings={(): tsam_result.clustering},
-        time_coords=pd.DatetimeIndex(da.coords[time_dim].values),
     )
 
     return AggregationResult(
@@ -670,11 +635,14 @@ def _apply_single(
 
 
 def _disaggregate_single(
-    time_coords: pd.DatetimeIndex,
     cr: tsam.ClusteringResult,
     data: xr.DataArray,
 ) -> xr.DataArray:
-    """Disaggregate a single (non-sliced) DataArray using a ClusteringResult."""
+    """Disaggregate a single (non-sliced) DataArray using a ClusteringResult.
+
+    Relies on tsam's ``cr.disaggregate()`` to return a DataFrame indexed
+    by the original ``DatetimeIndex`` stored on the clustering.
+    """
     other_dims = [str(d) for d in data.dims if d not in ("cluster", "timestep")]
     ordered = data.transpose("cluster", "timestep", *other_dims)
 
@@ -686,10 +654,11 @@ def _disaggregate_single(
     flat = ordered.values.reshape(n_clusters * n_timesteps, -1)
 
     if cr.segment_durations is not None:
-        idx_tuples = []
-        for c in clusters:
-            for seg, dur in enumerate(cr.segment_durations[int(c)]):
-                idx_tuples.append((int(c), seg, int(dur)))
+        idx_tuples = [
+            (int(c), seg, int(dur))
+            for c in clusters
+            for seg, dur in enumerate(cr.segment_durations[int(c)])
+        ]
         mi = pd.MultiIndex.from_tuples(
             idx_tuples, names=["cluster", "segment", "duration"]
         )
@@ -700,12 +669,10 @@ def _disaggregate_single(
 
     df = pd.DataFrame(flat, index=mi, columns=range(flat.shape[1]))
     expanded = cr.disaggregate(df)
-
-    n_original = len(time_coords)
-    vals = expanded.values[:n_original]
+    time_coords = expanded.index
 
     if other_dims:
-        vals = vals.reshape(n_original, *other_sizes)
+        vals = expanded.values.reshape(len(time_coords), *other_sizes)
         result = xr.DataArray(
             vals,
             dims=["time", *other_dims],
@@ -716,7 +683,7 @@ def _disaggregate_single(
                 result = result.assign_coords({d: data.coords[d]})
     else:
         result = xr.DataArray(
-            vals[:, 0],
+            expanded.values[:, 0],
             dims=["time"],
             coords={"time": time_coords},
         )

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/src/tsam_xarray/_core.py

@@ -14,6 +14,7 @@ import xarray as xr
 from tsam_xarray._result import AccuracyMetrics, AggregationResult
 
 Weights = dict[str, float] | dict[str, dict[str, float]] | None
+ClusterOn = str | Sequence[str] | dict[str, Sequence[str]] | None
 
 
 def aggregate(
@@ -23,6 +24,7 @@ def aggregate(
     cluster_dim: Sequence[str] | str,
     n_clusters: int,
     weights: Weights = None,
+    cluster_on: ClusterOn = None,
     **tsam_kwargs: Any,
 ) -> AggregationResult:
     """Aggregate an xarray DataArray using tsam.
@@ -55,6 +57,37 @@ def aggregate(
               e.g. ``("solar", "north")`` gets weight
               ``2.0 * 1.5 = 3.0``.
 
+        cluster_on: Restrict which coordinates drive the
+            clustering. Coordinates not selected are still
+            aggregated and reconstructed from the resulting
+            clusters, but have no influence on how the clusters
+            are formed. ``None`` (default) clusters on
+            everything. Mirrors the ``weights`` formats:
+
+            - **List** (single ``cluster_dim``)::
+
+                  cluster_on=["solar", "wind"]
+
+            - **Dict** (multiple ``cluster_dim``)::
+
+                  cluster_on={"variable": ["solar", "wind"]}
+
+              A dim omitted from the dict is unrestricted (all
+              its coordinates are eligible). When several dims
+              are listed, a column must match on *all* of them
+              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.
+
         **tsam_kwargs: Additional keyword arguments passed to
             ``tsam.aggregate()``.
     """
@@ -65,10 +98,18 @@ def aggregate(
     da = _validate_data(da, time_dim, col_dims, slice_dims)
     _validate_no_cluster_config_weights(tsam_kwargs)
     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)
 
     if not slice_dims:
         return _aggregate_single(
-            da, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
+            da,
+            n_clusters,
+            time_dim,
+            col_dims,
+            per_dim_weights,
+            active_coords,
+            tsam_kwargs,
         )
 
     slice_coords = {d: da.coords[d].values for d in slice_dims}
@@ -80,7 +121,13 @@ def aggregate(
         sel = dict(zip(slice_dims, key, strict=True))
         da_slice = da.sel(sel)
         r = _aggregate_single(
-            da_slice, n_clusters, time_dim, col_dims, per_dim_weights, tsam_kwargs
+            da_slice,
+            n_clusters,
+            time_dim,
+            col_dims,
+            per_dim_weights,
+            active_coords,
+            tsam_kwargs,
         )
         results.append(r)
 
@@ -129,6 +176,68 @@ def _validate_no_cluster_config_weights(
         raise ValueError(msg)
 
 
+def _validate_extremes_with_cluster_on(
+    tsam_kwargs: dict[str, Any],
+    active_coords: dict[str, set[str]] | None,
+    da: xr.DataArray,
+) -> None:
+    """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:
+
+    - ``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).
+    - 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
+      cause instead.
+    """
+    if active_coords is None:
+        return
+    extremes = tsam_kwargs.get("extremes")
+    if extremes is None:
+        return
+
+    if getattr(extremes, "method", None) == "replace":
+        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)."
+        )
+        raise ValueError(msg)
+
+    passive_coords: set[str] = set()
+    for dim, active_set in active_coords.items():
+        all_coords = {str(c) for c in da.coords[dim].values}
+        passive_coords |= all_coords - active_set
+
+    referenced = (
+        set(extremes.max_value)
+        | set(extremes.min_value)
+        | set(extremes.max_period)
+        | set(extremes.min_period)
+    )
+    excluded = referenced & passive_coords
+    if excluded:
+        msg = (
+            f"extremes references coordinates {excluded} that cluster_on "
+            "excluded from the clustering. Extreme periods can only be "
+            "identified on clustered-on columns — add these coordinates to "
+            "cluster_on, or remove them from the ExtremeConfig."
+        )
+        raise ValueError(msg)
+
+
 def _validate(
     da: xr.DataArray,
     time_dim: str,
@@ -265,7 +374,7 @@ def _representatives_to_da(
     df = df.copy()
     # With segmentation, index has 3 levels: (cluster, segment_step, segment_duration)
     # Without: 2 levels: (cluster, timestep)
-    if df.index.nlevels == 3:
+    if isinstance(df.index, pd.MultiIndex) and df.index.nlevels == 3:
         df.index = df.index.droplevel(2)  # drop segment_duration
     df.index.names = ["cluster", "timestep"]
 
@@ -339,6 +448,36 @@ def _metric_to_da(
     return xr.DataArray(series.to_xarray())
 
 
+def _validate_dim_coord_mapping(
+    mapping: dict[str, Any],
+    da: xr.DataArray,
+    col_dims: list[str],
+    label: str,
+) -> None:
+    """Validate a ``{dim: coords}`` mapping references real dims and coords.
+
+    ``coords`` may be any iterable of coordinate names (a list, or the keys
+    of a ``{coord: weight}`` dict). Shared by ``weights`` and ``cluster_on``.
+    """
+    extra_dims = set(mapping.keys()) - set(col_dims)
+    if extra_dims:
+        msg = (
+            f"{label} has unknown dims {extra_dims}, "
+            f"must be subset of cluster_dim {col_dims}"
+        )
+        raise ValueError(msg)
+
+    for dim_name, coords in mapping.items():
+        valid_coords = set(str(c) for c in da.coords[dim_name].values)
+        unknown = set(str(c) for c in coords) - valid_coords
+        if unknown:
+            msg = (
+                f"{label} has unknown coords {unknown} for dim {dim_name!r}, "
+                f"valid coords: {sorted(valid_coords)}"
+            )
+            raise ValueError(msg)
+
+
 def _normalize_weights(
     weights: dict[str, float] | dict[str, dict[str, float]] | None,
     da: xr.DataArray,
@@ -370,27 +509,51 @@ def _normalize_weights(
             raise ValueError(msg)
         per_dim_weights = {col_dims[0]: weights}  # type: ignore[dict-item]
 
-    # Validate dim names exist in cluster_dim
-    extra_dims = set(per_dim_weights.keys()) - set(col_dims)
-    if extra_dims:
-        msg = (
-            f"weights has unknown dims {extra_dims}, "
-            f"must be subset of cluster_dim {col_dims}"
-        )
-        raise ValueError(msg)
+    _validate_dim_coord_mapping(per_dim_weights, da, col_dims, "weights")
+    return per_dim_weights
 
-    # Validate coord values exist in the DataArray
-    for dim_name, coord_weights in per_dim_weights.items():
-        valid_coords = set(str(c) for c in da.coords[dim_name].values)
-        unknown = set(coord_weights.keys()) - valid_coords
-        if unknown:
+
+def _normalize_cluster_on(
+    cluster_on: ClusterOn,
+    da: xr.DataArray,
+    col_dims: list[str],
+) -> dict[str, set[str]] | None:
+    """Normalize cluster_on to a per-dim set of eligible coords.
+
+    A list/str form requires a single cluster_dim. Returns ``None``
+    when nothing is restricted (cluster on everything).
+    """
+    if cluster_on is None:
+        return None
+
+    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()
+        }
+    else:
+        if len(col_dims) != 1:
             msg = (
-                f"weights has unknown coords {unknown} for dim {dim_name!r}, "
-                f"valid coords: {sorted(valid_coords)}"
+                "List-form cluster_on requires a single cluster_dim. "
+                "For multiple cluster_dim, use a dict: "
+                '{"dim_name": ["coord", ...]}.'
             )
             raise ValueError(msg)
+        coords = [cluster_on] if isinstance(cluster_on, str) else list(cluster_on)
+        per_dim = {col_dims[0]: [str(c) for c in coords]}
 
-    return per_dim_weights
+    _validate_dim_coord_mapping(per_dim, da, col_dims, "cluster_on")
+
+    # An empty coord list for any restricted dim leaves no column to cluster
+    # on (the stacked columns are the full coordinate grid). Fail here, before
+    # any aggregation runs, rather than deep in the per-slice loop.
+    if any(not coords for coords in per_dim.values()):
+        msg = (
+            "cluster_on selected no columns — nothing to cluster on. "
+            "Select at least one coordinate present in the data."
+        )
+        raise ValueError(msg)
+
+    return {d: set(cs) for d, cs in per_dim.items()}
 
 
 def _translate_weights(
@@ -414,28 +577,103 @@ def _translate_weights(
     return flat
 
 
+# tsam.aggregate() kwargs that ClusteringResult.apply() also accepts. Only
+# these are forwarded to apply() when re-aggregating passive columns onto a
+# clustering built from the active subset; the rest (period_duration, cluster,
+# segments, extremes, ...) are already baked into the stored clustering.
+_APPLY_KWARGS = frozenset(
+    {"temporal_resolution", "round_decimals", "numerical_tolerance"}
+)
+
+
+def _active_columns(
+    df: pd.DataFrame,
+    active_coords: dict[str, set[str]] | None,
+    col_dims: list[str],
+) -> list[Hashable]:
+    """Columns that drive the clustering, given a cluster_on selection.
+
+    A column is active only if, for every restricted dim, its coordinate is
+    in that dim's eligible set; unrestricted dims impose no constraint. With
+    no restriction (``active_coords is None``) every column is active. The
+    remaining columns are *passive*: still aggregated and reconstructed, but
+    excluded from the cluster distances.
+    """
+    if active_coords is None:
+        return list(df.columns)
+
+    def is_active(col: Hashable) -> bool:
+        pairs = (
+            zip(col_dims, col, strict=True)
+            if isinstance(col, tuple)
+            else [(col_dims[0], col)]
+        )
+        return all(
+            str(coord) in active_coords[dim]
+            for dim, coord in pairs
+            if dim in active_coords
+        )
+
+    return [col for col in df.columns if is_active(col)]
+
+
 def _aggregate_single(
     da: xr.DataArray,
     n_clusters: int,
     time_dim: str,
     col_dims: list[str],
     weights: dict[str, dict[str, float]] | None,
+    active_coords: dict[str, set[str]] | None,
     tsam_kwargs: dict[str, Any],
 ) -> AggregationResult:
-    """Run a single tsam aggregation on a DataArray."""
+    """Run a single tsam aggregation on a DataArray.
+
+    When ``cluster_on`` restricts the clustering, the clustering runs on
+    the active subset only, then is applied back to the full frame so the
+    passive columns still get representatives and reconstruction.
+    """
     df = _to_dataframe(da, time_dim, col_dims)
 
     tsam_weights: dict[Hashable, float] | None = None
     if weights is not None:
         tsam_weights = _translate_weights(weights, df, col_dims)
 
-    tsam_result = tsam.aggregate(
-        df,
-        n_clusters,
-        weights=tsam_weights,  # type: ignore[arg-type]
-        **tsam_kwargs,
-    )
+    active_cols = _active_columns(df, active_coords, col_dims)
 
+    if len(active_cols) == len(df.columns):
+        # Nothing carried passively — cluster on the full frame directly.
+        tsam_result = tsam.aggregate(
+            df,
+            n_clusters,
+            weights=tsam_weights,  # type: ignore[arg-type]
+            **tsam_kwargs,
+        )
+    else:
+        # Cluster on the active subset, then apply that clustering back to the
+        # full frame so passive columns get representatives and reconstruction.
+        active_weights = (
+            None if tsam_weights is None else {c: tsam_weights[c] for c in active_cols}
+        )
+        clustering = tsam.aggregate(
+            df[active_cols],
+            n_clusters,
+            weights=active_weights,  # type: ignore[arg-type]
+            **tsam_kwargs,
+        ).clustering
+        apply_kwargs = {k: v for k, v in tsam_kwargs.items() if k in _APPLY_KWARGS}
+        tsam_result = clustering.apply(df, **apply_kwargs)
+
+    return _result_from_tsam(tsam_result, da, df, time_dim, col_dims)
+
+
+def _result_from_tsam(
+    tsam_result: Any,
+    da: xr.DataArray,
+    df: pd.DataFrame,
+    time_dim: str,
+    col_dims: list[str],
+) -> AggregationResult:
+    """Build an AggregationResult from a tsam aggregation result."""
     typical = _representatives_to_da(tsam_result.cluster_representatives, col_dims)
     reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims)
 
@@ -475,7 +713,6 @@ def _aggregate_single(
         cluster_dim=col_dims,
         slice_dims=[],
         clusterings={(): tsam_result.clustering},
-        time_coords=pd.DatetimeIndex(da.coords[time_dim].values),
     )
 
     return AggregationResult(
@@ -562,7 +799,6 @@ def _concat_results(
         cluster_dim=first.clustering.cluster_dim,
         slice_dims=slice_dims,
         clusterings=merged_clusterings,
-        time_coords=first.clustering.time_coords,
     )
 
     return AggregationResult(

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/src/tsam_xarray/_result.py

@@ -211,10 +211,7 @@ class AggregationResult:
 
     def _disaggregate_single(self, data: xr.DataArray) -> xr.DataArray:
         """Disaggregate without slice dims."""
-        import pandas as pd
-
         from tsam_xarray._clustering import _disaggregate_single
 
-        time_coords = pd.DatetimeIndex(self.original.coords["time"].values)
         cr = self.clustering.clusterings[()]
-        return _disaggregate_single(time_coords, cr, data)
+        return _disaggregate_single(cr, data)

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/src/tsam_xarray/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.2'
-__version_tuple__ = version_tuple = (0, 5, 2)
+__version__ = version = '0.6.1'
+__version_tuple__ = version_tuple = (0, 6, 1)
 
 __commit_id__ = commit_id = None

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/test/test_aggregate.py

@@ -306,6 +306,169 @@ class TestWeights:
             )
 
 
+@pytest.fixture
+def da_3var() -> xr.DataArray:
+    """Single cluster_dim: solar/wind/price over time (region dropped)."""
+    return _make_da(variables=["solar", "wind", "price"]).isel(region=0, drop=True)
+
+
+@pytest.fixture
+def da_3var_region() -> xr.DataArray:
+    """Multi cluster_dim: solar/wind/price over time and region."""
+    return _make_da(variables=["solar", "wind", "price"])
+
+
+class TestClusterOn:
+    """cluster_on restricts which coords drive clustering, keeping the rest."""
+
+    @pytest.mark.parametrize(
+        ("cluster_dim", "cluster_on", "subset"),
+        [
+            ("variable", ["solar", "wind"], ["solar", "wind"]),
+            ("variable", "solar", ["solar"]),
+            (
+                ["variable", "region"],
+                {"variable": ["solar", "wind"]},
+                ["solar", "wind"],
+            ),
+        ],
+        ids=["list", "string", "dict-multi"],
+    )
+    def test_matches_subset_clustering(
+        self,
+        da_3var: xr.DataArray,
+        da_3var_region: xr.DataArray,
+        cluster_dim: list[str] | str,
+        cluster_on: object,
+        subset: list[str],
+    ):
+        """cluster_on gives the same clustering as aggregating the subset alone,
+        while carried coords still appear in the output."""
+        da = da_3var if cluster_dim == "variable" else da_3var_region
+        restricted = tsam_xarray.aggregate(
+            da,
+            time_dim="time",
+            cluster_dim=cluster_dim,
+            n_clusters=4,
+            cluster_on=cluster_on,
+        )
+        reference = tsam_xarray.aggregate(
+            da.sel(variable=subset),
+            time_dim="time",
+            cluster_dim=cluster_dim,
+            n_clusters=4,
+        )
+        np.testing.assert_array_equal(
+            restricted.clustering.cluster_assignments.values,
+            reference.clustering.cluster_assignments.values,
+        )
+        reps = restricted.cluster_representatives
+        assert "price" in reps.coords["variable"].values
+        assert "price" in restricted.reconstructed.coords["variable"].values
+
+    def test_including_var_would_change_clustering(self, da_3var: xr.DataArray):
+        """Sanity: clustering on everything gives different assignments."""
+        restricted = tsam_xarray.aggregate(
+            da_3var,
+            time_dim="time",
+            cluster_dim="variable",
+            n_clusters=4,
+            cluster_on=["solar", "wind"],
+        )
+        full = tsam_xarray.aggregate(
+            da_3var, time_dim="time", cluster_dim="variable", n_clusters=4
+        )
+        assert (
+            restricted.clustering.cluster_assignments.values
+            != full.clustering.cluster_assignments.values
+        ).any()
+
+    def test_composes_with_weights(self, da_3var: xr.DataArray):
+        """cluster_on and weights coexist; weights apply to the active subset."""
+        result = tsam_xarray.aggregate(
+            da_3var,
+            time_dim="time",
+            cluster_dim="variable",
+            n_clusters=4,
+            cluster_on=["solar", "wind"],
+            weights={"solar": 5.0, "wind": 1.0},
+        )
+        assert result.n_clusters == 4
+        assert "price" in result.cluster_representatives.coords["variable"].values
+
+    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."""
+        import tsam
+
+        with pytest.raises(ValueError, match=r"replace.*not supported.*cluster_on"):
+            tsam_xarray.aggregate(
+                da_3var,
+                time_dim="time",
+                cluster_dim="variable",
+                n_clusters=4,
+                cluster_on=["solar", "wind"],
+                extremes=tsam.ExtremeConfig(method="replace", max_value=["solar"]),
+            )
+
+    def test_extremes_on_excluded_column_rejected(self, da_3var: xr.DataArray):
+        """An ExtremeConfig referencing a cluster_on-excluded coord errors with
+        a clear cause, not tsam's misleading 'not found in data'."""
+        import tsam
+
+        with pytest.raises(ValueError, match=r"cluster_on excluded"):
+            tsam_xarray.aggregate(
+                da_3var,
+                time_dim="time",
+                cluster_dim="variable",
+                n_clusters=4,
+                cluster_on=["solar", "wind"],
+                extremes=tsam.ExtremeConfig(method="append", max_value=["price"]),
+            )
+
+    def test_append_extremes_allowed(self, da_3var: xr.DataArray):
+        """extremes='append' on an active column composes with cluster_on."""
+        import tsam
+
+        result = tsam_xarray.aggregate(
+            da_3var,
+            time_dim="time",
+            cluster_dim="variable",
+            n_clusters=4,
+            cluster_on=["solar", "wind"],
+            extremes=tsam.ExtremeConfig(method="append", max_value=["solar"]),
+        )
+        assert "price" in result.cluster_representatives.coords["variable"].values
+
+    @pytest.mark.parametrize(
+        ("cluster_dim", "cluster_on", "match"),
+        [
+            (["variable", "region"], ["solar"], "single cluster_dim"),
+            (["variable", "region"], {"scenario": ["low"]}, "unknown dims"),
+            ("variable", ["nuclear"], "unknown coords"),
+            ("variable", {"variable": []}, "nothing to cluster on"),
+        ],
+        ids=["list-multi", "unknown-dim", "unknown-coord", "empty"],
+    )
+    def test_invalid_cluster_on_raises(
+        self,
+        da_3var: xr.DataArray,
+        da_3var_region: xr.DataArray,
+        cluster_dim: list[str] | str,
+        cluster_on: object,
+        match: str,
+    ):
+        da = da_3var if cluster_dim == "variable" else da_3var_region
+        with pytest.raises(ValueError, match=match):
+            tsam_xarray.aggregate(
+                da,
+                time_dim="time",
+                cluster_dim=cluster_dim,
+                n_clusters=4,
+                cluster_on=cluster_on,
+            )
+
+
 class TestWeightTranslation:
     """Verify weight values are correctly translated."""
 
@@ -1283,30 +1446,8 @@ class TestClusteringDisaggregate:
         expected = result.disaggregate(result.cluster_representatives)
         xr.testing.assert_allclose(dis, expected)
 
-    def test_disaggregate_no_time_coords_raises(self):
-        """ClusteringInfo without time_coords raises ValueError."""
-        from tsam_xarray._clustering import ClusteringInfo
-
-        da = _make_da()
-        da_flat = da.isel(region=0).drop_vars("region")
-        result = tsam_xarray.aggregate(
-            da_flat,
-            time_dim="time",
-            cluster_dim="variable",
-            n_clusters=4,
-        )
-        # Construct a ClusteringInfo without time_coords
-        ci = ClusteringInfo(
-            time_dim="time",
-            cluster_dim=["variable"],
-            slice_dims=[],
-            clusterings=result.clustering.clusterings,
-        )
-        with pytest.raises(ValueError, match="No time coordinates"):
-            ci.disaggregate(result.cluster_representatives)
-
-    def test_json_backward_compat(self, tmp_path):
-        """Old JSON without time_coords loads fine, disaggregate raises."""
+    def test_legacy_time_coords_json(self, tmp_path):
+        """Legacy JSON with outer ``time_coords`` loads with a warning."""
         import json
 
         da = _make_da()
@@ -1320,61 +1461,19 @@ class TestClusteringDisaggregate:
         path = tmp_path / "clustering.json"
         result.clustering.to_json(str(path))
 
-        # Strip time_coords from saved JSON
+        # Forge a pre-0.6 file: outer time_coords + inner blob without time_index.
         with open(path) as f:
             data = json.load(f)
-        del data["time_coords"]
+        tc = data["clusterings"][0]["clustering"].pop("time_index")
+        data["time_coords"] = tc
         with open(path, "w") as f:
             json.dump(data, f)
 
-        clustering = tsam_xarray.load_clustering(str(path))
-        assert clustering.time_coords is None
-        # apply() still works
-        new_result = clustering.apply(da_flat)
-        assert new_result.n_clusters == 4
-        # disaggregate() raises
-        with pytest.raises(ValueError, match="No time coordinates"):
-            clustering.disaggregate(result.cluster_representatives)
-
-    def test_time_coords_in_json(self, tmp_path):
-        """time_coords are serialized in JSON."""
-        import json
+        with pytest.warns(DeprecationWarning, match="legacy tsam_xarray JSON"):
+            clustering = tsam_xarray.load_clustering(str(path))
 
-        da = _make_da()
-        da_flat = da.isel(region=0).drop_vars("region")
-        result = tsam_xarray.aggregate(
-            da_flat,
-            time_dim="time",
-            cluster_dim="variable",
-            n_clusters=4,
-        )
-        path = tmp_path / "clustering.json"
-        result.clustering.to_json(str(path))
-        with open(path) as f:
-            data = json.load(f)
-        assert "time_coords" in data
-        tc = data["time_coords"]
-        # Regular index → compact dict format
-        assert isinstance(tc, dict)
-        assert tc["periods"] == da_flat.sizes["time"]
-
-    def test_time_coords_roundtrip(self, tmp_path):
-        """time_coords survive JSON round-trip."""
-        da = _make_da()
-        da_flat = da.isel(region=0).drop_vars("region")
-        result = tsam_xarray.aggregate(
-            da_flat,
-            time_dim="time",
-            cluster_dim="variable",
-            n_clusters=4,
-        )
-        path = tmp_path / "clustering.json"
-        result.clustering.to_json(str(path))
-        clustering = tsam_xarray.load_clustering(str(path))
-        np.testing.assert_array_equal(
-            result.clustering.time_coords,
-            clustering.time_coords,
-        )
+        dis = clustering.disaggregate(result.cluster_representatives)
+        assert dis.indexes["time"].equals(da_flat.indexes["time"])
 
     def test_1d_disaggregate(self, tmp_path):
         """Disaggregate works on 1D time series."""

{tsam_xarray-0.5.2 → tsam_xarray-0.6.1}/test/test_parametrized.py

@@ -505,61 +505,3 @@ class TestClusteringIORoundtrip:
         clustering = tsam_xarray.load_clustering(str(path))
         new_result = clustering.apply(agg_case.da)
         assert new_result.is_transferred
-
-
-class TestTimeCoordsSerialization:
-    """Compact time_coords serialization helpers."""
-
-    def test_regular_index_compact(self):
-        """Regular DatetimeIndex is stored as {start, periods, freq}."""
-        import pandas as pd
-
-        from tsam_xarray._clustering import _time_coords_to_dict
-
-        tc = pd.date_range("2025-01-01", periods=8760, freq="h")
-        d = _time_coords_to_dict(tc)
-        assert isinstance(d, dict)
-        assert set(d) == {"start", "periods", "freq"}
-        assert d["periods"] == 8760
-
-    def test_regular_index_roundtrip(self):
-        """Compact format round-trips exactly."""
-        import pandas as pd
-
-        from tsam_xarray._clustering import (
-            _time_coords_from_dict,
-            _time_coords_to_dict,
-        )
-
-        tc = pd.date_range("2025-01-01", periods=8760, freq="h")
-        restored = _time_coords_from_dict(_time_coords_to_dict(tc))
-        pd.testing.assert_index_equal(tc, restored)
-
-    def test_irregular_index_fallback(self):
-        """Irregular DatetimeIndex falls back to list of ISO strings."""
-        import pandas as pd
-
-        from tsam_xarray._clustering import (
-            _time_coords_from_dict,
-            _time_coords_to_dict,
-        )
-
-        tc = pd.DatetimeIndex(["2025-01-01", "2025-01-03", "2025-01-07"])
-        d = _time_coords_to_dict(tc)
-        assert isinstance(d, list)
-        restored = _time_coords_from_dict(d)
-        pd.testing.assert_index_equal(tc, restored)
-
-    def test_old_list_format_still_loads(self):
-        """List format (pre-compact) is still accepted by from_dict."""
-        from tsam_xarray._clustering import _time_coords_from_dict
-
-        raw = ["2025-01-01T00:00:00", "2025-01-01T01:00:00"]
-        restored = _time_coords_from_dict(raw)
-        assert len(restored) == 2
-
-    def test_dict_roundtrip_uses_compact(self, agg_case: AggregateCase):
-        """to_dict uses compact format for regular time indices."""
-        result = _aggregate(agg_case)
-        d = result.clustering.to_dict()
-        assert isinstance(d["time_coords"], dict)

tsam_xarray-0.5.2/.release-please-manifest.json

@@ -1,3 +0,0 @@
-{
-  ".": "0.5.2"
-}