flixopt 3.0.2__py3-none-any.whl → 3.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of flixopt might be problematic. Click here for more details.

flixopt/results.py CHANGED
@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Literal
10
10
  import linopy
11
11
  import numpy as np
12
12
  import pandas as pd
13
- import plotly
14
13
  import xarray as xr
15
14
  import yaml
16
15
 
@@ -20,6 +19,7 @@ from .flow_system import FlowSystem
20
19
 
21
20
  if TYPE_CHECKING:
22
21
  import matplotlib.pyplot as plt
22
+ import plotly
23
23
  import pyvis
24
24
 
25
25
  from .calculation import Calculation, SegmentedCalculation
@@ -195,8 +195,8 @@ class CalculationResults:
195
195
  if 'flow_system' in kwargs and flow_system_data is None:
196
196
  flow_system_data = kwargs.pop('flow_system')
197
197
  warnings.warn(
198
- "The 'flow_system' parameter is deprecated. Use 'flow_system_data' instead."
199
- "Acess is now by '.flow_system_data', while '.flow_system' returns the restored FlowSystem.",
198
+ "The 'flow_system' parameter is deprecated. Use 'flow_system_data' instead. "
199
+ "Access is now via '.flow_system_data', while '.flow_system' returns the restored FlowSystem.",
200
200
  DeprecationWarning,
201
201
  stacklevel=2,
202
202
  )
@@ -687,68 +687,117 @@ class CalculationResults:
687
687
 
688
688
  def plot_heatmap(
689
689
  self,
690
- variable_name: str,
691
- heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
692
- heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
693
- color_map: str = 'portland',
690
+ variable_name: str | list[str],
694
691
  save: bool | pathlib.Path = False,
695
692
  show: bool = True,
693
+ colors: plotting.ColorType = 'viridis',
696
694
  engine: plotting.PlottingEngine = 'plotly',
695
+ select: dict[FlowSystemDimensions, Any] | None = None,
696
+ facet_by: str | list[str] | None = 'scenario',
697
+ animate_by: str | None = 'period',
698
+ facet_cols: int = 3,
699
+ reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
700
+ | Literal['auto']
701
+ | None = 'auto',
702
+ fill: Literal['ffill', 'bfill'] | None = 'ffill',
703
+ # Deprecated parameters (kept for backwards compatibility)
697
704
  indexer: dict[FlowSystemDimensions, Any] | None = None,
705
+ heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] | None = None,
706
+ heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] | None = None,
707
+ color_map: str | None = None,
698
708
  ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
699
709
  """
700
- Plots a heatmap of the solution of a variable.
710
+ Plots a heatmap visualization of a variable using imshow or time-based reshaping.
711
+
712
+ Supports multiple visualization features that can be combined:
713
+ - **Multi-variable**: Plot multiple variables on a single heatmap (creates 'variable' dimension)
714
+ - **Time reshaping**: Converts 'time' dimension into 2D (e.g., hours vs days)
715
+ - **Faceting**: Creates subplots for different dimension values
716
+ - **Animation**: Animates through dimension values (Plotly only)
701
717
 
702
718
  Args:
703
- variable_name: The name of the variable to plot.
704
- heatmap_timeframes: The timeframes to use for the heatmap.
705
- heatmap_timesteps_per_frame: The timesteps per frame to use for the heatmap.
706
- color_map: The color map to use for the heatmap.
719
+ variable_name: The name of the variable to plot, or a list of variable names.
720
+ When a list is provided, variables are combined into a single DataArray
721
+ with a new 'variable' dimension.
707
722
  save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
708
723
  show: Whether to show the plot or not.
724
+ colors: Color scheme for the heatmap. See `flixopt.plotting.ColorType` for options.
709
725
  engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
710
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
711
- If None, uses first value for each dimension.
712
- If empty dict {}, uses all values.
726
+ select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
727
+ Applied BEFORE faceting/animation/reshaping.
728
+ facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
729
+ or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
730
+ animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
731
+ dimension values. Only one dimension can be animated. Ignored if not found.
732
+ facet_cols: Number of columns in the facet grid layout (default: 3).
733
+ reshape_time: Time reshaping configuration (default: 'auto'):
734
+ - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
735
+ - Tuple: Explicit reshaping, e.g. ('D', 'h') for days vs hours,
736
+ ('MS', 'D') for months vs days, ('W', 'h') for weeks vs hours
737
+ - None: Disable auto-reshaping (will error if only 1D time data)
738
+ Supported timeframes: 'YS', 'MS', 'W', 'D', 'h', '15min', 'min'
739
+ fill: Method to fill missing values after reshape: 'ffill' (forward fill) or 'bfill' (backward fill).
740
+ Default is 'ffill'.
713
741
 
714
742
  Examples:
715
- Basic usage (uses first scenario, first period, all time):
743
+ Direct imshow mode (default):
744
+
745
+ >>> results.plot_heatmap('Battery|charge_state', select={'scenario': 'base'})
746
+
747
+ Facet by scenario:
748
+
749
+ >>> results.plot_heatmap('Boiler(Qth)|flow_rate', facet_by='scenario', facet_cols=2)
716
750
 
717
- >>> results.plot_heatmap('Battery|charge_state')
751
+ Animate by period:
718
752
 
719
- Select specific scenario and period:
753
+ >>> results.plot_heatmap('Boiler(Qth)|flow_rate', select={'scenario': 'base'}, animate_by='period')
720
754
 
721
- >>> results.plot_heatmap('Boiler(Qth)|flow_rate', indexer={'scenario': 'base', 'period': 2024})
755
+ Time reshape mode - daily patterns:
722
756
 
723
- Time filtering (summer months only):
757
+ >>> results.plot_heatmap('Boiler(Qth)|flow_rate', select={'scenario': 'base'}, reshape_time=('D', 'h'))
758
+
759
+ Combined: time reshaping with faceting and animation:
724
760
 
725
761
  >>> results.plot_heatmap(
726
- ... 'Boiler(Qth)|flow_rate',
727
- ... indexer={
728
- ... 'scenario': 'base',
729
- ... 'time': results.solution.time[results.solution.time.dt.month.isin([6, 7, 8])],
730
- ... },
762
+ ... 'Boiler(Qth)|flow_rate', facet_by='scenario', animate_by='period', reshape_time=('D', 'h')
731
763
  ... )
732
764
 
733
- Save to specific location:
765
+ Multi-variable heatmap (variables as one axis):
734
766
 
735
767
  >>> results.plot_heatmap(
736
- ... 'Boiler(Qth)|flow_rate', indexer={'scenario': 'base'}, save='path/to/my_heatmap.html'
768
+ ... ['Boiler(Q_th)|flow_rate', 'CHP(Q_th)|flow_rate', 'HeatStorage|charge_state'],
769
+ ... select={'scenario': 'base', 'period': 1},
770
+ ... reshape_time=None,
737
771
  ... )
738
- """
739
- dataarray = self.solution[variable_name]
740
772
 
