Change fit_to_model_coords to work with a Collection of dims

Author: FBumann

flixopt/flow_system.py

@@ -7,7 +7,7 @@
 import pathlib
 import warnings
 from io import StringIO
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Collection, Dict, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -349,6 +349,7 @@ def fit_to_model_coords(
         name: str,
         data: Optional[Union[TemporalDataUser, NonTemporalDataUser]],
         has_time_dim: bool = True,
+        dims: Optional[Collection[FlowSystemDimensions]] = None,
     ) -> Optional[Union[TemporalData, NonTemporalData]]:
         """
         Fit data to model coordinate system (currently time, but extensible).
@@ -357,17 +358,27 @@ def fit_to_model_coords(
             name: Name of the data
             data: Data to fit to model coordinates
             has_time_dim: Wether to use the time dimension or not
+            dims: Collection of dimension names to use for fitting. If None, all dimensions are used.
 
         Returns:
             xr.DataArray aligned to model coordinate system. If data is None, returns None.
         """
         if data is None:
             return None
 
-        coords = self.coords
+        if dims is None:
+            coords = self.coords
 
-        if not has_time_dim:
-            coords.pop('time')
+            if not has_time_dim:
+                warnings.warn(
+                    'has_time_dim is deprecated. Please pass dims to fit_to_model_coords instead.',
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                coords.pop('time')
+        else:
+            coords = self.coords
+            coords = {k: coords[k] for k in dims if k in coords}
 
         # Rest of your method stays the same, just pass coords
         if isinstance(data, TimeSeriesData):
@@ -390,6 +401,7 @@ def fit_effects_to_model_coords(
         effect_values: Optional[Union[TemporalEffectsUser, NonTemporalEffectsUser]],
         label_suffix: Optional[str] = None,
         has_time_dim: bool = True,
+        dims: Optional[Collection[FlowSystemDimensions]] = None,
     ) -> Optional[Union[TemporalEffects, NonTemporalEffects]]:
         """
         Transform EffectValues from the user to Internal Datatypes aligned with model coordinates.
@@ -401,7 +413,10 @@ def fit_effects_to_model_coords(
 
         return {
             effect: self.fit_to_model_coords(
-                '|'.join(filter(None, [label_prefix, effect, label_suffix])), value, has_time_dim=has_time_dim
+                '|'.join(filter(None, [label_prefix, effect, label_suffix])),
+                value,
+                has_time_dim=has_time_dim,
+                dims=dims,
             )
             for effect, value in effect_values_dict.items()
         }