refactor: drop clustering dim-rename layer via tsam-xarray DimNames

flixopt/clustering/base.py

@@ -44,8 +44,6 @@ def __init__(
         original_timesteps: pd.DatetimeIndex | list[str] | None = None,
         # Internal: tsam_xarray AggregationResult for full data access
         _aggregation_result: TsamXarrayAggregationResult | None = None,
-        # Internal: mapping from renamed dims back to originals (e.g., _period -> period)
-        _unrename_map: dict[str, str] | None = None,
     ):
         # Handle ISO timestamp strings from serialization
         if (
@@ -69,14 +67,6 @@ def __init__(
         else:
             raise ValueError('Either clustering_result or _aggregation_result must be provided')
 
-        # Resolve unrename_map: if not explicitly provided, infer from slice_dims
-        # (e.g., '_period' in slice_dims → {'_period': 'period'})
-        if _unrename_map:
-            self._unrename_map = _unrename_map
-        else:
-            known_renames = {'_period': 'period', '_cluster': 'cluster'}
-            self._unrename_map = {k: v for k, v in known_renames.items() if k in self._clustering_result.slice_dims}
-
         # Flag indicating this was loaded from serialization (missing full AggregationResult data)
         self._from_serialization = _aggregation_result is None
 
@@ -89,17 +79,6 @@ def _clustering_result_from_dict(d: dict) -> ClusteringResult:
 
         return ClusteringResultClass.from_dict(d)
 
-    # ==========================================================================
-    # Helper for dim unrenaming
-    # ==========================================================================
-
-    def _unrename(self, da: xr.DataArray) -> xr.DataArray:
-        """Rename tsam_xarray output dims back to original names (e.g., _period -> period)."""
-        if not self._unrename_map:
-            return da
-        renames = {k: v for k, v in self._unrename_map.items() if k in da.dims}
-        return da.rename(renames) if renames else da
-
     # ==========================================================================
     # Core properties (delegated to ClusteringResult)
     # ==========================================================================
@@ -137,10 +116,10 @@ def is_segmented(self) -> bool:
     @property
     def dim_names(self) -> list[str]:
         """Names of extra dimensions, e.g., ['period', 'scenario']."""
-        return [self._unrename_map.get(d, d) for d in self._clustering_result.slice_dims]
+        return list(self._clustering_result.slice_dims)
 
     # ==========================================================================
-    # DataArray properties (delegated to ClusteringResult with unrename)
+    # DataArray properties (delegated to ClusteringResult)
     # ==========================================================================
 
     @property
@@ -151,11 +130,6 @@ def cluster_assignments(self) -> xr.DataArray:
             DataArray with dims [original_cluster, period?, scenario?].
         """
         da = self._clustering_result.cluster_assignments
-        # Rename tsam_xarray's 'period' dim to our 'original_cluster' convention
-        # (must happen before _unrename to avoid conflict with _period → period rename)
-        if 'period' in da.dims:
-            da = da.rename({'period': 'original_cluster'})
-        da = self._unrename(da)
         # Ensure original_cluster is first dim (tsam_xarray puts slice dims first)
         if 'original_cluster' in da.dims and da.dims[0] != 'original_cluster':
             other_dims = [d for d in da.dims if d != 'original_cluster']
@@ -169,7 +143,7 @@ def cluster_occurrences(self) -> xr.DataArray:
         Returns:
             DataArray with dims [cluster, period?, scenario?].
         """
-        return self._unrename(self._clustering_result.cluster_occurrences)
+        return self._clustering_result.cluster_occurrences
 
     @property
     def segment_assignments(self) -> xr.DataArray | None:
@@ -184,7 +158,7 @@ def segment_assignments(self) -> xr.DataArray | None:
         # tsam_xarray uses 'timestep', we use 'time'
         if 'timestep' in result.dims:
             result = result.rename({'timestep': 'time'})
-        return self._unrename(result)
+        return result
 
     @property
     def segment_durations(self) -> xr.DataArray | None:
@@ -199,7 +173,7 @@ def segment_durations(self) -> xr.DataArray | None:
         # tsam_xarray uses 'timestep', we use 'segment'
         if 'timestep' in result.dims:
             result = result.rename({'timestep': 'segment'})
-        return self._unrename(result)
+        return result
 
     # ==========================================================================
     # Methods
@@ -228,13 +202,7 @@ def disaggregate(self, data: xr.DataArray) -> xr.DataArray:
         renames_to_tsam = {k: v for k, v in flixopt_to_tsam.items() if k in data.dims}
         if renames_to_tsam:
             data = data.rename(renames_to_tsam)
-        # Rename period/scenario dims to internal names (_period, _scenario)
-        reverse_unrename = {v: k for k, v in self._unrename_map.items()}
-        renames = {k: v for k, v in reverse_unrename.items() if k in data.dims}
-        if renames:
-            data = data.rename(renames)
-        result = self._clustering_result.disaggregate(data)
-        return self._unrename(result)
+        return self._clustering_result.disaggregate(data)
 
     def apply(
         self,
@@ -350,7 +318,7 @@ def original(self) -> xr.DataArray:
         Only available before serialization.
         """
         self._require_full_data('original')
-        return self._unrename(self._aggregation_result.original)
+        return self._aggregation_result.original
 
     @property
     def reconstructed(self) -> xr.DataArray:
@@ -362,8 +330,7 @@ def reconstructed(self) -> xr.DataArray:
         Only available before serialization.
         """
         self._require_full_data('reconstructed')
-        da = self._unrename(self._aggregation_result.reconstructed)
-        return da.transpose(*self.original.dims)
+        return self._aggregation_result.reconstructed.transpose(*self.original.dims)
 
     @property
     def residuals(self) -> xr.DataArray:
@@ -372,25 +339,20 @@ def residuals(self) -> xr.DataArray:
         Only available before serialization.
         """
         self._require_full_data('residuals')
-        return self._unrename(self._aggregation_result.residuals)
+        return self._aggregation_result.residuals
 
     @property
     def accuracy(self) -> TsamXarrayAccuracyMetrics:
         """tsam_xarray ``AccuracyMetrics`` (per-variable and column-weighted).
 
         Exposes ``rmse`` / ``mae`` / ``rmse_duration`` (dims ``(variable, *dim_names)``)
         and the aggregate ``weighted_rmse`` / ``weighted_mae`` / ``weighted_rmse_duration``
-        (dims ``(*dim_names,)``). Dim names are un-renamed to match ``original``.
+        (dims ``(*dim_names,)``).
 
         Only available before serialization.
         """
-        import dataclasses
-
         self._require_full_data('accuracy')
-        acc = self._aggregation_result.accuracy
-        return dataclasses.replace(
-            acc, **{f.name: self._unrename(getattr(acc, f.name)) for f in dataclasses.fields(acc)}
-        )
+        return self._aggregation_result.accuracy
 
     def compare(self, variable: str | list[str] | None = None) -> xr.Dataset:
         """Tidy original-vs-clustered comparison, ready for plotting.

flixopt/transform_accessor.py

@@ -48,13 +48,11 @@ def __init__(
         agg_result: Any,  # tsam_xarray.AggregationResult
         timesteps_per_cluster: int,
         dt: float,
-        unrename_map: dict[str, str] | None = None,
     ):
         self._fs = fs
         self._agg_result = agg_result
         self._timesteps_per_cluster = timesteps_per_cluster
         self._dt = dt
-        self._unrename_map = unrename_map or {}
 
         self._n_clusters = agg_result.n_clusters
         self._is_segmented = agg_result.n_segments is not None
@@ -77,18 +75,13 @@ def __init__(
 
         self._base_coords = {'cluster': self._cluster_coords, 'time': self._time_coords}
 
-    def _unrename(self, da: xr.DataArray) -> xr.DataArray:
-        """Rename tsam_xarray output dims back to original names (e.g., _period -> period)."""
-        renames = {k: v for k, v in self._unrename_map.items() if k in da.dims}
-        return da.rename(renames) if renames else da
-
     def build_cluster_weights(self) -> xr.DataArray:
         """Build cluster_weight DataArray from aggregation result.
 
         Returns:
             DataArray with dims [cluster, period?, scenario?].
         """
-        return self._unrename(self._agg_result.cluster_weights.rename('cluster_weight'))
+        return self._agg_result.cluster_weights.rename('cluster_weight')
 
     def build_typical_periods(self) -> dict[str, xr.DataArray]:
         """Build typical periods DataArrays with (cluster, time, ...) shape.
@@ -97,11 +90,10 @@ def build_typical_periods(self) -> dict[str, xr.DataArray]:
             Dict mapping column names to DataArrays.
         """
         representatives = self._agg_result.cluster_representatives
-        # representatives has dims: (cluster, timestep, variable, _period?, scenario?)
+        # representatives has dims: (cluster, timestep, variable, period?, scenario?)
         # We need to split by variable and rename timestep -> time
         result = {}
-        # Exclude known dims (including renamed variants like _period, _cluster)
-        known_dims = {'cluster', 'timestep', 'period', 'scenario'} | set(self._unrename_map.keys())
+        known_dims = {'cluster', 'timestep', 'period', 'scenario'}
         unknown_dims = [d for d in representatives.dims if d not in known_dims]
         if len(unknown_dims) != 1:
             raise ValueError(
@@ -117,7 +109,7 @@ def build_typical_periods(self) -> dict[str, xr.DataArray]:
             # Ensure cluster and time are first two dims
             other_dims = [d for d in da.dims if d not in ('cluster', 'time')]
             da = da.transpose('cluster', 'time', *other_dims)
-            result[str(var_name)] = self._unrename(da)
+            result[str(var_name)] = da
         return result
 
     def build_segment_durations(self) -> xr.DataArray:
@@ -136,7 +128,7 @@ def build_segment_durations(self) -> xr.DataArray:
         da = da.rename({'timestep': 'time'})
         da = da.assign_coords(cluster=self._cluster_coords, time=self._time_coords)
         other_dims = [d for d in da.dims if d not in ('cluster', 'time')]
-        return self._unrename(da.transpose('cluster', 'time', *other_dims).rename('timestep_duration'))
+        return da.transpose('cluster', 'time', *other_dims).rename('timestep_duration')
 
     def build_reduced_dataset(self, ds: xr.Dataset, typical_das: dict[str, xr.DataArray]) -> xr.Dataset:
         """Build the reduced dataset with (cluster, time) structure.
@@ -241,7 +233,6 @@ def build(self, ds: xr.Dataset) -> FlowSystem:
         reduced_fs.clustering = Clustering(
             original_timesteps=self._fs.timesteps,
             _aggregation_result=self._agg_result,
-            _unrename_map=self._unrename_map,
         )
 
         return reduced_fs
@@ -1343,6 +1334,7 @@ def cluster(
             'rescale_exclude_columns',
             'round_decimals',
             'numerical_tolerance',
+            'dim_names',  # set explicitly to keep the 'period' slice dim from colliding
         }
         conflicts = reserved_tsam_keys & set(tsam_kwargs.keys())
         if conflicts:
@@ -1363,30 +1355,16 @@ def cluster(
                     "ExtremeConfig(..., method='replace')"
                 )
 
-        # Rename reserved dimension names to avoid conflict with tsam_xarray
-        # tsam_xarray reserves: 'period', 'cluster', 'timestep'
-        reserved_renames = {'period': '_period', 'cluster': '_cluster'}
-        # Check against full ds dims (period/cluster may only exist as coords, not in ds_for_clustering)
-        rename_map = {k: v for k, v in reserved_renames.items() if k in ds.dims}
-        unrename_map = {v: k for k, v in rename_map.items()}
-
-        if rename_map:
-            # Only rename dims that exist in each dataset
-            clustering_renames = {k: v for k, v in rename_map.items() if k in ds_for_clustering.dims}
-            if clustering_renames:
-                ds_for_clustering = ds_for_clustering.rename(clustering_renames)
-            ds = ds.rename(rename_map)
-
         # Stack Dataset into a single DataArray with 'variable' dimension
         da_for_clustering = ds_for_clustering.to_dataarray(dim='variable')
 
         # Ensure period/scenario dimensions are present in the DataArray
         # even if the data doesn't vary across them (tsam_xarray needs them for slicing)
         extra_dims = []
         if has_periods:
-            extra_dims.append(rename_map.get('period', 'period'))
+            extra_dims.append('period')
         if has_scenarios:
-            extra_dims.append(rename_map.get('scenario', 'scenario'))
+            extra_dims.append('scenario')
         for dim_name in extra_dims:
             if dim_name not in da_for_clustering.dims and dim_name in ds.dims:
                 # Drop as non-dim coordinate first (to_dataarray may keep it as scalar coord)
@@ -1451,15 +1429,12 @@ def cluster(
                 n_clusters=n_clusters,
                 weights=weights,
                 cluster_on=cluster_on,
+                dim_names=tsam_xarray.DimNames(period='original_cluster'),
                 **tsam_kwargs_full,
             )
 
-        # Rename reserved dims back to original names in the dataset
-        if unrename_map:
-            ds = ds.rename(unrename_map)
-
         # Build and return the reduced FlowSystem
-        builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt, unrename_map)
+        builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt)
         return builder.build(ds)
 
     def apply_clustering(
@@ -1521,47 +1496,23 @@ def apply_clustering(
                 f'Ensure self._fs.timesteps matches the original data used for clustering.'
             )
 
-        # Rename reserved dimension names to avoid conflict with tsam_xarray
-        reserved_renames = {'period': '_period', 'cluster': '_cluster'}
-        rename_map = {k: v for k, v in reserved_renames.items() if k in ds.dims}
-        unrename_map = {v: k for k, v in rename_map.items()}
-
-        if rename_map:
-            ds = ds.rename(rename_map)
-
         # Apply existing clustering to full data
         logger.info('Applying clustering...')
         with warnings.catch_warnings():
             warnings.filterwarnings('ignore', category=UserWarning, message='.*minimal value.*exceeds.*')
             da_full = ds.to_dataarray(dim='variable')
 
-            # Ensure extra dims are present in DataArray
-            for _orig_name, renamed in rename_map.items():
-                if renamed not in da_full.dims and renamed in ds.dims:
-                    if renamed in da_full.coords:
-                        da_full = da_full.drop_vars(renamed)
-                    da_full = da_full.expand_dims({renamed: ds.coords[renamed].values})
-
-            # Get clustering result with correct dim names for the renamed data
-            from tsam_xarray import ClusteringResult as ClusteringResultClass
-
-            cr_result = clustering.clustering_result
-            # Map dim names to renamed versions (e.g., period → _period)
-            slice_dims = [rename_map.get(d, d) for d in clustering.dim_names]
-            cr_result = ClusteringResultClass(
-                time_dim='time',
-                cluster_dim=['variable'],
-                slice_dims=slice_dims,
-                clusterings=dict(cr_result.clusterings),
-            )
-            agg_result = cr_result.apply(da_full)
+            # Ensure slice dims (period/scenario) are present in the DataArray
+            for dim_name in clustering.dim_names:
+                if dim_name not in da_full.dims and dim_name in ds.dims:
+                    if dim_name in da_full.coords:
+                        da_full = da_full.drop_vars(dim_name)
+                    da_full = da_full.expand_dims({dim_name: ds.coords[dim_name].values})
 
-        # Rename back
-        if unrename_map:
-            ds = ds.rename(unrename_map)
+            agg_result = clustering.clustering_result.apply(da_full)
 
         # Build and return the reduced FlowSystem
-        builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt, unrename_map)
+        builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt)
         return builder.build(ds)
 
     def _validate_for_expansion(self) -> Clustering:

pyproject.toml

@@ -70,7 +70,7 @@ tutorials = [
 
 # Full feature set (everything except dev tools)
 full = [
-    "tsam_xarray >= 0.6.1, < 1",  # Time series aggregation for clustering (wraps tsam); 0.6.1 adds aggregate(cluster_on=)
+    "tsam_xarray >= 0.6.4, < 1",  # Time series aggregation for clustering (wraps tsam); 0.6.4 adds aggregate(dim_names=)
     "tsam >= 3.4.0, < 4",  # Directly imported for ClusterConfig, ExtremeConfig, SegmentConfig
     "pyvis==0.3.2",  # Visualizing FlowSystem Network
     "scipy >= 1.15.1, < 2", # Used by tsam. Prior versions have conflict with highspy. See https://github.com/scipy/scipy/issues/22257
@@ -87,7 +87,7 @@ full = [
 # Development tools and testing
 dev = [
     "xarray<2026.5",  # TODO: drop once linopy ships xarray 2026.3+ compat fix
-    "tsam_xarray==0.6.1",  # Time series aggregation for clustering (wraps tsam)
+    "tsam_xarray==0.6.4",  # Time series aggregation for clustering (wraps tsam)
     "tsam==3.4.0",  # Directly imported for ClusterConfig, ExtremeConfig, SegmentConfig
     "pytest==9.1.1",
     "pytest-xdist==3.8.0",

tests/test_clustering/test_base.py

@@ -11,12 +11,17 @@
 
 
 def _make_clustering_result(clusterings: dict, dim_names: list[str]):
-    """Create a ClusteringResult from a dict of tsam ClusteringResult-like objects."""
+    """Create a ClusteringResult from a dict of tsam ClusteringResult-like objects.
+
+    Mirrors how ``transform.cluster()`` builds results: the original-period axis
+    is named ``original_cluster`` so it never collides with a ``period`` slice dim.
+    """
     return tsam_xarray.ClusteringResult(
         time_dim='time',
         cluster_dim=['variable'],
         slice_dims=dim_names,
         clusterings=clusterings,
+        dim_names=tsam_xarray.DimNames(period='original_cluster'),
     )