773
+ Multi-variable with time reshaping:
774
+
775
+ >>> results.plot_heatmap(
776
+ ... ['Boiler(Q_th)|flow_rate', 'CHP(Q_th)|flow_rate'],
777
+ ... facet_by='scenario',
778
+ ... animate_by='period',
779
+ ... reshape_time=('D', 'h'),
780
+ ... )
781
+ """
782
+ # Delegate to module-level plot_heatmap function
741
783
  return plot_heatmap(
742
- dataarray=dataarray,
743
- name=variable_name,
784
+ data=self.solution[variable_name],
785
+ name=variable_name if isinstance(variable_name, str) else None,
744
786
  folder=self.folder,
745
- heatmap_timeframes=heatmap_timeframes,
746
- heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
747
- color_map=color_map,
787
+ colors=colors,
748
788
  save=save,
749
789
  show=show,
750
790
  engine=engine,
791
+ select=select,
792
+ facet_by=facet_by,
793
+ animate_by=animate_by,
794
+ facet_cols=facet_cols,
795
+ reshape_time=reshape_time,
796
+ fill=fill,
751
797
  indexer=indexer,
798
+ heatmap_timeframes=heatmap_timeframes,
799
+ heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
800
+ color_map=color_map,
752
801
  )
753
802
 
754
803
  def plot_network(
@@ -920,51 +969,132 @@ class _NodeResults(_ElementResults):
920
969
  show: bool = True,
921
970
  colors: plotting.ColorType = 'viridis',
922
971
  engine: plotting.PlottingEngine = 'plotly',
923
- indexer: dict[FlowSystemDimensions, Any] | None = None,
924
- mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
925
- style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
972
+ select: dict[FlowSystemDimensions, Any] | None = None,
973
+ unit_type: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
974
+ mode: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
926
975
  drop_suffix: bool = True,
976
+ facet_by: str | list[str] | None = 'scenario',
977
+ animate_by: str | None = 'period',
978
+ facet_cols: int = 3,
979
+ # Deprecated parameter (kept for backwards compatibility)
980
+ indexer: dict[FlowSystemDimensions, Any] | None = None,
927
981
  ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
928
982
  """
929
- Plots the node balance of the Component or Bus.
983
+ Plots the node balance of the Component or Bus with optional faceting and animation.
984
+
930
985
  Args:
931
986
  save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
932
987
  show: Whether to show the plot or not.
933
988
  colors: The colors to use for the plot. See `flixopt.plotting.ColorType` for options.
934
989
  engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
935
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
936
- If None, uses first value for each dimension (except time).
937
- If empty dict {}, uses all values.
938
- style: The style to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
990
+ select: Optional data selection dict. Supports:
991
+ - Single values: {'scenario': 'base', 'period': 2024}
992
+ - Multiple values: {'scenario': ['base', 'high', 'renewable']}
993
+ - Slices: {'time': slice('2024-01', '2024-06')}
994
+ - Index arrays: {'time': time_array}
995
+ Note: Applied BEFORE faceting/animation.
996
+ unit_type: The unit type to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
939
997
  - 'flow_rate': Returns the flow_rates of the Node.
940
998
  - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
999
+ mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for stepped lines, or 'area' for stacked area charts.
941
1000
  drop_suffix: Whether to drop the suffix from the variable names.
1001
+ facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
1002
+ or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
1003
+ Example: 'scenario' creates one subplot per scenario.
1004
+ Example: ['scenario', 'period'] creates a grid of subplots for each scenario-period combination.
1005
+ animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
1006
+ dimension values. Only one dimension can be animated. Ignored if not found.
1007
+ facet_cols: Number of columns in the facet grid layout (default: 3).
1008
+
1009
+ Examples:
1010
+ Basic plot (current behavior):
1011
+
1012
+ >>> results['Boiler'].plot_node_balance()
1013
+
1014
+ Facet by scenario:
1015
+
1016
+ >>> results['Boiler'].plot_node_balance(facet_by='scenario', facet_cols=2)
1017
+
1018
+ Animate by period:
1019
+
1020
+ >>> results['Boiler'].plot_node_balance(animate_by='period')
1021
+
1022
+ Facet by scenario AND animate by period:
1023
+
1024
+ >>> results['Boiler'].plot_node_balance(facet_by='scenario', animate_by='period')
1025
+
1026
+ Select single scenario, then facet by period:
1027
+
1028
+ >>> results['Boiler'].plot_node_balance(select={'scenario': 'base'}, facet_by='period')
1029
+
1030
+ Select multiple scenarios and facet by them:
1031
+
1032
+ >>> results['Boiler'].plot_node_balance(
1033
+ ... select={'scenario': ['base', 'high', 'renewable']}, facet_by='scenario'
1034
+ ... )
1035
+
1036
+ Time range selection (summer months only):
1037
+
1038
+ >>> results['Boiler'].plot_node_balance(select={'time': slice('2024-06', '2024-08')}, facet_by='scenario')
942
1039
  """
943
- ds = self.node_balance(with_last_timestep=True, mode=mode, drop_suffix=drop_suffix, indexer=indexer)
1040
+ # Handle deprecated indexer parameter
1041
+ if indexer is not None:
1042
+ # Check for conflict with new parameter
1043
+ if select is not None:
1044
+ raise ValueError(
1045
+ "Cannot use both deprecated parameter 'indexer' and new parameter 'select'. Use only 'select'."
1046
+ )
1047
+
1048
+ import warnings
1049
+
1050
+ warnings.warn(
1051
+ "The 'indexer' parameter is deprecated and will be removed in a future version. Use 'select' instead.",
1052
+ DeprecationWarning,
1053
+ stacklevel=2,
1054
+ )
1055
+ select = indexer
944
1056
 
945
- ds, suffix_parts = _apply_indexer_to_data(ds, indexer, drop=True)
1057
+ if engine not in {'plotly', 'matplotlib'}:
1058
+ raise ValueError(f'Engine "{engine}" not supported. Use one of ["plotly", "matplotlib"]')
1059
+
1060
+ # Don't pass select/indexer to node_balance - we'll apply it afterwards
1061
+ ds = self.node_balance(with_last_timestep=True, unit_type=unit_type, drop_suffix=drop_suffix)
1062
+
1063
+ ds, suffix_parts = _apply_selection_to_data(ds, select=select, drop=True)
1064
+
1065
+ # Matplotlib requires only 'time' dimension; check for extras after selection
1066
+ if engine == 'matplotlib':
1067
+ extra_dims = [d for d in ds.dims if d != 'time']
1068
+ if extra_dims:
1069
+ raise ValueError(
1070
+ f'Matplotlib engine only supports a single time axis, but found extra dimensions: {extra_dims}. '
1071
+ f'Please use select={{...}} to reduce dimensions or switch to engine="plotly" for faceting/animation.'
1072
+ )
946
1073
  suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
947
1074
 
948
- title = f'{self.label} (flow rates){suffix}' if mode == 'flow_rate' else f'{self.label} (flow hours){suffix}'
1075
+ title = (
1076
+ f'{self.label} (flow rates){suffix}' if unit_type == 'flow_rate' else f'{self.label} (flow hours){suffix}'
1077
+ )
949
1078
 
950
1079
  if engine == 'plotly':
951
1080
  figure_like = plotting.with_plotly(
952
- ds.to_dataframe(),
1081
+ ds,
1082
+ facet_by=facet_by,
1083
+ animate_by=animate_by,
953
1084
  colors=colors,
954
- style=style,
1085
+ mode=mode,
955
1086
  title=title,
1087
+ facet_cols=facet_cols,
956
1088
  )
957
1089
  default_filetype = '.html'
958
- elif engine == 'matplotlib':
1090
+ else:
959
1091
  figure_like = plotting.with_matplotlib(
960
1092
  ds.to_dataframe(),
961
1093
  colors=colors,
962
- style=style,
1094
+ mode=mode,
963
1095
  title=title,
964
1096
  )
965
1097
  default_filetype = '.png'
966
- else:
967
- raise ValueError(f'Engine "{engine}" not supported. Use "plotly" or "matplotlib"')
968
1098
 
969
1099
  return plotting.export_figure(
970
1100
  figure_like=figure_like,
@@ -983,9 +1113,19 @@ class _NodeResults(_ElementResults):
983
1113
  save: bool | pathlib.Path = False,
984
1114
  show: bool = True,
985
1115
  engine: plotting.PlottingEngine = 'plotly',
1116
+ select: dict[FlowSystemDimensions, Any] | None = None,
1117
+ # Deprecated parameter (kept for backwards compatibility)
986
1118
  indexer: dict[FlowSystemDimensions, Any] | None = None,
987
1119
  ) -> plotly.graph_objs.Figure | tuple[plt.Figure, list[plt.Axes]]:
988
1120
  """Plot pie chart of flow hours distribution.
