tsam-xarray 0.6.3__tar.gz → 0.6.4__tar.gz

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

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

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [0.6.4](https://github.com/FBumann/tsam_xarray/compare/v0.6.3...v0.6.4) (2026-07-22)
+
+
+### Features
+
+* make output dimension names configurable via DimNames ([#101](https://github.com/FBumann/tsam_xarray/issues/101)) ([558fb9d](https://github.com/FBumann/tsam_xarray/commit/558fb9da4014d4bdd05515f79fef5c542c17239e)), closes [#99](https://github.com/FBumann/tsam_xarray/issues/99)
+
 ## [0.6.3](https://github.com/FBumann/tsam_xarray/compare/v0.6.2...v0.6.3) (2026-07-22)
 
 

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tsam_xarray
-Version: 0.6.3
+Version: 0.6.4
 Summary: Lightweight xarray wrapper for tsam time series aggregation
 License-Expression: MIT
 License-File: LICENSE

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/docs/data-model.md

@@ -41,6 +41,34 @@ graph LR
 
 </div>
 
+### Custom output dimension names
+
+`aggregate()` adds four structural dimensions to its results that do not
+exist in the input: `cluster`, `timestep`, `period` (in
+`cluster_assignments`), and `segment` (segmented runs). These names are
+reserved — an input dimension of the same name raises an error.
+
+Pass a `DimNames` to rename them, e.g. when a caller already has a
+`period` dimension (multi-period optimization models):
+
+```python
+from tsam_xarray import DimNames, aggregate
+
+result = aggregate(
+    da,  # has a slice dim literally named "period"
+    time_dim="time",
+    cluster_dim="variable",
+    n_clusters=8,
+    dim_names=DimNames(period="original_period"),
+)
+result.cluster_assignments.dims  # ("original_period", ...)
+```
+
+The resolved names are stored on `ClusteringResult`, so `apply()`,
+`disaggregate()`, and the JSON round-trip all reproduce them. `dim_names`
+defaults to `None`, which keeps today's names. The chosen names must be
+unique and must not collide with any input dimension.
+
 ## ClusteringResult
 
 The reusable part — knows *how* the time series was clustered,

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/src/tsam_xarray/__init__.py

@@ -2,6 +2,7 @@
 
 from tsam_xarray._clustering import ClusteringInfo, ClusteringResult
 from tsam_xarray._core import aggregate
+from tsam_xarray._dim_names import DimNames
 from tsam_xarray._result import AccuracyMetrics, AggregationResult
 from tsam_xarray._tuning import (
     TuningResult,
@@ -18,6 +19,7 @@ __all__ = [
     "AggregationResult",
     "ClusteringInfo",
     "ClusteringResult",
+    "DimNames",
     "TuningResult",
     "aggregate",
     "find_best_combination",

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/src/tsam_xarray/_clustering.py

@@ -19,6 +19,7 @@ from tsam_xarray._core import (
     _resolve_cluster_dim,
     _segment_durations_to_da,
 )
+from tsam_xarray._dim_names import DimNames
 
 
 @dataclass(frozen=True, repr=False)
@@ -53,12 +54,15 @@ class ClusteringResult:
         segment_centers: Representative timestep per segment,
             or ``None``.
             Dims: ``(cluster, segment, *slice_dims)``.
+        dim_names: Names of the structural output dimensions.
+            See `DimNames`.
     """
 
     time_dim: str
     cluster_dim: list[str]
     slice_dims: list[str]
     clusterings: dict[tuple[Hashable, ...], tsam.ClusteringResult]
+    dim_names: DimNames = field(default_factory=DimNames)
     _cache: dict[str, Any] = field(
         default_factory=dict, repr=False, init=False, compare=False
     )
@@ -125,14 +129,19 @@ class ClusteringResult:
     def _build_assignments(self) -> xr.DataArray:
         if not self.slice_dims:
             cr = self.clusterings[()]
-            return xr.DataArray(list(cr.cluster_assignments), dims=["period"])
+            return xr.DataArray(
+                list(cr.cluster_assignments), dims=[self.dim_names.period]
+            )
 
         import itertools
 
         sc = self._slice_coords
         keys = list(itertools.product(*(sc[d] for d in self.slice_dims)))
         arrays = [
-            xr.DataArray(list(self.clusterings[k].cluster_assignments), dims=["period"])
+            xr.DataArray(
+                list(self.clusterings[k].cluster_assignments),
+                dims=[self.dim_names.period],
+            )
             for k in keys
         ]
         return _concat_along_dims(arrays, self.slice_dims, sc)
@@ -153,8 +162,8 @@ class ClusteringResult:
             counts = np.bincount(cr.cluster_assignments, minlength=cr.n_clusters)
             return xr.DataArray(
                 counts,
-                dims=["cluster"],
-                coords={"cluster": np.arange(cr.n_clusters)},
+                dims=[self.dim_names.cluster],
+                coords={self.dim_names.cluster: np.arange(cr.n_clusters)},
             )
 
         if not self.slice_dims:
@@ -180,18 +189,24 @@ class ClusteringResult:
 
     def _build_segment_durations(self) -> xr.DataArray | None:
         if not self.slice_dims:
-            return _segment_durations_to_da(self.clusterings[()].segment_durations)
+            return _segment_durations_to_da(
+                self.clusterings[()].segment_durations, self.dim_names
+            )
 
         import itertools
 
         sc = self._slice_coords
         keys = list(itertools.product(*(sc[d] for d in self.slice_dims)))
-        first = _segment_durations_to_da(self.clusterings[keys[0]].segment_durations)
+        first = _segment_durations_to_da(
+            self.clusterings[keys[0]].segment_durations, self.dim_names
+        )
         if first is None:
             return None
         das: list[xr.DataArray] = [first]
         for k in keys[1:]:
-            da = _segment_durations_to_da(self.clusterings[k].segment_durations)
+            da = _segment_durations_to_da(
+                self.clusterings[k].segment_durations, self.dim_names
+            )
             if da is None:
                 msg = (
                     f"Slice {k} has no segment durations but the first "
@@ -220,8 +235,8 @@ class ClusteringResult:
                 raise ValueError(msg)
             return xr.DataArray(
                 list(centers),
-                dims=["cluster"],
-                coords={"cluster": np.arange(cr.n_clusters)},
+                dims=[self.dim_names.cluster],
+                coords={self.dim_names.cluster: np.arange(cr.n_clusters)},
             )
 
         if not self.slice_dims:
@@ -251,10 +266,10 @@ class ClusteringResult:
                 return None
             return xr.DataArray(
                 np.array(cr.segment_assignments),
-                dims=["cluster", "timestep"],
+                dims=[self.dim_names.cluster, self.dim_names.timestep],
                 coords={
-                    "cluster": np.arange(cr.n_clusters),
-                    "timestep": np.arange(cr.n_timesteps_per_period),
+                    self.dim_names.cluster: np.arange(cr.n_clusters),
+                    self.dim_names.timestep: np.arange(cr.n_timesteps_per_period),
                 },
             )
 
@@ -298,10 +313,10 @@ class ClusteringResult:
             n_segments = cr.n_segments or len(cr.segment_centers[0])
             return xr.DataArray(
                 np.array(cr.segment_centers),
-                dims=["cluster", "segment"],
+                dims=[self.dim_names.cluster, self.dim_names.segment],
                 coords={
-                    "cluster": np.arange(cr.n_clusters),
-                    "segment": np.arange(n_segments),
+                    self.dim_names.cluster: np.arange(cr.n_clusters),
+                    self.dim_names.segment: np.arange(n_segments),
                 },
             )
 
@@ -367,7 +382,7 @@ class ClusteringResult:
 
         if not slice_dims:
             cr = self.clusterings[()]
-            return _apply_single(da, cr, td, cd, tsam_kwargs)
+            return _apply_single(da, cr, td, cd, tsam_kwargs, self.dim_names)
 
         import itertools
 
@@ -380,7 +395,7 @@ class ClusteringResult:
             sel = dict(zip(slice_dims, key, strict=True))
             da_slice = da.sel(sel)
             cr = _lookup_clustering(self.clusterings, key)
-            r = _apply_single(da_slice, cr, td, cd, tsam_kwargs)
+            r = _apply_single(da_slice, cr, td, cd, tsam_kwargs, self.dim_names)
             results.append(r)
 
         return _concat_results(results, slice_dims, slice_coords, slice_keys)
@@ -408,7 +423,7 @@ class ClusteringResult:
         """
         slice_dims = self.slice_dims
         if not slice_dims:
-            return _disaggregate_single(self.clusterings[()], data)
+            return _disaggregate_single(self.clusterings[()], data, self.dim_names)
 
         import itertools
 
@@ -419,7 +434,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(cr, data_slice))
+            results.append(_disaggregate_single(cr, data_slice, self.dim_names))
 
         return _concat_along_dims(results, slice_dims, slice_coords)
 
@@ -442,6 +457,12 @@ class ClusteringResult:
             "time_dim": self.time_dim,
             "cluster_dim": self.cluster_dim,
             "slice_dims": self.slice_dims,
+            "dim_names": {
+                "cluster": self.dim_names.cluster,
+                "timestep": self.dim_names.timestep,
+                "period": self.dim_names.period,
+                "segment": self.dim_names.segment,
+            },
             "clusterings": entries,
         }
 
@@ -486,11 +507,15 @@ class ClusteringResult:
             key = tuple(entry["key"])
             clusterings[key] = tsam.ClusteringResult.from_dict(entry["clustering"])
 
+        dim_names_data = data.get("dim_names")
+        dim_names = DimNames(**dim_names_data) if dim_names_data else DimNames()
+
         return cls(
             time_dim=data["time_dim"],
             cluster_dim=data["cluster_dim"],
             slice_dims=data.get("slice_dims", []),
             clusterings=clusterings,
+            dim_names=dim_names,
         )
 
     @classmethod
@@ -603,6 +628,7 @@ def _apply_single(
     time_dim: str,
     col_dims: list[str],
     tsam_kwargs: dict[str, Any],
+    dim_names: DimNames,
 ) -> Any:
     """Apply a single ClusteringResult to a DataArray."""
     import pandas as pd
@@ -620,18 +646,22 @@ def _apply_single(
     cr = _drop_missing_weights(cr, df)
     tsam_result = cr.apply(df, **tsam_kwargs)
 
-    typical = _representatives_to_da(tsam_result.cluster_representatives, col_dims)
+    typical = _representatives_to_da(
+        tsam_result.cluster_representatives, col_dims, dim_names
+    )
     reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims)
 
     cw = tsam_result.cluster_weights
     cluster_ids = np.array(sorted(cw.keys()))
     cluster_weights_da = xr.DataArray(
         np.array([cw[k] for k in cluster_ids]),
-        dims=["cluster"],
-        coords={"cluster": cluster_ids},
+        dims=[dim_names.cluster],
+        coords={dim_names.cluster: cluster_ids},
     )
 
-    assignments_da = xr.DataArray(tsam_result.cluster_assignments, dims=["period"])
+    assignments_da = xr.DataArray(
+        tsam_result.cluster_assignments, dims=[dim_names.period]
+    )
 
     col_names: list[str] | None = None
     if isinstance(df.columns, pd.MultiIndex):
@@ -650,13 +680,14 @@ def _apply_single(
         ),
     )
 
-    seg_durations = _segment_durations_to_da(tsam_result.segment_durations)
+    seg_durations = _segment_durations_to_da(tsam_result.segment_durations, dim_names)
 
     clustering_info = ClusteringResult(
         time_dim=time_dim,
         cluster_dim=col_dims,
         slice_dims=[],
         clusterings={(): tsam_result.clustering},
+        dim_names=dim_names,
     )
 
     return AggregationResult(
@@ -675,18 +706,21 @@ def _apply_single(
 def _disaggregate_single(
     cr: tsam.ClusteringResult,
     data: xr.DataArray,
+    dim_names: DimNames,
 ) -> xr.DataArray:
     """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)
+    cluster_dim = dim_names.cluster
+    timestep_dim = dim_names.timestep
+    other_dims = [str(d) for d in data.dims if d not in (cluster_dim, timestep_dim)]
+    ordered = data.transpose(cluster_dim, timestep_dim, *other_dims)
 
-    clusters = ordered.coords["cluster"].values
+    clusters = ordered.coords[cluster_dim].values
     n_clusters = len(clusters)
-    n_timesteps = ordered.sizes["timestep"]
+    n_timesteps = ordered.sizes[timestep_dim]
     other_sizes = ordered.shape[2:]
 
     flat = ordered.values.reshape(n_clusters * n_timesteps, -1)

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/src/tsam_xarray/_core.py

@@ -12,6 +12,7 @@ import pandas as pd
 import tsam
 import xarray as xr
 
+from tsam_xarray._dim_names import DimNames
 from tsam_xarray._result import AccuracyMetrics, AggregationResult
 
 Weights = dict[str, float] | dict[str, dict[str, float]] | None
@@ -26,6 +27,7 @@ def aggregate(
     n_clusters: int,
     weights: Weights = None,
     cluster_on: ClusterOn = None,
+    dim_names: DimNames | None = None,
     **tsam_kwargs: Any,
 ) -> AggregationResult:
     """Aggregate an xarray DataArray using tsam.
@@ -89,13 +91,21 @@ def aggregate(
               extreme periods are identified only on the
               clustered-on columns.
 
+        dim_names: Names for the structural output dimensions
+            (``cluster``, ``timestep``, ``period``, ``segment``).
+            ``None`` (default) keeps today's names. Override to
+            avoid collisions with the caller's own dimension
+            names. See `DimNames`.
+
         **tsam_kwargs: Additional keyword arguments passed to
             ``tsam.aggregate()``.
     """
+    resolved_dim_names = dim_names if dim_names is not None else DimNames()
     _validate_time_dim(da, time_dim)
     col_dims = _resolve_cluster_dim(cluster_dim)
     slice_dims = _infer_slice_dims(da, time_dim, col_dims)
     _validate(da, time_dim, col_dims, slice_dims)
+    _validate_dim_names(resolved_dim_names, time_dim, col_dims, slice_dims)
     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)
@@ -111,6 +121,7 @@ def aggregate(
             per_dim_weights,
             active_coords,
             tsam_kwargs,
+            resolved_dim_names,
         )
 
     slice_coords = {d: da.coords[d].values for d in slice_dims}
@@ -129,6 +140,7 @@ def aggregate(
             per_dim_weights,
             active_coords,
             tsam_kwargs,
+            resolved_dim_names,
         )
         results.append(r)
 
@@ -255,7 +267,22 @@ def _validate(
             raise ValueError(msg)
 
 
-_RESERVED_DIMS = {"cluster", "timestep", "period"}
+def _validate_dim_names(
+    dim_names: DimNames,
+    time_dim: str,
+    col_dims: list[str],
+    slice_dims: list[str],
+) -> None:
+    """Reject output dimension names that collide with input dimensions."""
+    input_dims = {time_dim, *col_dims, *slice_dims}
+    conflict = input_dims & set(dim_names.as_tuple())
+    if conflict:
+        msg = (
+            f"Output dimension names {conflict} collide with input dimension "
+            "names. Rename them in your input DataArray, or pass a DimNames "
+            "with different names via dim_names=."
+        )
+        raise ValueError(msg)
 
 
 def _validate_data(
@@ -268,17 +295,6 @@ def _validate_data(
 
     Returns the (possibly computed) DataArray.
     """
-    # Reserved dimension names
-    all_user_dims = {time_dim, *col_dims, *slice_dims}
-    reserved_conflict = all_user_dims & _RESERVED_DIMS
-    if reserved_conflict:
-        msg = (
-            f"Dimension names {reserved_conflict} are reserved by "
-            "tsam_xarray for output dimensions. Rename them in "
-            "your input DataArray."
-        )
-        raise ValueError(msg)
-
     # Dask arrays — compute before other checks
     if hasattr(da.data, "dask"):
         warnings.warn(
@@ -368,6 +384,7 @@ def _to_dataframe(
 def _representatives_to_da(
     df: pd.DataFrame,
     col_dims: list[str],
+    dim_names: DimNames,
 ) -> xr.DataArray:
     """Convert cluster_representatives DataFrame to DataArray."""
     df = df.copy()
@@ -375,7 +392,7 @@ def _representatives_to_da(
     # Without: 2 levels: (cluster, timestep)
     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"]
+    df.index.names = [dim_names.cluster, dim_names.timestep]
 
     if not col_dims:
         clusters = df.index.get_level_values(0).unique()
@@ -383,8 +400,8 @@ def _representatives_to_da(
         values = df.values.squeeze(axis=1).reshape(len(clusters), len(timesteps))
         return xr.DataArray(
             values,
-            dims=["cluster", "timestep"],
-            coords={"cluster": clusters, "timestep": timesteps},
+            dims=[dim_names.cluster, dim_names.timestep],
+            coords={dim_names.cluster: clusters, dim_names.timestep: timesteps},
         )
 
     stacked = df.stack(df.columns.names, future_stack=True)
@@ -394,6 +411,7 @@ def _representatives_to_da(
 
 def _segment_durations_to_da(
     raw_durations: tuple[tuple[int, ...], ...] | None,
+    dim_names: DimNames,
 ) -> xr.DataArray | None:
     """Convert tsam segment_durations to DataArray."""
     if raw_durations is None:
@@ -401,10 +419,10 @@ def _segment_durations_to_da(
     data = np.array(raw_durations)  # (n_clusters, n_segments)
     return xr.DataArray(
         data,
-        dims=["cluster", "timestep"],
+        dims=[dim_names.cluster, dim_names.timestep],
         coords={
-            "cluster": np.arange(data.shape[0]),
-            "timestep": np.arange(data.shape[1]),
+            dim_names.cluster: np.arange(data.shape[0]),
+            dim_names.timestep: np.arange(data.shape[1]),
         },
     )
 
@@ -624,6 +642,7 @@ def _aggregate_single(
     weights: dict[str, dict[str, float]] | None,
     active_coords: dict[str, set[str]] | None,
     tsam_kwargs: dict[str, Any],
+    dim_names: DimNames,
 ) -> AggregationResult:
     """Run a single tsam aggregation on a DataArray.
 
@@ -667,7 +686,7 @@ def _aggregate_single(
             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)
+    return _result_from_tsam(tsam_result, da, df, time_dim, col_dims, dim_names)
 
 
 def _result_from_tsam(
@@ -676,9 +695,12 @@ def _result_from_tsam(
     df: pd.DataFrame,
     time_dim: str,
     col_dims: list[str],
+    dim_names: DimNames,
 ) -> AggregationResult:
     """Build an AggregationResult from a tsam aggregation result."""
-    typical = _representatives_to_da(tsam_result.cluster_representatives, col_dims)
+    typical = _representatives_to_da(
+        tsam_result.cluster_representatives, col_dims, dim_names
+    )
     reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims)
     reconstructed = reconstructed.transpose(*da.dims).reindex_like(da)
 
@@ -686,11 +708,13 @@ def _result_from_tsam(
     cluster_ids = np.array(sorted(cw.keys()))
     cluster_weights_da = xr.DataArray(
         np.array([cw[k] for k in cluster_ids]),
-        dims=["cluster"],
-        coords={"cluster": cluster_ids},
+        dims=[dim_names.cluster],
+        coords={dim_names.cluster: cluster_ids},
     )
 
-    assignments_da = xr.DataArray(tsam_result.cluster_assignments, dims=["period"])
+    assignments_da = xr.DataArray(
+        tsam_result.cluster_assignments, dims=[dim_names.period]
+    )
 
     col_names: list[str] | None = None
     if isinstance(df.columns, pd.MultiIndex):
@@ -709,7 +733,7 @@ def _result_from_tsam(
         ),
     )
 
-    seg_durations = _segment_durations_to_da(tsam_result.segment_durations)
+    seg_durations = _segment_durations_to_da(tsam_result.segment_durations, dim_names)
 
     from tsam_xarray._clustering import ClusteringResult
 
@@ -718,6 +742,7 @@ def _result_from_tsam(
         cluster_dim=col_dims,
         slice_dims=[],
         clusterings={(): tsam_result.clustering},
+        dim_names=dim_names,
     )
 
     return AggregationResult(
@@ -804,6 +829,7 @@ def _concat_results(
         cluster_dim=first.clustering.cluster_dim,
         slice_dims=slice_dims,
         clusterings=merged_clusterings,
+        dim_names=first.clustering.dim_names,
     )
 
     return AggregationResult(

tsam_xarray-0.6.4/src/tsam_xarray/_dim_names.py

@@ -0,0 +1,39 @@
+"""Configurable output dimension names for tsam_xarray."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DimNames:
+    """Names of the structural output dimensions produced by aggregation.
+
+    tsam_xarray adds four dimensions to its results that do not exist in the
+    input: the cluster/representative axis, the intra-period timestep axis,
+    the original-period axis (in ``cluster_assignments``), and the segment
+    axis (segmented runs). By default these are ``cluster``, ``timestep``,
+    ``period``, and ``segment``; override them when they would collide with
+    the caller's own dimension names.
+
+    Attributes:
+        cluster: Cluster/representative axis.
+        timestep: Intra-period timestep axis.
+        period: Original-period axis (in ``cluster_assignments``).
+        segment: Segment axis (segmented runs).
+    """
+
+    cluster: str = "cluster"
+    timestep: str = "timestep"
+    period: str = "period"
+    segment: str = "segment"
+
+    def __post_init__(self) -> None:
+        names = self.as_tuple()
+        if len(set(names)) != len(names):
+            msg = f"DimNames must be unique, got {names}"
+            raise ValueError(msg)
+
+    def as_tuple(self) -> tuple[str, str, str, str]:
+        """The four names as a tuple, in declaration order."""
+        return (self.cluster, self.timestep, self.period, self.segment)

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/src/tsam_xarray/_result.py

@@ -10,6 +10,7 @@ import xarray as xr
 
 if TYPE_CHECKING:
     from tsam_xarray._clustering import ClusteringResult
+    from tsam_xarray._dim_names import DimNames
 
 
 @dataclass(frozen=True, repr=False)
@@ -103,15 +104,20 @@ class AggregationResult:
             f"weighted_rmse={float(self.accuracy.weighted_rmse.mean()):.4f})"
         )
 
+    @property
+    def dim_names(self) -> DimNames:
+        """Names of the structural output dimensions. See `DimNames`."""
+        return self.clustering.dim_names
+
     @property
     def n_clusters(self) -> int:
         """Number of cluster representative clusters."""
-        return int(self.cluster_weights.sizes["cluster"])
+        return int(self.cluster_weights.sizes[self.dim_names.cluster])
 
     @property
     def n_timesteps_per_period(self) -> int:
         """Number of timesteps per cluster representative."""
-        return int(self.cluster_representatives.sizes["timestep"])
+        return int(self.cluster_representatives.sizes[self.dim_names.timestep])
 
     @property
     def n_segments(self) -> int | None:
@@ -264,6 +270,7 @@ class AggregationResult:
                 cluster_dim=self.clustering.cluster_dim,
                 slice_dims=[],
                 clusterings={(): cr},
+                dim_names=self.clustering.dim_names,
             ),
         )
 
@@ -272,4 +279,4 @@ class AggregationResult:
         from tsam_xarray._clustering import _disaggregate_single
 
         cr = self.clustering.clusterings[()]
-        return _disaggregate_single(cr, data)
+        return _disaggregate_single(cr, data, self.clustering.dim_names)

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/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.6.3'
-__version_tuple__ = version_tuple = (0, 6, 3)
+__version__ = version = '0.6.4'
+__version_tuple__ = version_tuple = (0, 6, 4)
 
 __commit_id__ = commit_id = None

{tsam_xarray-0.6.3 → tsam_xarray-0.6.4}/test/test_aggregate.py

@@ -833,7 +833,7 @@ class TestDataValidation:
             dims=["time", "cluster"],
             coords={"time": time, "cluster": ["a", "b"]},
         )
-        with pytest.raises(ValueError, match="reserved"):
+        with pytest.raises(ValueError, match="collide"):
             tsam_xarray.aggregate(
                 da, time_dim="time", cluster_dim="cluster", n_clusters=4
             )
@@ -845,7 +845,7 @@ class TestDataValidation:
             dims=["time", "timestep"],
             coords={"time": time, "timestep": [0, 1]},
         )
-        with pytest.raises(ValueError, match="reserved"):
+        with pytest.raises(ValueError, match="collide"):
             tsam_xarray.aggregate(
                 da, time_dim="time", cluster_dim="timestep", n_clusters=4
             )
@@ -861,7 +861,7 @@ class TestDataValidation:
                 "period": [0, 1],
             },
         )
-        with pytest.raises(ValueError, match="reserved"):
+        with pytest.raises(ValueError, match="collide"):
             tsam_xarray.aggregate(
                 da,
                 time_dim="time",
@@ -869,6 +869,102 @@ class TestDataValidation:
                 n_clusters=4,
             )
 
+
+class TestDimNames:
+    def _period_da(self) -> xr.DataArray:
+        time = pd.date_range("2020-01-01", periods=30 * 24, freq="h")
+        return xr.DataArray(
+            np.random.default_rng(42).random((len(time), 2, 2)),
+            dims=["time", "variable", "period"],
+            coords={"time": time, "variable": ["a", "b"], "period": [0, 1]},
+        )
+
+    def test_default_reproduces_current_names(self):
+        da = _make_da()
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable"
+        )
+        assert result.dim_names == tsam_xarray.DimNames()
+        assert set(result.cluster_representatives.dims) >= {"cluster", "timestep"}
+        assert set(result.cluster_assignments.dims) == {"period", "region"}
+
+    def test_custom_names_applied(self):
+        da = _make_da()
+        dn = tsam_xarray.DimNames(cluster="rep", timestep="intra", period="orig_period")
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable", dim_names=dn
+        )
+        assert set(result.cluster_representatives.dims) >= {"rep", "intra"}
+        assert "orig_period" in result.cluster_assignments.dims
+        assert set(result.cluster_weights.dims) == {"rep", "region"}
+        assert result.n_clusters == 4
+        assert result.n_timesteps_per_period == 24
+
+    def test_renamed_period_frees_input_dim(self):
+        da = self._period_da()
+        dn = tsam_xarray.DimNames(period="orig_period")
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable", dim_names=dn
+        )
+        assert "period" in result.cluster_representatives.dims
+        assert "orig_period" in result.cluster_assignments.dims
+
+    def test_collision_still_raises_with_partial_rename(self):
+        da = self._period_da()
+        dn = tsam_xarray.DimNames(cluster="rep")
+        with pytest.raises(ValueError, match="collide"):
+            tsam_xarray.aggregate(
+                da,
+                n_clusters=4,
+                time_dim="time",
+                cluster_dim="variable",
+                dim_names=dn,
+            )
+
+    def test_dim_names_must_be_unique(self):
+        with pytest.raises(ValueError, match="unique"):
+            tsam_xarray.DimNames(cluster="x", timestep="x")
+
+    def test_apply_reuses_stored_names(self):
+        da = _make_da()
+        dn = tsam_xarray.DimNames(cluster="rep", timestep="intra")
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable", dim_names=dn
+        )
+        applied = result.clustering.apply(da)
+        assert set(applied.cluster_representatives.dims) >= {"rep", "intra"}
+
+    def test_disaggregate_roundtrip_custom_names(self):
+        da = _make_da()
+        dn = tsam_xarray.DimNames(cluster="rep", timestep="intra")
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable", dim_names=dn
+        )
+        back = result.disaggregate(result.cluster_representatives)
+        assert "time" in back.dims
+
+    def test_json_roundtrip_preserves_names(self, tmp_path):
+        da = _make_da()
+        dn = tsam_xarray.DimNames(period="orig_period", cluster="rep")
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable", dim_names=dn
+        )
+        path = tmp_path / "clustering.json"
+        result.clustering.to_json(path)
+        loaded = tsam_xarray.ClusteringResult.from_json(path)
+        assert loaded.dim_names == dn
+        assert "orig_period" in loaded.cluster_assignments.dims
+
+    def test_from_dict_defaults_names_for_legacy_blob(self):
+        da = _make_da()
+        result = tsam_xarray.aggregate(
+            da, n_clusters=4, time_dim="time", cluster_dim="variable"
+        )
+        blob = result.clustering.to_dict()
+        del blob["dim_names"]
+        loaded = tsam_xarray.ClusteringResult.from_dict(blob)
+        assert loaded.dim_names == tsam_xarray.DimNames()
+
     def test_dask_array_warns(self):
         pytest.importorskip("dask")
         da = _make_da()

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

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