From c4a1915c956233b7eeb1d18643bb4e4eb9492207 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:48:55 +0200 Subject: [PATCH] refactor: drop clustering dim-rename layer via tsam-xarray DimNames tsam-xarray 0.6.4 makes output dimension names configurable through DimNames. Pass DimNames(period='original_cluster') to aggregate() so the original-period axis no longer collides with a real 'period' slice dim. This removes the entire rename/unrename adapter that existed only to work around the fixed 'period'/'cluster' output names: - cluster()/apply_clustering() no longer rename period->_period etc. in and back out, and apply_clustering() applies the stored ClusteringResult directly instead of rebuilding one with renamed slice dims. - _ReducedFlowSystemBuilder and Clustering lose _unrename_map/_unrename. The only renames left are the honest tsam<->flixopt boundary translations timestep->time / timestep->segment. Bumps the tsam_xarray pin to 0.6.4. All public clustering outputs (cluster_assignments, cluster_occurrences, original/reconstructed/residuals/accuracy, dim_names, disaggregate, the reduced FlowSystem's (cluster, time) vars and cluster_weight) keep identical dim names, shapes and values in single- and multi-period systems. BREAKING CHANGE: the serialized clustering format changed - slice dims are now stored as 'period' (was '_period') plus a new 'dim_names' key. Clustering artifacts saved to JSON/netCDF by an older flixopt no longer load for apply_clustering()/expand() in this version; re-run transform.cluster() to regenerate them. Fresh clustering and same-version round-trips are unaffected. Also raises the minimum tsam_xarray to 0.6.4. Co-Authored-By: Claude Opus 4.8 --- flixopt/clustering/base.py | 60 ++++----------------- flixopt/transform_accessor.py | 85 +++++++----------------------- pyproject.toml | 4 +- tests/test_clustering/test_base.py | 7 ++- 4 files changed, 37 insertions(+), 119 deletions(-) diff --git a/flixopt/clustering/base.py b/flixopt/clustering/base.py index f50122142..a6fc2dce7 100644 --- a/flixopt/clustering/base.py +++ b/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,7 +339,7 @@ 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: @@ -380,17 +347,12 @@ def accuracy(self) -> TsamXarrayAccuracyMetrics: 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. diff --git a/flixopt/transform_accessor.py b/flixopt/transform_accessor.py index 30fc12a9d..f8d079012 100644 --- a/flixopt/transform_accessor.py +++ b/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,20 +1355,6 @@ 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') @@ -1384,9 +1362,9 @@ def cluster( # 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: diff --git a/pyproject.toml b/pyproject.toml index d9f6bc18b..fbc374e19 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_clustering/test_base.py b/tests/test_clustering/test_base.py index f69de4cdf..a945c05d8 100644 --- a/tests/test_clustering/test_base.py +++ b/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'), )