1121
+
1122
+ Note:
1123
+ Pie charts require scalar data (no extra dimensions beyond time).
1124
+ If your data has dimensions like 'scenario' or 'period', either:
1125
+
1126
+ - Use `select` to choose specific values: `select={'scenario': 'base', 'period': 2024}`
1127
+ - Let auto-selection choose the first value (a warning will be logged)
1128
+
989
1129
  Args:
990
1130
  lower_percentage_group: Percentage threshold for "Others" grouping.
991
1131
  colors: Color scheme. Also see plotly.
@@ -993,10 +1133,35 @@ class _NodeResults(_ElementResults):
993
1133
  save: Whether to save plot.
994
1134
  show: Whether to display plot.
995
1135
  engine: Plotting engine ('plotly' or 'matplotlib').
996
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
997
- If None, uses first value for each dimension.
998
- If empty dict {}, uses all values.
1136
+ select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
1137
+ Use this to select specific scenario/period before creating the pie chart.
1138
+
1139
+ Examples:
1140
+ Basic usage (auto-selects first scenario/period if present):
1141
+
1142
+ >>> results['Bus'].plot_node_balance_pie()
1143
+
1144
+ Explicitly select a scenario and period:
1145
+
1146
+ >>> results['Bus'].plot_node_balance_pie(select={'scenario': 'high_demand', 'period': 2030})
999
1147
  """
1148
+ # Handle deprecated indexer parameter
1149
+ if indexer is not None:
1150
+ # Check for conflict with new parameter
1151
+ if select is not None:
1152
+ raise ValueError(
1153
+ "Cannot use both deprecated parameter 'indexer' and new parameter 'select'. Use only 'select'."
1154
+ )
1155
+
1156
+ import warnings
1157
+
1158
+ warnings.warn(
1159
+ "The 'indexer' parameter is deprecated and will be removed in a future version. Use 'select' instead.",
1160
+ DeprecationWarning,
1161
+ stacklevel=2,
1162
+ )
1163
+ select = indexer
1164
+
1000
1165
  inputs = sanitize_dataset(
1001
1166
  ds=self.solution[self.inputs] * self._calculation_results.hours_per_timestep,
1002
1167
  threshold=1e-5,
@@ -1012,15 +1177,46 @@ class _NodeResults(_ElementResults):
1012
1177
  drop_suffix='|',
1013
1178
  )
1014
1179
 
1015
- inputs, suffix_parts = _apply_indexer_to_data(inputs, indexer, drop=True)
1016
- outputs, suffix_parts = _apply_indexer_to_data(outputs, indexer, drop=True)
1017
- suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
1018
-
1019
- title = f'{self.label} (total flow hours){suffix}'
1180
+ inputs, suffix_parts = _apply_selection_to_data(inputs, select=select, drop=True)
1181
+ outputs, suffix_parts = _apply_selection_to_data(outputs, select=select, drop=True)
1020
1182
 
1183
+ # Sum over time dimension
1021
1184
  inputs = inputs.sum('time')
1022
1185
  outputs = outputs.sum('time')
1023
1186
 
1187
+ # Auto-select first value for any remaining dimensions (scenario, period, etc.)
1188
+ # Pie charts need scalar data, so we automatically reduce extra dimensions
1189
+ extra_dims_inputs = [dim for dim in inputs.dims if dim != 'time']
1190
+ extra_dims_outputs = [dim for dim in outputs.dims if dim != 'time']
1191
+ extra_dims = list(set(extra_dims_inputs + extra_dims_outputs))
1192
+
1193
+ if extra_dims:
1194
+ auto_select = {}
1195
+ for dim in extra_dims:
1196
+ # Get first value of this dimension
1197
+ if dim in inputs.coords:
1198
+ first_val = inputs.coords[dim].values[0]
1199
+ elif dim in outputs.coords:
1200
+ first_val = outputs.coords[dim].values[0]
1201
+ else:
1202
+ continue
1203
+ auto_select[dim] = first_val
1204
+ logger.info(
1205
+ f'Pie chart auto-selected {dim}={first_val} (first value). '
1206
+ f'Use select={{"{dim}": value}} to choose a different value.'
1207
+ )
1208
+
1209
+ # Apply auto-selection
1210
+ inputs = inputs.sel(auto_select)
1211
+ outputs = outputs.sel(auto_select)
1212
+
1213
+ # Update suffix with auto-selected values
1214
+ auto_suffix_parts = [f'{dim}={val}' for dim, val in auto_select.items()]
1215
+ suffix_parts.extend(auto_suffix_parts)
1216
+
1217
+ suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
1218
+ title = f'{self.label} (total flow hours){suffix}'
1219
+
1024
1220
  if engine == 'plotly':
1025
1221
  figure_like = plotting.dual_pie_with_plotly(
1026
1222
  data_left=inputs.to_pandas(),
@@ -1063,8 +1259,10 @@ class _NodeResults(_ElementResults):
1063
1259
  negate_outputs: bool = False,
1064
1260
  threshold: float | None = 1e-5,
1065
1261
  with_last_timestep: bool = False,
1066
- mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
1262
+ unit_type: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
1067
1263
  drop_suffix: bool = False,
1264
+ select: dict[FlowSystemDimensions, Any] | None = None,
1265
+ # Deprecated parameter (kept for backwards compatibility)
1068
1266
  indexer: dict[FlowSystemDimensions, Any] | None = None,
1069
1267
  ) -> xr.Dataset:
1070
1268
  """
