From f058fdc3a50cb3ceb5006f4cba5213d070e4c915 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 23 Jul 2026 17:02:12 +0200 Subject: [PATCH] Revert "refactor: drop clustering dim-rename layer via tsam-xarray DimNames (#739)" This reverts commit cc56774d31e05d637ee40424ec64af276716cc2f. --- flixopt/clustering/base.py | 60 +++++++++++++++++---- flixopt/transform_accessor.py | 85 +++++++++++++++++++++++------- pyproject.toml | 4 +- tests/test_clustering/test_base.py | 7 +-- 4 files changed, 119 insertions(+), 37 deletions(-) diff --git a/flixopt/clustering/base.py b/flixopt/clustering/base.py index a6fc2dce7..f50122142 100644 --- a/flixopt/clustering/base.py +++ b/flixopt/clustering/base.py @@ -44,6 +44,8 @@ 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 ( @@ -67,6 +69,14 @@ 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 @@ -79,6 +89,17 @@ 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) # ========================================================================== @@ -116,10 +137,10 @@ def is_segmented(self) -> bool: @property def dim_names(self) -> list[str]: """Names of extra dimensions, e.g., ['period', 'scenario'].""" - return list(self._clustering_result.slice_dims) + return [self._unrename_map.get(d, d) for d in self._clustering_result.slice_dims] # ========================================================================== - # DataArray properties (delegated to ClusteringResult) + # DataArray properties (delegated to ClusteringResult with unrename) # ========================================================================== @property @@ -130,6 +151,11 @@ 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'] @@ -143,7 +169,7 @@ def cluster_occurrences(self) -> xr.DataArray: Returns: DataArray with dims [cluster, period?, scenario?]. """ - return self._clustering_result.cluster_occurrences + return self._unrename(self._clustering_result.cluster_occurrences) @property def segment_assignments(self) -> xr.DataArray | None: @@ -158,7 +184,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 result + return self._unrename(result) @property def segment_durations(self) -> xr.DataArray | None: @@ -173,7 +199,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 result + return self._unrename(result) # ========================================================================== # Methods @@ -202,7 +228,13 @@ 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) - return self._clustering_result.disaggregate(data) + # 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) def apply( self, @@ -318,7 +350,7 @@ def original(self) -> xr.DataArray: Only available before serialization. """ self._require_full_data('original') - return self._aggregation_result.original + return self._unrename(self._aggregation_result.original) @property def reconstructed(self) -> xr.DataArray: @@ -330,7 +362,8 @@ def reconstructed(self) -> xr.DataArray: Only available before serialization. """ self._require_full_data('reconstructed') - return self._aggregation_result.reconstructed.transpose(*self.original.dims) + da = self._unrename(self._aggregation_result.reconstructed) + return da.transpose(*self.original.dims) @property def residuals(self) -> xr.DataArray: @@ -339,7 +372,7 @@ def residuals(self) -> xr.DataArray: Only available before serialization. """ self._require_full_data('residuals') - return self._aggregation_result.residuals + return self._unrename(self._aggregation_result.residuals) @property def accuracy(self) -> TsamXarrayAccuracyMetrics: @@ -347,12 +380,17 @@ 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,)``). + (dims ``(*dim_names,)``). Dim names are un-renamed to match ``original``. Only available before serialization. """ + import dataclasses + self._require_full_data('accuracy') - return self._aggregation_result.accuracy + acc = self._aggregation_result.accuracy + return dataclasses.replace( + acc, **{f.name: self._unrename(getattr(acc, f.name)) for f in dataclasses.fields(acc)} + ) 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 f8d079012..30fc12a9d 100644 --- a/flixopt/transform_accessor.py +++ b/flixopt/transform_accessor.py @@ -48,11 +48,13 @@ 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 @@ -75,13 +77,18 @@ 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._agg_result.cluster_weights.rename('cluster_weight') + return self._unrename(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. @@ -90,10 +97,11 @@ 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 = {} - known_dims = {'cluster', 'timestep', 'period', 'scenario'} + # Exclude known dims (including renamed variants like _period, _cluster) + known_dims = {'cluster', 'timestep', 'period', 'scenario'} | set(self._unrename_map.keys()) unknown_dims = [d for d in representatives.dims if d not in known_dims] if len(unknown_dims) != 1: raise ValueError( @@ -109,7 +117,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)] = da + result[str(var_name)] = self._unrename(da) return result def build_segment_durations(self) -> xr.DataArray: @@ -128,7 +136,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 da.transpose('cluster', 'time', *other_dims).rename('timestep_duration') + return self._unrename(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. @@ -233,6 +241,7 @@ 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 @@ -1334,7 +1343,6 @@ 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: @@ -1355,6 +1363,20 @@ 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') @@ -1362,9 +1384,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('period') + extra_dims.append(rename_map.get('period', 'period')) if has_scenarios: - extra_dims.append('scenario') + extra_dims.append(rename_map.get('scenario', '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) @@ -1429,12 +1451,15 @@ 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) + builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt, unrename_map) return builder.build(ds) def apply_clustering( @@ -1496,23 +1521,47 @@ 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 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}) + # 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) - agg_result = clustering.clustering_result.apply(da_full) + # Rename back + if unrename_map: + ds = ds.rename(unrename_map) # Build and return the reduced FlowSystem - builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt) + builder = _ReducedFlowSystemBuilder(self._fs, agg_result, timesteps_per_cluster, dt, unrename_map) return builder.build(ds) def _validate_for_expansion(self) -> Clustering: diff --git a/pyproject.toml b/pyproject.toml index fbc374e19..d9f6bc18b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ tutorials = [ # Full feature set (everything except dev tools) full = [ - "tsam_xarray >= 0.6.4, < 1", # Time series aggregation for clustering (wraps tsam); 0.6.4 adds aggregate(dim_names=) + "tsam_xarray >= 0.6.1, < 1", # Time series aggregation for clustering (wraps tsam); 0.6.1 adds aggregate(cluster_on=) "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.4", # Time series aggregation for clustering (wraps tsam) + "tsam_xarray==0.6.1", # 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 a945c05d8..f69de4cdf 100644 --- a/tests/test_clustering/test_base.py +++ b/tests/test_clustering/test_base.py @@ -11,17 +11,12 @@ def _make_clustering_result(clusterings: dict, dim_names: list[str]): - """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. - """ + """Create a ClusteringResult from a dict of tsam ClusteringResult-like objects.""" return tsam_xarray.ClusteringResult( time_dim='time', cluster_dim=['variable'], slice_dims=dim_names, clusterings=clusterings, - dim_names=tsam_xarray.DimNames(period='original_cluster'), )