@@ -1074,14 +1272,29 @@ class _NodeResults(_ElementResults):
1074
1272
  negate_outputs: Whether to negate the output flow_rates of the Node.
1075
1273
  threshold: The threshold for small values. Variables with all values below the threshold are dropped.
1076
1274
  with_last_timestep: Whether to include the last timestep in the dataset.
1077
- mode: The mode to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
1275
+ unit_type: The unit type to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
1078
1276
  - 'flow_rate': Returns the flow_rates of the Node.
1079
1277
  - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
1080
1278
  drop_suffix: Whether to drop the suffix from the variable names.
1081
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
1082
- If None, uses first value for each dimension.
1083
- If empty dict {}, uses all values.
1279
+ select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
1084
1280
  """
1281
+ # Handle deprecated indexer parameter
1282
+ if indexer is not None:
1283
+ # Check for conflict with new parameter
1284
+ if select is not None:
1285
+ raise ValueError(
1286
+ "Cannot use both deprecated parameter 'indexer' and new parameter 'select'. Use only 'select'."
1287
+ )
1288
+
1289
+ import warnings
1290
+
1291
+ warnings.warn(
1292
+ "The 'indexer' parameter is deprecated and will be removed in a future version. Use 'select' instead.",
1293
+ DeprecationWarning,
1294
+ stacklevel=2,
1295
+ )
1296
+ select = indexer
1297
+
1085
1298
  ds = self.solution[self.inputs + self.outputs]
1086
1299
 
1087
1300
  ds = sanitize_dataset(
@@ -1100,9 +1313,9 @@ class _NodeResults(_ElementResults):
1100
1313
  drop_suffix='|' if drop_suffix else None,
1101
1314
  )
1102
1315
 
1103
- ds, _ = _apply_indexer_to_data(ds, indexer, drop=True)
1316
+ ds, _ = _apply_selection_to_data(ds, select=select, drop=True)
1104
1317
 
1105
- if mode == 'flow_hours':
1318
+ if unit_type == 'flow_hours':
1106
1319
  ds = ds * self._calculation_results.hours_per_timestep
1107
1320
  ds = ds.rename_vars({var: var.replace('flow_rate', 'flow_hours') for var in ds.data_vars})
1108
1321
 
@@ -1137,69 +1350,162 @@ class ComponentResults(_NodeResults):
1137
1350
  show: bool = True,
1138
1351
  colors: plotting.ColorType = 'viridis',
1139
1352
  engine: plotting.PlottingEngine = 'plotly',
1140
- style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
1353
+ mode: Literal['area', 'stacked_bar', 'line'] = 'area',
1354
+ select: dict[FlowSystemDimensions, Any] | None = None,
1355
+ facet_by: str | list[str] | None = 'scenario',
1356
+ animate_by: str | None = 'period',
1357
+ facet_cols: int = 3,
1358
+ # Deprecated parameter (kept for backwards compatibility)
1141
1359
  indexer: dict[FlowSystemDimensions, Any] | None = None,
1142
1360
  ) -> plotly.graph_objs.Figure:
1143
- """Plot storage charge state over time, combined with the node balance.
1361
+ """Plot storage charge state over time, combined with the node balance with optional faceting and animation.
1144
1362
 
1145
1363
  Args:
1146
1364
  save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
1147
1365
  show: Whether to show the plot or not.
1148
1366
  colors: Color scheme. Also see plotly.
1149
1367
  engine: Plotting engine to use. Only 'plotly' is implemented atm.
1150
- style: The colors to use for the plot. See `flixopt.plotting.ColorType` for options.
1151
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
1152
- If None, uses first value for each dimension.
1153
- If empty dict {}, uses all values.
1368
+ mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for stepped lines, or 'area' for stacked area charts.
1369
+ select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
1370
+ Applied BEFORE faceting/animation.
1371
+ facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
1372
+ or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
1373
+ animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
1374
+ dimension values. Only one dimension can be animated. Ignored if not found.
1375
+ facet_cols: Number of columns in the facet grid layout (default: 3).
1154
1376
 
1155
1377
  Raises:
1156
1378
  ValueError: If component is not a storage.
1379
+
1380
+ Examples:
1381
+ Basic plot:
1382
+
1383
+ >>> results['Storage'].plot_charge_state()
1384
+
1385
+ Facet by scenario:
1386
+
1387
+ >>> results['Storage'].plot_charge_state(facet_by='scenario', facet_cols=2)
1388
+
1389
+ Animate by period:
1390
+
1391
+ >>> results['Storage'].plot_charge_state(animate_by='period')
1392
+
1393
+ Facet by scenario AND animate by period:
1394
+
1395
+ >>> results['Storage'].plot_charge_state(facet_by='scenario', animate_by='period')
1157
1396
  """
1397
+ # Handle deprecated indexer parameter
1398
+ if indexer is not None:
1399
+ # Check for conflict with new parameter
1400
+ if select is not None:
1401
+ raise ValueError(
1402
+ "Cannot use both deprecated parameter 'indexer' and new parameter 'select'. Use only 'select'."
1403
+ )
1404
+
1405
+ import warnings
1406
+
1407
+ warnings.warn(
1408
+ "The 'indexer' parameter is deprecated and will be removed in a future version. Use 'select' instead.",
1409
+ DeprecationWarning,
1410
+ stacklevel=2,
1411
+ )
1412
+ select = indexer
1413
+
1158
1414
  if not self.is_storage:
1159
1415
  raise ValueError(f'Cant plot charge_state. "{self.label}" is not a storage')
1160
1416
 
1161
- ds = self.node_balance(with_last_timestep=True, indexer=indexer)
1162
- charge_state = self.charge_state
1417
+ # Get node balance and charge state
1418
+ ds = self.node_balance(with_last_timestep=True)
1419
+ charge_state_da = self.charge_state
1163
1420
 
1164
- ds, suffix_parts = _apply_indexer_to_data(ds, indexer, drop=True)
1165
- charge_state, suffix_parts = _apply_indexer_to_data(charge_state, indexer, drop=True)
1421
+ # Apply select filtering
1422
+ ds, suffix_parts = _apply_selection_to_data(ds, select=select, drop=True)
1423
+ charge_state_da, _ = _apply_selection_to_data(charge_state_da, select=select, drop=True)
1166
1424
  suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
1167
1425
 
1168
1426
  title = f'Operation Balance of {self.label}{suffix}'
1169
1427
 
1170
1428
  if engine == 'plotly':
1171
- fig = plotting.with_plotly(
1172
- ds.to_dataframe(),
1429
+ # Plot flows (node balance) with the specified mode
1430
+ figure_like = plotting.with_plotly(
1431
+ ds,
1432
+ facet_by=facet_by,
1433
+ animate_by=animate_by,
1173
1434
  colors=colors,
1174
- style=style,
1435
+ mode=mode,
1175
1436
  title=title,
1437
+ facet_cols=facet_cols,
1176
1438
  )
1177
1439
 
1178
- # TODO: Use colors for charge state?
1440
+ # Create a dataset with just charge_state and plot it as lines
1441
+ # This ensures proper handling of facets and animation
1442
+ charge_state_ds = charge_state_da.to_dataset(name=self._charge_state)
1179
1443
 
1180
- charge_state = charge_state.to_dataframe()
1181
- fig.add_trace(
1182
- plotly.graph_objs.Scatter(
1183
- x=charge_state.index, y=charge_state.values.flatten(), mode='lines', name=self._charge_state
1184
- )
1444
+ # Plot charge_state with mode='line' to get Scatter traces
1445
+ charge_state_fig = plotting.with_plotly(
1446
+ charge_state_ds,
1447
+ facet_by=facet_by,
1448
+ animate_by=animate_by,
1449
+ colors=colors,
1450
+ mode='line', # Always line for charge_state
1451
+ title='', # No title needed for this temp figure
1452
+ facet_cols=facet_cols,
1185
1453
  )
1454
+
1455
+ # Add charge_state traces to the main figure
1456
+ # This preserves subplot assignments and animation frames
1457
+ for trace in charge_state_fig.data:
1458
+ trace.line.width = 2 # Make charge_state line more prominent
1459
+ trace.line.shape = 'linear' # Smooth line for charge state (not stepped like flows)
1460
+ figure_like.add_trace(trace)
1461
+
1462
+ # Also add traces from animation frames if they exist
1463
+ # Both figures use the same animate_by parameter, so they should have matching frames
1464
+ if hasattr(charge_state_fig, 'frames') and charge_state_fig.frames:
1465
+ # Add charge_state traces to each frame
1466
+ for i, frame in enumerate(charge_state_fig.frames):
1467
+ if i < len(figure_like.frames):
1468
+ for trace in frame.data:
1469
+ trace.line.width = 2
1470
+ trace.line.shape = 'linear' # Smooth line for charge state
1471
+ figure_like.frames[i].data = figure_like.frames[i].data + (trace,)
1472
+
1473
+ default_filetype = '.html'
1186
1474
  elif engine == 'matplotlib':
1475
+ # Matplotlib requires only 'time' dimension; check for extras after selection
1476
+ extra_dims = [d for d in ds.dims if d != 'time']
1477
+ if extra_dims:
1478
+ raise ValueError(
1479
+ f'Matplotlib engine only supports a single time axis, but found extra dimensions: {extra_dims}. '
1480
+ f'Please use select={{...}} to reduce dimensions or switch to engine="plotly" for faceting/animation.'
1481
+ )
1482
+ # For matplotlib, plot flows (node balance), then add charge_state as line
1187
1483
  fig, ax = plotting.with_matplotlib(
1188
1484
  ds.to_dataframe(),
1189
1485
  colors=colors,
1190
- style=style,
1486
+ mode=mode,
1191
1487
  title=title,
1192
1488
  )
1193
1489
 
1194
- charge_state = charge_state.to_dataframe()
1195
- ax.plot(charge_state.index, charge_state.values.flatten(), label=self._charge_state)
1490
+ # Add charge_state as a line overlay
1491
+ charge_state_df = charge_state_da.to_dataframe()
1492
+ ax.plot(
1493
+ charge_state_df.index,
1494
+ charge_state_df.values.flatten(),
1495
+ label=self._charge_state,
1496
+ linewidth=2,
1497
+ color='black',
1498
+ )
1499
+ ax.legend()
1196
1500
  fig.tight_layout()
1197
- fig = fig, ax
1501
+
1502
+ figure_like = fig, ax
1503
+ default_filetype = '.png'
1198
1504
 
1199
1505
  return plotting.export_figure(
1200
- fig,
1506
+ figure_like=figure_like,
1201
1507
  default_path=self._calculation_results.folder / title,
1202
- default_filetype='.html',
1508
+ default_filetype=default_filetype,
1203
1509
  user_path=None if isinstance(save, bool) else pathlib.Path(save),
1204
1510
  show=show,
1205
1511
  save=True if save else False,
@@ -1473,37 +1779,95 @@ class SegmentedCalculationResults:
1473
1779
  def plot_heatmap(
1474
1780
  self,
1475
1781
  variable_name: str,
1476
- heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
1477
- heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
1478
- color_map: str = 'portland',
1782
+ reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
1783
+ | Literal['auto']
1784
+ | None = 'auto',
1785
+ colors: str = 'portland',
1479
1786
  save: bool | pathlib.Path = False,
1480
1787
  show: bool = True,
1481
1788
  engine: plotting.PlottingEngine = 'plotly',
1789
+ facet_by: str | list[str] | None = None,
1790
+ animate_by: str | None = None,
1791
+ facet_cols: int = 3,
1792
+ fill: Literal['ffill', 'bfill'] | None = 'ffill',
1793
+ # Deprecated parameters (kept for backwards compatibility)
1794
+ heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] | None = None,
1795
+ heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] | None = None,
1796
+ color_map: str | None = None,
1482
1797
  ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
1483
1798
  """Plot heatmap of variable solution across segments.
1484
1799
 
1485
1800
  Args:
1486
1801
  variable_name: Variable to plot.
1487
- heatmap_timeframes: Time aggregation level.
1488
- heatmap_timesteps_per_frame: Timesteps per frame.
1489
- color_map: Color scheme. Also see plotly.
1802
+ reshape_time: Time reshaping configuration (default: 'auto'):
1803
+ - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
1804
+ - Tuple like ('D', 'h'): Explicit reshaping (days vs hours)
1805
+ - None: Disable time reshaping
1806
+ colors: Color scheme. See plotting.ColorType for options.
1490
1807
  save: Whether to save plot.
1491
1808
  show: Whether to display plot.
1492
1809
  engine: Plotting engine.
1810
+ facet_by: Dimension(s) to create facets (subplots) for.
1811
+ animate_by: Dimension to animate over (Plotly only).
1812
+ facet_cols: Number of columns in the facet grid layout.
1813
+ fill: Method to fill missing values: 'ffill' or 'bfill'.
1814
+ heatmap_timeframes: (Deprecated) Use reshape_time instead.
1815
+ heatmap_timesteps_per_frame: (Deprecated) Use reshape_time instead.
1816
+ color_map: (Deprecated) Use colors instead.
1493
1817
 
1494
1818
  Returns:
1495
1819
  Figure object.
1496
1820
  """
1821
+ # Handle deprecated parameters
1822
+ if heatmap_timeframes is not None or heatmap_timesteps_per_frame is not None:
1823
+ # Check for conflict with new parameter
1824
+ if reshape_time != 'auto': # Check if user explicitly set reshape_time
1825
+ raise ValueError(
1826
+ "Cannot use both deprecated parameters 'heatmap_timeframes'/'heatmap_timesteps_per_frame' "
1827
+ "and new parameter 'reshape_time'. Use only 'reshape_time'."
1828
+ )
1829
+
1830
+ import warnings
1831
+
1832
+ warnings.warn(
1833
+ "The 'heatmap_timeframes' and 'heatmap_timesteps_per_frame' parameters are deprecated. "
1834
+ "Use 'reshape_time=(timeframes, timesteps_per_frame)' instead.",
1835
+ DeprecationWarning,
1836
+ stacklevel=2,
1837
+ )
1838
+ # Override reshape_time if old parameters provided
1839
+ if heatmap_timeframes is not None and heatmap_timesteps_per_frame is not None:
1840
+ reshape_time = (heatmap_timeframes, heatmap_timesteps_per_frame)
1841
+
1842
+ if color_map is not None:
1843
+ # Check for conflict with new parameter
1844
+ if colors != 'portland': # Check if user explicitly set colors
1845
+ raise ValueError(
1846
+ "Cannot use both deprecated parameter 'color_map' and new parameter 'colors'. Use only 'colors'."
1847
+ )
1848
+
1849
+ import warnings
1850
+
1851
+ warnings.warn(
1852
+ "The 'color_map' parameter is deprecated. Use 'colors' instead.",
1853
+ DeprecationWarning,
1854
+ stacklevel=2,
1855
+ )
1856
+ colors = color_map
1857
+
1497
1858
  return plot_heatmap(
1498
- dataarray=self.solution_without_overlap(variable_name),
1859
+ data=self.solution_without_overlap(variable_name),
1499
1860
  name=variable_name,
1500
1861
  folder=self.folder,
1501
- heatmap_timeframes=heatmap_timeframes,
1502
- heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
1503
- color_map=color_map,
1862
+ reshape_time=reshape_time,
1863
+ colors=colors,
1504
1864
  save=save,
1505
1865
  show=show,
1506
1866
  engine=engine,
1867
+ facet_by=facet_by,
1868
+ animate_by=animate_by,
1869
+ facet_cols=facet_cols,
1870
+ fill=fill,
1507
1871
  )
1508
1872
 
1509
1873
  def to_file(self, folder: str | pathlib.Path | None = None, name: str | None = None, compression: int = 5):
@@ -1533,59 +1897,212 @@ class SegmentedCalculationResults:
1533
1897
 
1534
1898
 
1535
1899
  def plot_heatmap(
1536
- dataarray: xr.DataArray,
1537
- name: str,
1538
- folder: pathlib.Path,
1539
- heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
1540
- heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
1541
- color_map: str = 'portland',
1900
+ data: xr.DataArray | xr.Dataset,
1901
+ name: str | None = None,
1902
+ folder: pathlib.Path | None = None,
1903
+ colors: plotting.ColorType = 'viridis',
1542
1904
  save: bool | pathlib.Path = False,
1543
1905
  show: bool = True,
1544
1906
  engine: plotting.PlottingEngine = 'plotly',
1907
+ select: dict[str, Any] | None = None,
1908
+ facet_by: str | list[str] | None = None,
1909
+ animate_by: str | None = None,
1910
+ facet_cols: int = 3,
1911
+ reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
1912
+ | Literal['auto']
1913
+ | None = 'auto',
1914
+ fill: Literal['ffill', 'bfill'] | None = 'ffill',
1915
+ # Deprecated parameters (kept for backwards compatibility)
1545
1916
  indexer: dict[str, Any] | None = None,
1917
+ heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] | None = None,
1918
+ heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] | None = None,
1919
+ color_map: str | None = None,
1546
1920
  ):
1547
- """Plot heatmap of time series data.
1921
+ """Plot heatmap visualization with support for multi-variable, faceting, and animation.
1922
+
1923
+ This function provides a standalone interface to the heatmap plotting capabilities,
1924
+ supporting the same modern features as CalculationResults.plot_heatmap().
1548
1925
 
1549
1926
  Args:
1550
- dataarray: Data to plot.
1551
- name: Variable name for title.
1552
- folder: Save folder.
1553
- heatmap_timeframes: Time aggregation level.
1554
- heatmap_timesteps_per_frame: Timesteps per frame.
1555
- color_map: Color scheme. Also see plotly.
1556
- save: Whether to save plot.
1557
- show: Whether to display plot.
1558
- engine: Plotting engine.
1559
- indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
1560
- If None, uses first value for each dimension.
1561
- If empty dict {}, uses all values.
1927
+ data: Data to plot. Can be a single DataArray or an xarray Dataset.
1928
+ When a Dataset is provided, all data variables are combined along a new 'variable' dimension.
1929
+ name: Optional name for the title. If not provided, uses the DataArray name or
1930
+ generates a default title for Datasets.
1931
+ folder: Save folder for the plot. Defaults to current directory if not provided.
1932
+ colors: Color scheme for the heatmap. See `flixopt.plotting.ColorType` for options.
1933
+ save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
1934
+ show: Whether to show the plot or not.
1935
+ engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
1936
+ select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
1937
+ facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
1938
+ or list of dimensions. Each unique value combination creates a subplot.
1939
+ animate_by: Dimension to animate over (Plotly only). Creates animation frames.
1940
+ facet_cols: Number of columns in the facet grid layout (default: 3).
1941
+ reshape_time: Time reshaping configuration (default: 'auto'):
1942
+ - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
1943
+ - Tuple: Explicit reshaping, e.g. ('D', 'h') for days vs hours
1944
+ - None: Disable auto-reshaping
1945
+ fill: Method to fill missing values after reshape: 'ffill' (forward fill) or 'bfill' (backward fill).
1946
+ Default is 'ffill'.
1947
+
1948
+ Examples:
1949
+ Single DataArray with time reshaping:
1950
+
1951
+ >>> plot_heatmap(data, name='Temperature', folder=Path('.'), reshape_time=('D', 'h'))
1952
+
1953
+ Dataset with multiple variables (facet by variable):
1954
+
1955
+ >>> dataset = xr.Dataset({'Boiler': data1, 'CHP': data2, 'Storage': data3})
1956
+ >>> plot_heatmap(
1957
+ ... dataset,
1958
+ ... folder=Path('.'),
1959
+ ... facet_by='variable',
1960
+ ... reshape_time=('D', 'h'),
1961
+ ... )
1962
+
1963
+ Dataset with animation by variable:
1964
+
1965
+ >>> plot_heatmap(dataset, animate_by='variable', reshape_time=('D', 'h'))
1562
1966
  """
1563
- dataarray, suffix_parts = _apply_indexer_to_data(dataarray, indexer, drop=True)
1967
+ # Handle deprecated heatmap time parameters
1968
+ if heatmap_timeframes is not None or heatmap_timesteps_per_frame is not None:
1969
+ # Check for conflict with new parameter
1970
+ if reshape_time != 'auto': # User explicitly set reshape_time
1971
+ raise ValueError(
1972
+ "Cannot use both deprecated parameters 'heatmap_timeframes'/'heatmap_timesteps_per_frame' "
1973
+ "and new parameter 'reshape_time'. Use only 'reshape_time'."
1974
+ )
1975
+
1976
+ import warnings
1977
+
1978
+ warnings.warn(
1979
+ "The 'heatmap_timeframes' and 'heatmap_timesteps_per_frame' parameters are deprecated. "
1980
+ "Use 'reshape_time=(timeframes, timesteps_per_frame)' instead.",
1981
+ DeprecationWarning,
1982
+ stacklevel=2,
1983
+ )
1984
+ # Override reshape_time if both old parameters provided
1985
+ if heatmap_timeframes is not None and heatmap_timesteps_per_frame is not None:
1986
+ reshape_time = (heatmap_timeframes, heatmap_timesteps_per_frame)
1987
+
1988
+ # Handle deprecated color_map parameter
1989
+ if color_map is not None:
1990
+ # Check for conflict with new parameter
1991
+ if colors != 'viridis': # User explicitly set colors
1992
+ raise ValueError(
1993
+ "Cannot use both deprecated parameter 'color_map' and new parameter 'colors'. Use only 'colors'."
1994
+ )
1995
+
1996
+ import warnings
1997
+
1998
+ warnings.warn(
1999
+ "The 'color_map' parameter is deprecated. Use 'colors' instead.",
2000
+ DeprecationWarning,
2001
+ stacklevel=2,
2002
+ )
2003
+ colors = color_map
2004
+
2005
+ # Handle deprecated indexer parameter
2006
+ if indexer is not None:
2007
+ # Check for conflict with new parameter
2008
+ if select is not None: # User explicitly set select
2009
+ raise ValueError(
2010
+ "Cannot use both deprecated parameter 'indexer' and new parameter 'select'. Use only 'select'."
2011
+ )
2012
+
2013
+ import warnings
2014
+
2015
+ warnings.warn(
2016
+ "The 'indexer' parameter is deprecated. Use 'select' instead.",
2017
+ DeprecationWarning,
2018
+ stacklevel=2,
2019
+ )
2020
+ select = indexer
2021
+
2022
+ # Convert Dataset to DataArray with 'variable' dimension
2023
+ if isinstance(data, xr.Dataset):
2024
+ # Extract all data variables from the Dataset
2025
+ variable_names = list(data.data_vars)
2026
+ dataarrays = [data[var] for var in variable_names]
2027
+
2028
+ # Combine into single DataArray with 'variable' dimension
2029
+ data = xr.concat(dataarrays, dim='variable')
2030
+ data = data.assign_coords(variable=variable_names)
2031
+
2032
+ # Use Dataset variable names for title if name not provided
2033
+ if name is None:
2034
+ title_name = f'Heatmap of {len(variable_names)} variables'
2035
+ else:
2036
+ title_name = name
2037
+ else:
2038
+ # Single DataArray
2039
+ if name is None:
2040
+ title_name = data.name if data.name else 'Heatmap'
2041
+ else:
2042
+ title_name = name
2043
+
2044
+ # Apply select filtering
2045
+ data, suffix_parts = _apply_selection_to_data(data, select=select, drop=True)
1564
2046
  suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
1565
- name = name if not suffix_parts else name + suffix
1566
2047
 
1567
- heatmap_data = plotting.heat_map_data_from_df(
1568
- dataarray.to_dataframe(name), heatmap_timeframes, heatmap_timesteps_per_frame, 'ffill'
1569
- )
2048
+ # Matplotlib heatmaps require at most 2D data
2049
+ # Time dimension will be reshaped to 2D (timeframe × timestep), so can't have other dims alongside it
2050
+ if engine == 'matplotlib':
2051
+ dims = list(data.dims)
1570
2052
 
1571
- xlabel, ylabel = f'timeframe [{heatmap_timeframes}]', f'timesteps [{heatmap_timesteps_per_frame}]'
2053
+ # If 'time' dimension exists and will be reshaped, we can't have any other dimensions
2054
+ if 'time' in dims and len(dims) > 1 and reshape_time is not None:
2055
+ extra_dims = [d for d in dims if d != 'time']
2056
+ raise ValueError(
2057
+ f'Matplotlib heatmaps with time reshaping cannot have additional dimensions. '
2058
+ f'Found extra dimensions: {extra_dims}. '
2059
+ f'Use select={{...}} to reduce to time only, use "reshape_time=None" or switch to engine="plotly" or use for multi-dimensional support.'
2060
+ )
2061
+ # If no 'time' dimension (already reshaped or different data), allow at most 2 dimensions
2062
+ elif 'time' not in dims and len(dims) > 2:
2063
+ raise ValueError(
2064
+ f'Matplotlib heatmaps support at most 2 dimensions, but data has {len(dims)}: {dims}. '
2065
+ f'Use select={{...}} to reduce dimensions or switch to engine="plotly".'
2066
+ )
1572
2067
 
2068
+ # Build title
2069
+ title = f'{title_name}{suffix}'
2070
+ if isinstance(reshape_time, tuple):
2071
+ timeframes, timesteps_per_frame = reshape_time
2072
+ title += f' ({timeframes} vs {timesteps_per_frame})'
2073
+
2074
+ # Plot with appropriate engine
1573
2075
  if engine == 'plotly':
1574
- figure_like = plotting.heat_map_plotly(
1575
- heatmap_data, title=name, color_map=color_map, xlabel=xlabel, ylabel=ylabel
2076
+ figure_like = plotting.heatmap_with_plotly(
2077
+ data=data,
2078
+ facet_by=facet_by,
2079
+ animate_by=animate_by,
2080
+ colors=colors,
2081
+ title=title,
2082
+ facet_cols=facet_cols,
2083
+ reshape_time=reshape_time,
2084
+ fill=fill,
1576
2085
  )
1577
2086
  default_filetype = '.html'
1578
2087
  elif engine == 'matplotlib':
1579
- figure_like = plotting.heat_map_matplotlib(
1580
- heatmap_data, title=name, color_map=color_map, xlabel=xlabel, ylabel=ylabel
2088
+ figure_like = plotting.heatmap_with_matplotlib(
2089
+ data=data,
2090
+ colors=colors,
2091
+ title=title,
2092
+ reshape_time=reshape_time,
2093
+ fill=fill,
1581
2094
  )
1582
2095
  default_filetype = '.png'
1583
2096
  else:
1584
2097
  raise ValueError(f'Engine "{engine}" not supported. Use "plotly" or "matplotlib"')
1585
2098
 
2099
+ # Set default folder if not provided
2100
+ if folder is None:
2101
+ folder = pathlib.Path('.')
2102
+
1586
2103
  return plotting.export_figure(
1587
2104
  figure_like=figure_like,
1588
- default_path=folder / f'{name} ({heatmap_timeframes}-{heatmap_timesteps_per_frame})',
2105
+ default_path=folder / title,
1589
2106
  default_filetype=default_filetype,
1590
2107
  user_path=None if isinstance(save, bool) else pathlib.Path(save),
1591
2108
  show=show,
@@ -1787,8 +2304,13 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
1787
2304
  if coord_name not in array.coords:
1788
2305
  raise AttributeError(f"Missing required coordinate '{coord_name}'")
1789
2306
 
1790
- # Convert single value to list
1791
- val_list = [coord_values] if isinstance(coord_values, str) else coord_values
2307
+ # Normalize to list for sequence-like inputs (excluding strings)
2308
+ if isinstance(coord_values, str):
2309
+ val_list = [coord_values]
2310
+ elif isinstance(coord_values, (list, tuple, np.ndarray, pd.Index)):
2311
+ val_list = list(coord_values)
2312
+ else:
2313
+ val_list = [coord_values]
1792
2314
 
1793
2315
  # Verify coord_values exist
1794
2316
  available = set(array[coord_name].values)
@@ -1798,7 +2320,7 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
1798
2320
 
1799
2321
  # Apply filter
1800
2322
  return array.where(
1801
- array[coord_name].isin(val_list) if isinstance(coord_values, list) else array[coord_name] == coord_values,
2323
+ array[coord_name].isin(val_list) if len(val_list) > 1 else array[coord_name] == val_list[0],
1802
2324
  drop=True,
1803
2325
  )
1804
2326
 
@@ -1817,36 +2339,26 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
1817
2339
  return da
1818
2340
 
1819
2341
 
1820
- def _apply_indexer_to_data(
1821
- data: xr.DataArray | xr.Dataset, indexer: dict[str, Any] | None = None, drop=False
2342
+ def _apply_selection_to_data(
2343
+ data: xr.DataArray | xr.Dataset,
2344
+ select: dict[str, Any] | None = None,
2345
+ drop=False,
1822
2346
  ) -> tuple[xr.DataArray | xr.Dataset, list[str]]:
1823
2347
  """
1824
- Apply indexer selection or auto-select first values for non-time dimensions.
2348
+ Apply selection to data.
1825
2349
 
1826
2350
  Args:
1827
2351
  data: xarray Dataset or DataArray
1828
- indexer: Optional selection dict
1829
- If None, uses first value for each dimension (except time).
1830
- If empty dict {}, uses all values.
2352
+ select: Optional selection dict
2353
+ drop: Whether to drop dimensions after selection
1831
2354
 
1832
2355
  Returns:
1833
2356
  Tuple of (selected_data, selection_string)
1834
2357
  """
1835
2358
  selection_string = []
1836
2359
 
1837
- if indexer is not None:
1838
- # User provided indexer
1839
- data = data.sel(indexer, drop=drop)
1840
- selection_string.extend(f'{v}[{k}]' for k, v in indexer.items())
1841
- else:
1842
- # Auto-select first value for each dimension except 'time'
1843
- selection = {}
1844
- for dim in data.dims:
1845
- if dim != 'time' and dim in data.coords:
1846
- first_value = data.coords[dim].values[0]
1847
- selection[dim] = first_value
1848
- selection_string.append(f'{first_value}[{dim}]')
1849
- if selection:
1850
- data = data.sel(selection, drop=drop)
2360
+ if select:
2361
+ data = data.sel(select, drop=drop)
2362
+ selection_string.extend(f'{dim}={val}' for dim, val in select.items())
1851
2363
 
1852
2364
  return data, selection_string