flixopt 3.0.1__py3-none-any.whl → 6.0.0rc7__py3-none-any.whl

flixopt/results.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import datetime
 import json
 import logging
@@ -10,33 +11,63 @@ from typing import TYPE_CHECKING, Any, Literal
 import linopy
 import numpy as np
 import pandas as pd
-import plotly
 import xarray as xr
-import yaml
 
 from . import io as fx_io
 from . import plotting
+from .color_processing import process_colors
+from .config import CONFIG, DEPRECATION_REMOVAL_VERSION, SUCCESS_LEVEL
 from .flow_system import FlowSystem
+from .structure import CompositeContainerMixin, ResultsContainer
 
 if TYPE_CHECKING:
     import matplotlib.pyplot as plt
+    import plotly
     import pyvis
 
-    from .calculation import Calculation, SegmentedCalculation
     from .core import FlowSystemDimensions
-
+    from .optimization import Optimization, SegmentedOptimization
 
 logger = logging.getLogger('flixopt')
 
 
+def load_mapping_from_file(path: pathlib.Path) -> dict[str, str | list[str]]:
+    """Load color mapping from JSON or YAML file.
+
+    Tries loader based on file suffix first, with fallback to the other format.
+
+    Args:
+        path: Path to config file (.json or .yaml/.yml)
+
+    Returns:
+        Dictionary mapping components to colors or colorscales to component lists
+
+    Raises:
+        ValueError: If file cannot be loaded as JSON or YAML
+    """
+    return fx_io.load_config_file(path)
+
+
+def _get_solution_attr(solution: xr.Dataset, key: str) -> dict:
+    """Get an attribute from solution, decoding JSON if necessary.
+
+    Solution attrs are stored as JSON strings for netCDF compatibility.
+    This helper handles both JSON strings and dicts (for backward compatibility).
+    """
+    value = solution.attrs.get(key, {})
+    if isinstance(value, str):
+        return json.loads(value)
+    return value
+
+
 class _FlowSystemRestorationError(Exception):
     """Exception raised when a FlowSystem cannot be restored from dataset."""
 
     pass
 
 
-class CalculationResults:
-    """Comprehensive container for optimization calculation results and analysis tools.
+class Results(CompositeContainerMixin['ComponentResults | BusResults | EffectResults | FlowResults']):
+    """Comprehensive container for optimization results and analysis tools.
 
     This class provides unified access to all optimization results including flow rates,
     component states, bus balances, and system effects. It offers powerful analysis
@@ -55,27 +86,27 @@ class CalculationResults:
         - **Buses**: Network node balances and energy flows
         - **Effects**: System-wide impacts (costs, emissions, resource consumption)
         - **Solution**: Raw optimization variables and their values
-        - **Metadata**: Calculation parameters, timing, and system configuration
+        - **Metadata**: Optimization parameters, timing, and system configuration
 
     Attributes:
         solution: Dataset containing all optimization variable solutions
         flow_system_data: Dataset with complete system configuration and parameters. Restore the used FlowSystem for further analysis.
-        summary: Calculation metadata including solver status, timing, and statistics
-        name: Unique identifier for this calculation
+        summary: Optimization metadata including solver status, timing, and statistics
+        name: Unique identifier for this optimization
         model: Original linopy optimization model (if available)
         folder: Directory path for result storage and loading
         components: Dictionary mapping component labels to ComponentResults objects
         buses: Dictionary mapping bus labels to BusResults objects
         effects: Dictionary mapping effect names to EffectResults objects
         timesteps_extra: Extended time index including boundary conditions
-        hours_per_timestep: Duration of each timestep for proper energy calculations
+        timestep_duration: Duration of each timestep in hours for proper energy calculations
 
     Examples:
         Load and analyze saved results:
 
         ```python
         # Load results from file
-        results = CalculationResults.from_file('results', 'annual_optimization')
+        results = Results.from_file('results', 'annual_optimization')
 
         # Access specific component results
         boiler_results = results['Boiler_01']
@@ -107,27 +138,43 @@ class CalculationResults:
         ).mean()
         ```
 
+        Configure automatic color management for plots:
+
+        ```python
+        # Dict-based configuration:
+        results.setup_colors({'Solar*': 'Oranges', 'Wind*': 'Blues', 'Battery': 'green'})
+
+        # All plots automatically use configured colors (colors=None is the default)
+        results['ElectricityBus'].plot_node_balance()
+        results['Battery'].plot_charge_state()
+
+        # Override when needed
+        results['ElectricityBus'].plot_node_balance(colors='turbo')  # Ignores setup
+        ```
+
     Design Patterns:
-        **Factory Methods**: Use `from_file()` and `from_calculation()` for creation or access directly from `Calculation.results`
+        **Factory Methods**: Use `from_file()` and `from_optimization()` for creation or access directly from `Optimization.results`
         **Dictionary Access**: Use `results[element_label]` for element-specific results
         **Lazy Loading**: Results objects created on-demand for memory efficiency
         **Unified Interface**: Consistent API across different result types
 
     """
 
+    model: linopy.Model | None
+
     @classmethod
-    def from_file(cls, folder: str | pathlib.Path, name: str) -> CalculationResults:
-        """Load CalculationResults from saved files.
+    def from_file(cls, folder: str | pathlib.Path, name: str) -> Results:
+        """Load Results from saved files.
 
         Args:
             folder: Directory containing saved files.
             name: Base name of saved files (without extensions).
 
         Returns:
-            CalculationResults: Loaded instance.
+            Results: Loaded instance.
         """
         folder = pathlib.Path(folder)
-        paths = fx_io.CalculationResultsPaths(folder, name)
+        paths = fx_io.ResultsPaths(folder, name)
 
         model = None
         if paths.linopy_model.exists():
@@ -137,8 +184,7 @@ class CalculationResults:
             except Exception as e:
                 logger.critical(f'Could not load the linopy model "{name}" from file ("{paths.linopy_model}"): {e}')
 
-        with open(paths.summary, encoding='utf-8') as f:
-            summary = yaml.load(f, Loader=yaml.FullLoader)
+        summary = fx_io.load_yaml(paths.summary)
 
         return cls(
             solution=fx_io.load_dataset_from_netcdf(paths.solution),
@@ -150,22 +196,22 @@ class CalculationResults:
         )
 
     @classmethod
-    def from_calculation(cls, calculation: Calculation) -> CalculationResults:
-        """Create CalculationResults from a Calculation object.
+    def from_optimization(cls, optimization: Optimization) -> Results:
+        """Create Results from an Optimization instance.
 
         Args:
-            calculation: Calculation object with solved model.
+            optimization: The Optimization instance to extract results from.
 
         Returns:
-            CalculationResults: New instance with extracted results.
+            Results: New instance containing the optimization results.
         """
         return cls(
-            solution=calculation.model.solution,
-            flow_system_data=calculation.flow_system.to_dataset(),
-            summary=calculation.summary,
-            model=calculation.model,
-            name=calculation.name,
-            folder=calculation.folder,
+            solution=optimization.model.solution,
+            flow_system_data=optimization.flow_system.to_dataset(),
+            summary=optimization.summary,
+            model=optimization.model,
+            name=optimization.name,
+            folder=optimization.folder,
         )
 
     def __init__(
@@ -176,30 +222,27 @@ class CalculationResults:
         summary: dict,
         folder: pathlib.Path | None = None,
         model: linopy.Model | None = None,
-        **kwargs,  # To accept old "flow_system" parameter
     ):
-        """Initialize CalculationResults with optimization data.
-        Usually, this class is instantiated by the Calculation class, or by loading from file.
+        """Initialize Results with optimization data.
+        Usually, this class is instantiated by an Optimization object via `Results.from_optimization()`
+        or by loading from file using `Results.from_file()`.
 
         Args:
             solution: Optimization solution dataset.
             flow_system_data: Flow system configuration dataset.
-            name: Calculation name.
-            summary: Calculation metadata.
+            name: Optimization name.
+            summary: Optimization metadata.
             folder: Results storage folder.
             model: Linopy optimization model.
-        Deprecated:
-            flow_system: Use flow_system_data instead.
         """
-        # Handle potential old "flow_system" parameter for backward compatibility
-        if 'flow_system' in kwargs and flow_system_data is None:
-            flow_system_data = kwargs.pop('flow_system')
-            warnings.warn(
-                "The 'flow_system' parameter is deprecated. Use 'flow_system_data' instead."
-                "Acess is now by '.flow_system_data', while '.flow_system' returns the restored FlowSystem.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
+        warnings.warn(
+            f'Results is deprecated and will be removed in v{DEPRECATION_REMOVAL_VERSION}. '
+            'Access results directly via FlowSystem.solution after optimization, or use the '
+            '.plot accessor on FlowSystem and its components (e.g., flow_system.plot.heatmap(...)). '
+            'To load old result files, use FlowSystem.from_old_results(folder, name).',
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
         self.solution = solution
         self.flow_system_data = flow_system_data
@@ -207,29 +250,44 @@ class CalculationResults:
         self.name = name
         self.model = model
         self.folder = pathlib.Path(folder) if folder is not None else pathlib.Path.cwd() / 'results'
-        self.components = {
-            label: ComponentResults(self, **infos) for label, infos in self.solution.attrs['Components'].items()
+
+        # Create ResultsContainers for better access patterns
+        components_dict = {
+            label: ComponentResults(self, **infos)
+            for label, infos in _get_solution_attr(self.solution, 'Components').items()
         }
+        self.components = ResultsContainer(
+            elements=components_dict, element_type_name='component results', truncate_repr=10
+        )
 
-        self.buses = {label: BusResults(self, **infos) for label, infos in self.solution.attrs['Buses'].items()}
+        buses_dict = {
+            label: BusResults(self, **infos) for label, infos in _get_solution_attr(self.solution, 'Buses').items()
+        }
+        self.buses = ResultsContainer(elements=buses_dict, element_type_name='bus results', truncate_repr=10)
 
-        self.effects = {label: EffectResults(self, **infos) for label, infos in self.solution.attrs['Effects'].items()}
+        effects_dict = {
+            label: EffectResults(self, **infos) for label, infos in _get_solution_attr(self.solution, 'Effects').items()
+        }
+        self.effects = ResultsContainer(elements=effects_dict, element_type_name='effect results', truncate_repr=10)
 
-        if 'Flows' not in self.solution.attrs:
+        flows_attr = _get_solution_attr(self.solution, 'Flows')
+        if not flows_attr:
             warnings.warn(
                 'No Data about flows found in the results. This data is only included since v2.2.0. Some functionality '
                 'is not availlable. We recommend to evaluate your results with a version <2.2.0.',
                 stacklevel=2,
             )
-            self.flows = {}
+            flows_dict = {}
+            self._has_flow_data = False
         else:
-            self.flows = {
-                label: FlowResults(self, **infos) for label, infos in self.solution.attrs.get('Flows', {}).items()
-            }
+            flows_dict = {label: FlowResults(self, **infos) for label, infos in flows_attr.items()}
+            self._has_flow_data = True
+        self.flows = ResultsContainer(elements=flows_dict, element_type_name='flow results', truncate_repr=10)
 
         self.timesteps_extra = self.solution.indexes['time']
-        self.hours_per_timestep = FlowSystem.calculate_hours_per_timestep(self.timesteps_extra)
+        self.timestep_duration = FlowSystem.calculate_timestep_duration(self.timesteps_extra)
         self.scenarios = self.solution.indexes['scenario'] if 'scenario' in self.solution.indexes else None
+        self.periods = self.solution.indexes['period'] if 'period' in self.solution.indexes else None
 
         self._effect_share_factors = None
         self._flow_system = None
@@ -239,16 +297,24 @@ class CalculationResults:
         self._sizes = None
         self._effects_per_component = None
 
-    def __getitem__(self, key: str) -> ComponentResults | BusResults | EffectResults:
-        if key in self.components:
-            return self.components[key]
-        if key in self.buses:
-            return self.buses[key]
-        if key in self.effects:
-            return self.effects[key]
-        if key in self.flows:
-            return self.flows[key]
-        raise KeyError(f'No element with label {key} found.')
+        self.colors: dict[str, str] = {}
+
+    def _get_container_groups(self) -> dict[str, ResultsContainer]:
+        """Return ordered container groups for CompositeContainerMixin."""
+        return {
+            'Components': self.components,
+            'Buses': self.buses,
+            'Effects': self.effects,
+            'Flows': self.flows,
+        }
+
+    def __repr__(self) -> str:
+        """Return grouped representation of all results."""
+        r = fx_io.format_title_with_underline(self.__class__.__name__, '=')
+        r += f'Name: "{self.name}"\nFolder: {self.folder}\n'
+        # Add grouped container view
+        r += '\n' + self._format_grouped_containers()
+        return r
 
     @property
     def storages(self) -> list[ComponentResults]:
@@ -288,23 +354,151 @@ class CalculationResults:
 
     @property
     def flow_system(self) -> FlowSystem:
-        """The restored flow_system that was used to create the calculation.
+        """The restored flow_system that was used to create the optimization.
         Contains all input parameters."""
         if self._flow_system is None:
-            old_level = logger.level
-            logger.level = logging.CRITICAL
+            # Temporarily disable all logging to suppress messages during restoration
+            flixopt_logger = logging.getLogger('flixopt')
+            original_level = flixopt_logger.level
+            flixopt_logger.setLevel(logging.CRITICAL + 1)  # Disable all logging
             try:
                 self._flow_system = FlowSystem.from_dataset(self.flow_system_data)
                 self._flow_system._connect_network()
             except Exception as e:
+                flixopt_logger.setLevel(original_level)  # Re-enable before logging
                 logger.critical(
                     f'Not able to restore FlowSystem from dataset. Some functionality is not availlable. {e}'
                 )
                 raise _FlowSystemRestorationError(f'Not able to restore FlowSystem from dataset. {e}') from e
             finally:
-                logger.level = old_level
+                flixopt_logger.setLevel(original_level)  # Restore original level
         return self._flow_system
 
+    def setup_colors(
+        self,
+        config: dict[str, str | list[str]] | str | pathlib.Path | None = None,
+        default_colorscale: str | None = None,
+    ) -> dict[str, str]:
+        """
+        Setup colors for all variables across all elements. Overwrites existing ones.
+
+        Args:
+            config: Configuration for color assignment. Can be:
+                - dict: Maps components to colors/colorscales:
+                    * 'component1': 'red'  # Single component to single color
+                    * 'component1': '#FF0000'  # Single component to hex color
+                    - OR maps colorscales to multiple components:
+                    * 'colorscale_name': ['component1', 'component2']  # Colorscale across components
+                - str: Path to a JSON/YAML config file or a colorscale name to apply to all
+                - Path: Path to a JSON/YAML config file
+                - None: Use default_colorscale for all components
+            default_colorscale: Default colorscale for unconfigured components (default: 'turbo')
+
+        Examples:
+            setup_colors({
+                # Direct component-to-color mappings
+                'Boiler1': '#FF0000',
+                'CHP': 'darkred',
+                # Colorscale for multiple components
+                'Oranges': ['Solar1', 'Solar2'],
+                'Blues': ['Wind1', 'Wind2'],
+                'Greens': ['Battery1', 'Battery2', 'Battery3'],
+            })
+
+        Returns:
+            Complete variable-to-color mapping dictionary
+        """
+
+        def get_all_variable_names(comp: str) -> list[str]:
+            """Collect all variables from the component, including flows and flow_hours."""
+            comp_object = self.components[comp]
+            var_names = [comp] + list(comp_object.variable_names)
+            for flow in comp_object.flows:
+                var_names.extend([flow, f'{flow}|flow_hours'])
+            return var_names
+
+        # Set default colorscale if not provided
+        if default_colorscale is None:
+            default_colorscale = CONFIG.Plotting.default_qualitative_colorscale
+
+        # Handle different config input types
+        if config is None:
+            # Apply default colorscale to all components
+            config_dict = {}
+        elif isinstance(config, (str, pathlib.Path)):
+            # Try to load from file first
+            config_path = pathlib.Path(config)
+            if config_path.exists():
+                # Load config from file using helper
+                config_dict = load_mapping_from_file(config_path)
+            else:
+                # Treat as colorscale name to apply to all components
+                all_components = list(self.components.keys())
+                config_dict = {config: all_components}
+        elif isinstance(config, dict):
+            config_dict = config
+        else:
+            raise TypeError(f'config must be dict, str, Path, or None, got {type(config)}')
+
+        # Step 1: Build component-to-color mapping
+        component_colors: dict[str, str] = {}
+
+        # Track which components are configured
+        configured_components = set()
+
+        # Process each configuration entry
+        for key, value in config_dict.items():
+            # Check if value is a list (colorscale -> [components])
+            # or a string (component -> color OR colorscale -> [components])
+
+            if isinstance(value, list):
+                # key is colorscale, value is list of components
+                # Format: 'Blues': ['Wind1', 'Wind2']
+                components = value
+                colorscale_name = key
+
+                # Validate components exist
+                for component in components:
+                    if component not in self.components:
+                        raise ValueError(f"Component '{component}' not found")
+
+                configured_components.update(components)
+
+                # Use process_colors to get one color per component from the colorscale
+                colors_for_components = process_colors(colorscale_name, components)
+                component_colors.update(colors_for_components)
+
+            elif isinstance(value, str):
+                # Check if key is an existing component
+                if key in self.components:
+                    # Format: 'CHP': 'red' (component -> color)
+                    component, color = key, value
+
+                    configured_components.add(component)
+                    component_colors[component] = color
+                else:
+                    raise ValueError(f"Component '{key}' not found")
+            else:
+                raise TypeError(f'Config value must be str or list, got {type(value)}')
+
+        # Step 2: Assign colors to remaining unconfigured components
+        remaining_components = list(set(self.components.keys()) - configured_components)
+        if remaining_components:
+            # Use default colorscale to assign one color per remaining component
+            default_colors = process_colors(default_colorscale, remaining_components)
+            component_colors.update(default_colors)
+
+        # Step 3: Build variable-to-color mapping
+        # Clear existing colors to avoid stale keys
+        self.colors = {}
+        # Each component's variables all get the same color as the component
+        for component, color in component_colors.items():
+            variable_names = get_all_variable_names(component)
+            for var_name in variable_names:
+                self.colors[var_name] = color
+
+        return self.colors
+
     def filter_solution(
         self,
         variable_dims: Literal['scalar', 'time', 'scenario', 'timeonly', 'scenarioonly'] | None = None,
@@ -373,21 +567,42 @@ class CalculationResults:
     ) -> xr.DataArray:
         """Returns a DataArray containing the flow rates of each Flow.
 
-        Args:
-            start: Optional source node(s) to filter by. Can be a single node name or a list of names.
-            end: Optional destination node(s) to filter by. Can be a single node name or a list of names.
-            component: Optional component(s) to filter by. Can be a single component name or a list of names.
+        .. deprecated::
+            Use `results.plot.all_flow_rates` (Dataset) or
+            `results.flows['FlowLabel'].flow_rate` (DataArray) instead.
 
-        Further usage:
-            Convert the dataarray to a dataframe:
-            >>>results.flow_rates().to_pandas()
-            Get the max or min over time:
-            >>>results.flow_rates().max('time')
-            Sum up the flow rates of flows with the same start and end:
-            >>>results.flow_rates(end='Fernwärme').groupby('start').sum(dim='flow')
-            To recombine filtered dataarrays, use `xr.concat` with dim 'flow':
-            >>>xr.concat([results.flow_rates(start='Fernwärme'), results.flow_rates(end='Fernwärme')], dim='flow')
+            **Note**: The new API differs from this method:
+
+            - Returns ``xr.Dataset`` (not ``DataArray``) with flow labels as variable names
+            - No ``'flow'`` dimension - each flow is a separate variable
+            - No filtering parameters - filter using these alternatives::
+
+                # Select specific flows by label
+                ds = results.plot.all_flow_rates
+                ds[['Boiler(Q_th)', 'CHP(Q_th)']]
+
+                # Filter by substring in label
+                ds[[v for v in ds.data_vars if 'Boiler' in v]]
+
+                # Filter by bus (start/end) - get flows connected to a bus
+                results['Fernwärme'].inputs  # list of input flow labels
+                results['Fernwärme'].outputs  # list of output flow labels
+                ds[results['Fernwärme'].inputs]  # Dataset with only inputs to bus
+
+                # Filter by component - get flows of a component
+                results['Boiler'].inputs  # list of input flow labels
+                results['Boiler'].outputs  # list of output flow labels
         """
+        warnings.warn(
+            'results.flow_rates() is deprecated. '
+            'Use results.plot.all_flow_rates instead (returns Dataset, not DataArray). '
+            'Note: The new API has no filtering parameters and uses flow labels as variable names. '
+            f'Will be removed in v{DEPRECATION_REMOVAL_VERSION}.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if not self._has_flow_data:
+            raise ValueError('Flow data is not available in this results object (pre-v2.2.0).')
         if self._flow_rates is None:
             self._flow_rates = self._assign_flow_coords(
                 xr.concat(
@@ -406,6 +621,32 @@ class CalculationResults:
     ) -> xr.DataArray:
         """Returns a DataArray containing the flow hours of each Flow.
 
+        .. deprecated::
+            Use `results.plot.all_flow_hours` (Dataset) or
+            `results.flows['FlowLabel'].flow_rate * results.timestep_duration` instead.
+
+            **Note**: The new API differs from this method:
+
+            - Returns ``xr.Dataset`` (not ``DataArray``) with flow labels as variable names
+            - No ``'flow'`` dimension - each flow is a separate variable
+            - No filtering parameters - filter using these alternatives::
+
+                # Select specific flows by label
+                ds = results.plot.all_flow_hours
+                ds[['Boiler(Q_th)', 'CHP(Q_th)']]
+
+                # Filter by substring in label
+                ds[[v for v in ds.data_vars if 'Boiler' in v]]
+
+                # Filter by bus (start/end) - get flows connected to a bus
+                results['Fernwärme'].inputs  # list of input flow labels
+                results['Fernwärme'].outputs  # list of output flow labels
+                ds[results['Fernwärme'].inputs]  # Dataset with only inputs to bus
+
+                # Filter by component - get flows of a component
+                results['Boiler'].inputs  # list of input flow labels
+                results['Boiler'].outputs  # list of output flow labels
+
         Flow hours represent the total energy/material transferred over time,
         calculated by multiplying flow rates by the duration of each timestep.
 
@@ -425,8 +666,16 @@ class CalculationResults:
             >>>xr.concat([results.flow_hours(start='Fernwärme'), results.flow_hours(end='Fernwärme')], dim='flow')
 
         """
+        warnings.warn(
+            'results.flow_hours() is deprecated. '
+            'Use results.plot.all_flow_hours instead (returns Dataset, not DataArray). '
+            'Note: The new API has no filtering parameters and uses flow labels as variable names. '
+            f'Will be removed in v{DEPRECATION_REMOVAL_VERSION}.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
         if self._flow_hours is None:
-            self._flow_hours = (self.flow_rates() * self.hours_per_timestep).rename('flow_hours')
+            self._flow_hours = (self.flow_rates() * self.timestep_duration).rename('flow_hours')
         filters = {k: v for k, v in {'start': start, 'end': end, 'component': component}.items() if v is not None}
         return filter_dataarray_by_coord(self._flow_hours, **filters)
 
@@ -437,18 +686,43 @@ class CalculationResults:
         component: str | list[str] | None = None,
     ) -> xr.DataArray:
         """Returns a dataset with the sizes of the Flows.
-        Args:
-            start: Optional source node(s) to filter by. Can be a single node name or a list of names.
-            end: Optional destination node(s) to filter by. Can be a single node name or a list of names.
-            component: Optional component(s) to filter by. Can be a single component name or a list of names.
 
-        Further usage:
-            Convert the dataarray to a dataframe:
-            >>>results.sizes().to_pandas()
-            To recombine filtered dataarrays, use `xr.concat` with dim 'flow':
-            >>>xr.concat([results.sizes(start='Fernwärme'), results.sizes(end='Fernwärme')], dim='flow')
+        .. deprecated::
+            Use `results.plot.all_sizes` (Dataset) or
+            `results.flows['FlowLabel'].size` (DataArray) instead.
+
+            **Note**: The new API differs from this method:
+
+            - Returns ``xr.Dataset`` (not ``DataArray``) with flow labels as variable names
+            - No ``'flow'`` dimension - each flow is a separate variable
+            - No filtering parameters - filter using these alternatives::
 
+                # Select specific flows by label
+                ds = results.plot.all_sizes
+                ds[['Boiler(Q_th)', 'CHP(Q_th)']]
+
+                # Filter by substring in label
+                ds[[v for v in ds.data_vars if 'Boiler' in v]]
+
+                # Filter by bus (start/end) - get flows connected to a bus
+                results['Fernwärme'].inputs  # list of input flow labels
+                results['Fernwärme'].outputs  # list of output flow labels
+                ds[results['Fernwärme'].inputs]  # Dataset with only inputs to bus
+
+                # Filter by component - get flows of a component
+                results['Boiler'].inputs  # list of input flow labels
+                results['Boiler'].outputs  # list of output flow labels
         """
+        warnings.warn(
+            'results.sizes() is deprecated. '
+            'Use results.plot.all_sizes instead (returns Dataset, not DataArray). '
+            'Note: The new API has no filtering parameters and uses flow labels as variable names. '
+            f'Will be removed in v{DEPRECATION_REMOVAL_VERSION}.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if not self._has_flow_data:
+            raise ValueError('Flow data is not available in this results object (pre-v2.2.0).')
         if self._sizes is None:
             self._sizes = self._assign_flow_coords(
                 xr.concat(
@@ -461,11 +735,12 @@ class CalculationResults:
 
     def _assign_flow_coords(self, da: xr.DataArray):
         # Add start and end coordinates
+        flows_list = list(self.flows.values())
         da = da.assign_coords(
             {
-                'start': ('flow', [flow.start for flow in self.flows.values()]),
-                'end': ('flow', [flow.end for flow in self.flows.values()]),
-                'component': ('flow', [flow.component for flow in self.flows.values()]),
+                'start': ('flow', [flow.start for flow in flows_list]),
+                'end': ('flow', [flow.end for flow in flows_list]),
+                'component': ('flow', [flow.component for flow in flows_list]),
             }
         )
 
@@ -553,7 +828,7 @@ class CalculationResults:
         Args:
             element: The element identifier for which to calculate total effects.
             effect: The effect identifier to calculate.
-            mode: The calculation mode. Options are:
+            mode: The optimization mode. Options are:
                 'temporal': Returns temporal effects.
                 'periodic': Returns investment-specific effects.
                 'total': Returns the sum of temporal effects and periodic effects. Defaults to 'total'.
@@ -584,8 +859,6 @@ class CalculationResults:
             temporal = temporal.sum('time')
             if periodic.isnull().all():
                 return temporal.rename(f'{element}->{effect}')
-            if 'time' in temporal.indexes:
-                temporal = temporal.sum('time')
             return periodic + temporal
 
         total = xr.DataArray(0)
@@ -619,42 +892,57 @@ class CalculationResults:
             total = xr.DataArray(np.nan)
         return total.rename(f'{element}->{effect}({mode})')
 
+    def _create_template_for_mode(self, mode: Literal['temporal', 'periodic', 'total']) -> xr.DataArray:
+        """Create a template DataArray with the correct dimensions for a given mode.
+
+        Args:
+            mode: The optimization mode ('temporal', 'periodic', or 'total').
+
+        Returns:
+            A DataArray filled with NaN, with dimensions appropriate for the mode.
+        """
+        coords = {}
+        if mode == 'temporal':
+            coords['time'] = self.timesteps_extra
+        if self.periods is not None:
+            coords['period'] = self.periods
+        if self.scenarios is not None:
+            coords['scenario'] = self.scenarios
+
+        # Create template with appropriate shape
+        if coords:
+            shape = tuple(len(coords[dim]) for dim in coords)
+            return xr.DataArray(np.full(shape, np.nan, dtype=float), coords=coords, dims=list(coords.keys()))
+        else:
+            return xr.DataArray(np.nan)
+
     def _create_effects_dataset(self, mode: Literal['temporal', 'periodic', 'total']) -> xr.Dataset:
         """Creates a dataset containing effect totals for all components (including their flows).
         The dataset does contain the direct as well as the indirect effects of each component.
 
         Args:
-            mode: The calculation mode ('temporal', 'periodic', or 'total').
+            mode: The optimization mode ('temporal', 'periodic', or 'total').
 
         Returns:
             An xarray Dataset with components as dimension and effects as variables.
         """
+        # Create template with correct dimensions for this mode
+        template = self._create_template_for_mode(mode)
+
         ds = xr.Dataset()
         all_arrays = {}
-        template = None  # Template is needed to determine the dimensions of the arrays. This handles the case of no shares for an effect
-
         components_list = list(self.components)
 
-        # First pass: collect arrays and find template
+        # Collect arrays for all effects and components
         for effect in self.effects:
             effect_arrays = []
             for component in components_list:
                 da = self._compute_effect_total(element=component, effect=effect, mode=mode, include_flows=True)
                 effect_arrays.append(da)
 
-                if template is None and (da.dims or not da.isnull().all()):
-                    template = da
-
             all_arrays[effect] = effect_arrays
 
-        # Ensure we have a template
-        if template is None:
-            raise ValueError(
-                f"No template with proper dimensions found for mode '{mode}'. "
-                f'All computed arrays are scalars, which indicates a data issue.'
-            )
-
-        # Second pass: process all effects (guaranteed to include all)
+        # Process all effects: expand scalar NaN arrays to match template dimensions
         for effect in self.effects:
             dataarrays = all_arrays[effect]
             component_arrays = []
@@ -687,68 +975,136 @@ class CalculationResults:
 
     def plot_heatmap(
         self,
-        variable_name: str,
-        heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
-        heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
-        color_map: str = 'portland',
+        variable_name: str | list[str],
         save: bool | pathlib.Path = False,
-        show: bool = True,
+        show: bool | None = None,
+        colors: plotting.ColorType | None = None,
         engine: plotting.PlottingEngine = 'plotly',
-        indexer: dict[FlowSystemDimensions, Any] | None = None,
+        select: dict[FlowSystemDimensions, Any] | None = None,
+        facet_by: str | list[str] | None = 'scenario',
+        animate_by: str | None = 'period',
+        facet_cols: int | None = None,
+        reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
+        | Literal['auto']
+        | None = 'auto',
+        fill: Literal['ffill', 'bfill'] | None = 'ffill',
+        **plot_kwargs: Any,
     ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
         """
-        Plots a heatmap of the solution of a variable.
+        Plots a heatmap visualization of a variable using imshow or time-based reshaping.
+
+        Supports multiple visualization features that can be combined:
+        - **Multi-variable**: Plot multiple variables on a single heatmap (creates 'variable' dimension)
+        - **Time reshaping**: Converts 'time' dimension into 2D (e.g., hours vs days)
+        - **Faceting**: Creates subplots for different dimension values
+        - **Animation**: Animates through dimension values (Plotly only)
 
         Args:
-            variable_name: The name of the variable to plot.
-            heatmap_timeframes: The timeframes to use for the heatmap.
-            heatmap_timesteps_per_frame: The timesteps per frame to use for the heatmap.
-            color_map: The color map to use for the heatmap.
+            variable_name: The name of the variable to plot, or a list of variable names.
+                When a list is provided, variables are combined into a single DataArray
+                with a new 'variable' dimension.
             save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
             show: Whether to show the plot or not.
+            colors: Color scheme for the heatmap. See `flixopt.plotting.ColorType` for options.
             engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
-            indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-                 If None, uses first value for each dimension.
-                 If empty dict {}, uses all values.
+            select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
+                Applied BEFORE faceting/animation/reshaping.
+            facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
+                or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
+            animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
+                dimension values. Only one dimension can be animated. Ignored if not found.
+            facet_cols: Number of columns in the facet grid layout (default: 3).
+            reshape_time: Time reshaping configuration (default: 'auto'):
+                - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
+                - Tuple: Explicit reshaping, e.g. ('D', 'h') for days vs hours,
+                         ('MS', 'D') for months vs days, ('W', 'h') for weeks vs hours
+                - None: Disable auto-reshaping (will error if only 1D time data)
+                Supported timeframes: 'YS', 'MS', 'W', 'D', 'h', '15min', 'min'
+            fill: Method to fill missing values after reshape: 'ffill' (forward fill) or 'bfill' (backward fill).
+                Default is 'ffill'.
+            **plot_kwargs: Additional plotting customization options.
+                Common options:
+
+                - **dpi** (int): Export resolution for saved plots. Default: 300.
+
+                For heatmaps specifically:
+
+                - **vmin** (float): Minimum value for color scale (both engines).
+                - **vmax** (float): Maximum value for color scale (both engines).
+
+                For Matplotlib heatmaps:
+
+                - **imshow_kwargs** (dict): Additional kwargs for matplotlib's imshow (e.g., interpolation, aspect).
+                - **cbar_kwargs** (dict): Additional kwargs for colorbar customization.
 
         Examples:
-            Basic usage (uses first scenario, first period, all time):
+            Direct imshow mode (default):
+
+            >>> results.plot_heatmap('Battery|charge_state', select={'scenario': 'base'})
 
-            >>> results.plot_heatmap('Battery|charge_state')
+            Facet by scenario:
 
-            Select specific scenario and period:
+            >>> results.plot_heatmap('Boiler(Qth)|flow_rate', facet_by='scenario', facet_cols=2)
 
-            >>> results.plot_heatmap('Boiler(Qth)|flow_rate', indexer={'scenario': 'base', 'period': 2024})
+            Animate by period:
 
-            Time filtering (summer months only):
+            >>> results.plot_heatmap('Boiler(Qth)|flow_rate', select={'scenario': 'base'}, animate_by='period')
+
+            Time reshape mode - daily patterns:
+
+            >>> results.plot_heatmap('Boiler(Qth)|flow_rate', select={'scenario': 'base'}, reshape_time=('D', 'h'))
+
+            Combined: time reshaping with faceting and animation:
 
             >>> results.plot_heatmap(
-            ...     'Boiler(Qth)|flow_rate',
-            ...     indexer={
-            ...         'scenario': 'base',
-            ...         'time': results.solution.time[results.solution.time.dt.month.isin([6, 7, 8])],
-            ...     },
+            ...     'Boiler(Qth)|flow_rate', facet_by='scenario', animate_by='period', reshape_time=('D', 'h')
             ... )
 
-            Save to specific location:
+            Multi-variable heatmap (variables as one axis):
 
             >>> results.plot_heatmap(
-            ...     'Boiler(Qth)|flow_rate', indexer={'scenario': 'base'}, save='path/to/my_heatmap.html'
+            ...     ['Boiler(Q_th)|flow_rate', 'CHP(Q_th)|flow_rate', 'HeatStorage|charge_state'],
+            ...     select={'scenario': 'base', 'period': 1},
+            ...     reshape_time=None,
             ... )
-        """
-        dataarray = self.solution[variable_name]
 
+            Multi-variable with time reshaping:
+
+            >>> results.plot_heatmap(
+            ...     ['Boiler(Q_th)|flow_rate', 'CHP(Q_th)|flow_rate'],
+            ...     facet_by='scenario',
+            ...     animate_by='period',
+            ...     reshape_time=('D', 'h'),
+            ... )
+
+            High-resolution export with custom color range:
+
+            >>> results.plot_heatmap('Battery|charge_state', save=True, dpi=600, vmin=0, vmax=100)
+
+            Matplotlib heatmap with custom imshow settings:
+
+            >>> results.plot_heatmap(
+            ...     'Boiler(Q_th)|flow_rate',
+            ...     engine='matplotlib',
+            ...     imshow_kwargs={'interpolation': 'bilinear', 'aspect': 'auto'},
+            ... )
+        """
+        # Delegate to module-level plot_heatmap function
         return plot_heatmap(
-            dataarray=dataarray,
-            name=variable_name,
+            data=self.solution[variable_name],
+            name=variable_name if isinstance(variable_name, str) else None,
             folder=self.folder,
-            heatmap_timeframes=heatmap_timeframes,
-            heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
-            color_map=color_map,
+            colors=colors,
             save=save,
             show=show,
             engine=engine,
-            indexer=indexer,
+            select=select,
+            facet_by=facet_by,
+            animate_by=animate_by,
+            facet_cols=facet_cols,
+            reshape_time=reshape_time,
+            fill=fill,
+            **plot_kwargs,
         )
 
     def plot_network(
@@ -760,19 +1116,74 @@ class CalculationResults:
             ]
         ) = True,
         path: pathlib.Path | None = None,
-        show: bool = False,
+        show: bool | None = None,
     ) -> pyvis.network.Network | None:
         """Plot interactive network visualization of the system.
 
         Args:
             controls: Enable/disable interactive controls.
             path: Save path for network HTML.
-            show: Whether to display the plot.
+            show: Whether to display the plot. If None, uses CONFIG.Plotting.default_show.
         """
         if path is None:
             path = self.folder / f'{self.name}--network.html'
         return self.flow_system.plot_network(controls=controls, path=path, show=show)
 
+    def to_flow_system(self) -> FlowSystem:
+        """Convert Results to a FlowSystem with solution attached.
+
+        This method migrates results from the deprecated Results format to the
+        new FlowSystem-based format, enabling use of the modern API.
+
+        Note:
+            For loading old results files directly, consider using
+            ``FlowSystem.from_old_results(folder, name)`` instead.
+
+        Returns:
+            FlowSystem: A FlowSystem instance with the solution data attached.
+
+        Caveats:
+            - The linopy model is NOT attached (only the solution data)
+            - Element submodels are NOT recreated (no re-optimization without
+              calling build_model() first)
+            - Variable/constraint names on elements are NOT restored
+
+        Examples:
+            Convert loaded Results to FlowSystem:
+
+            ```python
+            # Load old results
+            results = Results.from_file('results', 'my_optimization')
+
+            # Convert to FlowSystem
+            flow_system = results.to_flow_system()
+
+            # Use new API
+            flow_system.plot.heatmap()
+            flow_system.solution.to_netcdf('solution.nc')
+
+            # Save in new single-file format
+            flow_system.to_netcdf('my_optimization.nc')
+            ```
+        """
+        from flixopt.io import convert_old_dataset
+
+        # Convert flow_system_data to new parameter names
+        convert_old_dataset(self.flow_system_data)
+
+        # Reconstruct FlowSystem from stored data
+        flow_system = FlowSystem.from_dataset(self.flow_system_data)
+
+        # Convert solution attrs from dicts to JSON strings for consistency with new format
+        # The _get_solution_attr helper handles both formats, but we normalize here
+        solution = self.solution.copy()
+        for key in ['Components', 'Buses', 'Effects', 'Flows']:
+            if key in solution.attrs and isinstance(solution.attrs[key], dict):
+                solution.attrs[key] = json.dumps(solution.attrs[key])
+
+        flow_system.solution = solution
+        return flow_system
+
     def to_file(
         self,
         folder: str | pathlib.Path | None = None,
@@ -780,59 +1191,70 @@ class CalculationResults:
         compression: int = 5,
         document_model: bool = True,
         save_linopy_model: bool = False,
+        overwrite: bool = False,
     ):
         """Save results to files.
 
         Args:
-            folder: Save folder (defaults to calculation folder).
-            name: File name (defaults to calculation name).
+            folder: Save folder (defaults to optimization folder).
+            name: File name (defaults to optimization name).
             compression: Compression level 0-9.
             document_model: Whether to document model formulations as yaml.
             save_linopy_model: Whether to save linopy model file.
+            overwrite: If False, raise error if results files already exist. If True, overwrite existing files.
+
+        Raises:
+            FileExistsError: If overwrite=False and result files already exist.
         """
         folder = self.folder if folder is None else pathlib.Path(folder)
         name = self.name if name is None else name
-        if not folder.exists():
-            try:
-                folder.mkdir(parents=False)
-            except FileNotFoundError as e:
-                raise FileNotFoundError(
-                    f'Folder {folder} and its parent do not exist. Please create them first.'
-                ) from e
 
-        paths = fx_io.CalculationResultsPaths(folder, name)
+        # Ensure folder exists, creating parent directories as needed
+        folder.mkdir(parents=True, exist_ok=True)
+
+        paths = fx_io.ResultsPaths(folder, name)
+
+        # Check if files already exist (unless overwrite is True)
+        if not overwrite:
+            existing_files = []
+            for file_path in paths.all_paths().values():
+                if file_path.exists():
+                    existing_files.append(file_path.name)
+
+            if existing_files:
+                raise FileExistsError(
+                    f'Results files already exist in {folder}: {", ".join(existing_files)}. '
+                    f'Use overwrite=True to overwrite existing files.'
+                )
 
         fx_io.save_dataset_to_netcdf(self.solution, paths.solution, compression=compression)
         fx_io.save_dataset_to_netcdf(self.flow_system_data, paths.flow_system, compression=compression)
 
-        with open(paths.summary, 'w', encoding='utf-8') as f:
-            yaml.dump(self.summary, f, allow_unicode=True, sort_keys=False, indent=4, width=1000)
+        fx_io.save_yaml(self.summary, paths.summary, compact_numeric_lists=True)
 
         if save_linopy_model:
             if self.model is None:
-                logger.critical('No model in the CalculationResults. Saving the model is not possible.')
+                logger.critical('No model in the Results. Saving the model is not possible.')
             else:
-                self.model.to_netcdf(paths.linopy_model, engine='h5netcdf')
+                self.model.to_netcdf(paths.linopy_model, engine='netcdf4')
 
         if document_model:
             if self.model is None:
-                logger.critical('No model in the CalculationResults. Documenting the model is not possible.')
+                logger.critical('No model in the Results. Documenting the model is not possible.')
             else:
                 fx_io.document_linopy_model(self.model, path=paths.model_documentation)
 
-        logger.info(f'Saved calculation results "{name}" to {paths.model_documentation.parent}')
+        logger.log(SUCCESS_LEVEL, f'Saved optimization results "{name}" to {paths.model_documentation.parent}')
 
 
 class _ElementResults:
-    def __init__(
-        self, calculation_results: CalculationResults, label: str, variables: list[str], constraints: list[str]
-    ):
-        self._calculation_results = calculation_results
+    def __init__(self, results: Results, label: str, variables: list[str], constraints: list[str]):
+        self._results = results
         self.label = label
-        self._variable_names = variables
+        self.variable_names = variables
         self._constraint_names = constraints
 
-        self.solution = self._calculation_results.solution[self._variable_names]
+        self.solution = self._results.solution[self.variable_names]
 
     @property
     def variables(self) -> linopy.Variables:
@@ -841,9 +1263,9 @@ class _ElementResults:
         Raises:
             ValueError: If linopy model is unavailable.
         """
-        if self._calculation_results.model is None:
+        if self._results.model is None:
             raise ValueError('The linopy model is not available.')
-        return self._calculation_results.model.variables[self._variable_names]
+        return self._results.model.variables[self.variable_names]
 
     @property
     def constraints(self) -> linopy.Constraints:
@@ -852,9 +1274,17 @@ class _ElementResults:
         Raises:
             ValueError: If linopy model is unavailable.
         """
-        if self._calculation_results.model is None:
+        if self._results.model is None:
             raise ValueError('The linopy model is not available.')
-        return self._calculation_results.model.constraints[self._constraint_names]
+        return self._results.model.constraints[self._constraint_names]
+
+    def __repr__(self) -> str:
+        """Return string representation with element info and dataset preview."""
+        class_name = self.__class__.__name__
+        header = f'{class_name}: "{self.label}"'
+        sol = self.solution.copy(deep=False)
+        sol.attrs = {}
+        return f'{header}\n{"-" * len(header)}\n{repr(sol)}'
 
     def filter_solution(
         self,
@@ -901,7 +1331,7 @@ class _ElementResults:
 class _NodeResults(_ElementResults):
     def __init__(
         self,
-        calculation_results: CalculationResults,
+        results: Results,
         label: str,
         variables: list[str],
         constraints: list[str],
@@ -909,7 +1339,7 @@ class _NodeResults(_ElementResults):
         outputs: list[str],
         flows: list[str],
     ):
-        super().__init__(calculation_results, label, variables, constraints)
+        super().__init__(results, label, variables, constraints)
         self.inputs = inputs
         self.outputs = outputs
         self.flows = flows
@@ -917,75 +1347,194 @@ class _NodeResults(_ElementResults):
     def plot_node_balance(
         self,
         save: bool | pathlib.Path = False,
-        show: bool = True,
-        colors: plotting.ColorType = 'viridis',
+        show: bool | None = None,
+        colors: plotting.ColorType | None = None,
         engine: plotting.PlottingEngine = 'plotly',
-        indexer: dict[FlowSystemDimensions, Any] | None = None,
-        mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
-        style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
+        select: dict[FlowSystemDimensions, Any] | None = None,
+        unit_type: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
+        mode: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
         drop_suffix: bool = True,
+        facet_by: str | list[str] | None = 'scenario',
+        animate_by: str | None = 'period',
+        facet_cols: int | None = None,
+        **plot_kwargs: Any,
     ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
         """
-        Plots the node balance of the Component or Bus.
+        Plots the node balance of the Component or Bus with optional faceting and animation.
+
         Args:
             save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
             show: Whether to show the plot or not.
             colors: The colors to use for the plot. See `flixopt.plotting.ColorType` for options.
             engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
-            indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-                 If None, uses first value for each dimension (except time).
-                 If empty dict {}, uses all values.
-            style: The style to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
+            select: Optional data selection dict. Supports:
+                - Single values: {'scenario': 'base', 'period': 2024}
+                - Multiple values: {'scenario': ['base', 'high', 'renewable']}
+                - Slices: {'time': slice('2024-01', '2024-06')}
+                - Index arrays: {'time': time_array}
+                Note: Applied BEFORE faceting/animation.
+            unit_type: The unit type to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
                 - 'flow_rate': Returns the flow_rates of the Node.
                 - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
+            mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for stepped lines, or 'area' for stacked area charts.
             drop_suffix: Whether to drop the suffix from the variable names.
+            facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
+                or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
+                Example: 'scenario' creates one subplot per scenario.
+                Example: ['period', 'scenario'] creates a grid of subplots for each scenario-period combination.
+            animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
+                dimension values. Only one dimension can be animated. Ignored if not found.
+            facet_cols: Number of columns in the facet grid layout (default: 3).
+            **plot_kwargs: Additional plotting customization options passed to underlying plotting functions.
+
+                Common options:
+
+                - **dpi** (int): Export resolution in dots per inch. Default: 300.
+
+                **For Plotly engine** (`engine='plotly'`):
+
+                - Any Plotly Express parameter for px.bar()/px.line()/px.area()
+                  Example: `range_y=[0, 100]`, `line_shape='linear'`
+
+                **For Matplotlib engine** (`engine='matplotlib'`):
+
+                - **plot_kwargs** (dict): Customize plot via `ax.bar()` or `ax.step()`.
+                  Example: `plot_kwargs={'linewidth': 3, 'alpha': 0.7, 'edgecolor': 'black'}`
+
+                See :func:`flixopt.plotting.with_plotly` and :func:`flixopt.plotting.with_matplotlib`
+                for complete parameter reference.
+
+                Note: For Plotly, you can further customize the returned figure using `fig.update_traces()`
+                and `fig.update_layout()` after calling this method.
+
+        Examples:
+            Basic plot (current behavior):
+
+            >>> results['Boiler'].plot_node_balance()
+
+            Facet by scenario:
+
+            >>> results['Boiler'].plot_node_balance(facet_by='scenario', facet_cols=2)
+
+            Animate by period:
+
+            >>> results['Boiler'].plot_node_balance(animate_by='period')
+
+            Facet by scenario AND animate by period:
+
+            >>> results['Boiler'].plot_node_balance(facet_by='scenario', animate_by='period')
+
+            Select single scenario, then facet by period:
+
+            >>> results['Boiler'].plot_node_balance(select={'scenario': 'base'}, facet_by='period')
+
+            Select multiple scenarios and facet by them:
+
+            >>> results['Boiler'].plot_node_balance(
+            ...     select={'scenario': ['base', 'high', 'renewable']}, facet_by='scenario'
+            ... )
+
+            Time range selection (summer months only):
+
+            >>> results['Boiler'].plot_node_balance(select={'time': slice('2024-06', '2024-08')}, facet_by='scenario')
+
+            High-resolution export for publication:
+
+            >>> results['Boiler'].plot_node_balance(engine='matplotlib', save='figure.png', dpi=600)
+
+            Plotly Express customization (e.g., set y-axis range):
+
+            >>> results['Boiler'].plot_node_balance(range_y=[0, 100])
+
+            Custom matplotlib appearance:
+
+            >>> results['Boiler'].plot_node_balance(engine='matplotlib', plot_kwargs={'linewidth': 3, 'alpha': 0.7})
+
+            Further customize Plotly figure after creation:
+
+            >>> fig = results['Boiler'].plot_node_balance(mode='line', show=False)
+            >>> fig.update_traces(line={'width': 5, 'dash': 'dot'})
+            >>> fig.update_layout(template='plotly_dark', width=1200, height=600)
+            >>> fig.show()
         """
-        ds = self.node_balance(with_last_timestep=True, mode=mode, drop_suffix=drop_suffix, indexer=indexer)
+        if engine not in {'plotly', 'matplotlib'}:
+            raise ValueError(f'Engine "{engine}" not supported. Use one of ["plotly", "matplotlib"]')
+
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+
+        # Don't pass select/indexer to node_balance - we'll apply it afterwards
+        ds = self.node_balance(with_last_timestep=False, unit_type=unit_type, drop_suffix=drop_suffix)
+
+        ds, suffix_parts = _apply_selection_to_data(ds, select=select, drop=True)
 
-        ds, suffix_parts = _apply_indexer_to_data(ds, indexer, drop=True)
+        # Matplotlib requires only 'time' dimension; check for extras after selection
+        if engine == 'matplotlib':
+            extra_dims = [d for d in ds.dims if d != 'time']
+            if extra_dims:
+                raise ValueError(
+                    f'Matplotlib engine only supports a single time axis, but found extra dimensions: {extra_dims}. '
+                    f'Please use select={{...}} to reduce dimensions or switch to engine="plotly" for faceting/animation.'
+                )
         suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
 
-        title = f'{self.label} (flow rates){suffix}' if mode == 'flow_rate' else f'{self.label} (flow hours){suffix}'
+        title = (
+            f'{self.label} (flow rates){suffix}' if unit_type == 'flow_rate' else f'{self.label} (flow hours){suffix}'
+        )
 
         if engine == 'plotly':
             figure_like = plotting.with_plotly(
-                ds.to_dataframe(),
-                colors=colors,
-                style=style,
+                ds,
+                facet_by=facet_by,
+                animate_by=animate_by,
+                colors=colors if colors is not None else self._results.colors,
+                mode=mode,
                 title=title,
+                facet_cols=facet_cols,
+                xlabel='Time in h',
+                **plot_kwargs,
             )
             default_filetype = '.html'
-        elif engine == 'matplotlib':
+        else:
             figure_like = plotting.with_matplotlib(
-                ds.to_dataframe(),
-                colors=colors,
-                style=style,
+                ds,
+                colors=colors if colors is not None else self._results.colors,
+                mode=mode,
                 title=title,
+                **plot_kwargs,
             )
             default_filetype = '.png'
-        else:
-            raise ValueError(f'Engine "{engine}" not supported. Use "plotly" or "matplotlib"')
 
         return plotting.export_figure(
             figure_like=figure_like,
-            default_path=self._calculation_results.folder / title,
+            default_path=self._results.folder / title,
             default_filetype=default_filetype,
             user_path=None if isinstance(save, bool) else pathlib.Path(save),
             show=show,
             save=True if save else False,
+            dpi=dpi,
         )
 
     def plot_node_balance_pie(
         self,
         lower_percentage_group: float = 5,
-        colors: plotting.ColorType = 'viridis',
+        colors: plotting.ColorType | None = None,
         text_info: str = 'percent+label+value',
         save: bool | pathlib.Path = False,
-        show: bool = True,
+        show: bool | None = None,
         engine: plotting.PlottingEngine = 'plotly',
-        indexer: dict[FlowSystemDimensions, Any] | None = None,
+        select: dict[FlowSystemDimensions, Any] | None = None,
+        **plot_kwargs: Any,
     ) -> plotly.graph_objs.Figure | tuple[plt.Figure, list[plt.Axes]]:
         """Plot pie chart of flow hours distribution.
+
+        Note:
+            Pie charts require scalar data (no extra dimensions beyond time).
+            If your data has dimensions like 'scenario' or 'period', either:
+
+            - Use `select` to choose specific values: `select={'scenario': 'base', 'period': 2024}`
+            - Let auto-selection choose the first value (a warning will be logged)
+
         Args:
             lower_percentage_group: Percentage threshold for "Others" grouping.
             colors: Color scheme. Also see plotly.
@@ -993,44 +1542,107 @@ class _NodeResults(_ElementResults):
             save: Whether to save plot.
             show: Whether to display plot.
             engine: Plotting engine ('plotly' or 'matplotlib').
-            indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-                 If None, uses first value for each dimension.
-                 If empty dict {}, uses all values.
+            select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
+                Use this to select specific scenario/period before creating the pie chart.
+            **plot_kwargs: Additional plotting customization options.
+
+                Common options:
+
+                - **dpi** (int): Export resolution in dots per inch. Default: 300.
+                - **hover_template** (str): Hover text template (Plotly only).
+                  Example: `hover_template='%{label}: %{value} (%{percent})'`
+                - **text_position** (str): Text position ('inside', 'outside', 'auto').
+                - **hole** (float): Size of donut hole (0.0 to 1.0).
+
+                See :func:`flixopt.plotting.dual_pie_with_plotly` for complete reference.
+
+        Examples:
+            Basic usage (auto-selects first scenario/period if present):
+
+            >>> results['Bus'].plot_node_balance_pie()
+
+            Explicitly select a scenario and period:
+
+            >>> results['Bus'].plot_node_balance_pie(select={'scenario': 'high_demand', 'period': 2030})
+
+            Create a donut chart with custom hover text:
+
+            >>> results['Bus'].plot_node_balance_pie(hole=0.4, hover_template='%{label}: %{value:.2f} (%{percent})')
+
+            High-resolution export:
+
+            >>> results['Bus'].plot_node_balance_pie(save='figure.png', dpi=600)
         """
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+
         inputs = sanitize_dataset(
-            ds=self.solution[self.inputs] * self._calculation_results.hours_per_timestep,
+            ds=self.solution[self.inputs] * self._results.timestep_duration,
             threshold=1e-5,
             drop_small_vars=True,
             zero_small_values=True,
             drop_suffix='|',
         )
         outputs = sanitize_dataset(
-            ds=self.solution[self.outputs] * self._calculation_results.hours_per_timestep,
+            ds=self.solution[self.outputs] * self._results.timestep_duration,
             threshold=1e-5,
             drop_small_vars=True,
             zero_small_values=True,
             drop_suffix='|',
         )
 
-        inputs, suffix_parts = _apply_indexer_to_data(inputs, indexer, drop=True)
-        outputs, suffix_parts = _apply_indexer_to_data(outputs, indexer, drop=True)
-        suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
-
-        title = f'{self.label} (total flow hours){suffix}'
+        inputs, suffix_parts_in = _apply_selection_to_data(inputs, select=select, drop=True)
+        outputs, suffix_parts_out = _apply_selection_to_data(outputs, select=select, drop=True)
+        suffix_parts = suffix_parts_in + suffix_parts_out
 
+        # Sum over time dimension
         inputs = inputs.sum('time')
         outputs = outputs.sum('time')
 
+        # Auto-select first value for any remaining dimensions (scenario, period, etc.)
+        # Pie charts need scalar data, so we automatically reduce extra dimensions
+        extra_dims_inputs = [dim for dim in inputs.dims if dim != 'time']
+        extra_dims_outputs = [dim for dim in outputs.dims if dim != 'time']
+        extra_dims = sorted(set(extra_dims_inputs + extra_dims_outputs))
+
+        if extra_dims:
+            auto_select = {}
+            for dim in extra_dims:
+                # Get first value of this dimension
+                if dim in inputs.coords:
+                    first_val = inputs.coords[dim].values[0]
+                elif dim in outputs.coords:
+                    first_val = outputs.coords[dim].values[0]
+                else:
+                    continue
+                auto_select[dim] = first_val
+                logger.info(
+                    f'Pie chart auto-selected {dim}={first_val} (first value). '
+                    f'Use select={{"{dim}": value}} to choose a different value.'
+                )
+
+            # Apply auto-selection only for coords present in each dataset
+            inputs = inputs.sel({k: v for k, v in auto_select.items() if k in inputs.coords})
+            outputs = outputs.sel({k: v for k, v in auto_select.items() if k in outputs.coords})
+
+            # Update suffix with auto-selected values
+            auto_suffix_parts = [f'{dim}={val}' for dim, val in auto_select.items()]
+            suffix_parts.extend(auto_suffix_parts)
+
+        suffix = '--' + '-'.join(sorted(set(suffix_parts))) if suffix_parts else ''
+        title = f'{self.label} (total flow hours){suffix}'
+
         if engine == 'plotly':
             figure_like = plotting.dual_pie_with_plotly(
-                data_left=inputs.to_pandas(),
-                data_right=outputs.to_pandas(),
-                colors=colors,
+                data_left=inputs,
+                data_right=outputs,
+                colors=colors if colors is not None else self._results.colors,
                 title=title,
                 text_info=text_info,
                 subtitles=('Inputs', 'Outputs'),
                 legend_title='Flows',
                 lower_percentage_group=lower_percentage_group,
+                **plot_kwargs,
             )
             default_filetype = '.html'
         elif engine == 'matplotlib':
@@ -1038,11 +1650,12 @@ class _NodeResults(_ElementResults):
             figure_like = plotting.dual_pie_with_matplotlib(
                 data_left=inputs.to_pandas(),
                 data_right=outputs.to_pandas(),
-                colors=colors,
+                colors=colors if colors is not None else self._results.colors,
                 title=title,
                 subtitles=('Inputs', 'Outputs'),
                 legend_title='Flows',
                 lower_percentage_group=lower_percentage_group,
+                **plot_kwargs,
             )
             default_filetype = '.png'
         else:
@@ -1050,11 +1663,12 @@ class _NodeResults(_ElementResults):
 
         return plotting.export_figure(
             figure_like=figure_like,
-            default_path=self._calculation_results.folder / title,
+            default_path=self._results.folder / title,
             default_filetype=default_filetype,
             user_path=None if isinstance(save, bool) else pathlib.Path(save),
             show=show,
             save=True if save else False,
+            dpi=dpi,
         )
 
     def node_balance(
@@ -1063,9 +1677,9 @@ class _NodeResults(_ElementResults):
         negate_outputs: bool = False,
         threshold: float | None = 1e-5,
         with_last_timestep: bool = False,
-        mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
+        unit_type: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
         drop_suffix: bool = False,
-        indexer: dict[FlowSystemDimensions, Any] | None = None,
+        select: dict[FlowSystemDimensions, Any] | None = None,
     ) -> xr.Dataset:
         """
         Returns a dataset with the node balance of the Component or Bus.
@@ -1074,20 +1688,18 @@ class _NodeResults(_ElementResults):
             negate_outputs: Whether to negate the output flow_rates of the Node.
             threshold: The threshold for small values. Variables with all values below the threshold are dropped.
             with_last_timestep: Whether to include the last timestep in the dataset.
-            mode: The mode to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
+            unit_type: The unit type to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
                 - 'flow_rate': Returns the flow_rates of the Node.
                 - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
             drop_suffix: Whether to drop the suffix from the variable names.
-            indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-                 If None, uses first value for each dimension.
-                 If empty dict {}, uses all values.
+            select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
         """
         ds = self.solution[self.inputs + self.outputs]
 
         ds = sanitize_dataset(
             ds=ds,
             threshold=threshold,
-            timesteps=self._calculation_results.timesteps_extra if with_last_timestep else None,
+            timesteps=self._results.timesteps_extra if with_last_timestep else None,
             negate=(
                 self.outputs + self.inputs
                 if negate_outputs and negate_inputs
@@ -1100,10 +1712,10 @@ class _NodeResults(_ElementResults):
             drop_suffix='|' if drop_suffix else None,
         )
 
-        ds, _ = _apply_indexer_to_data(ds, indexer, drop=True)
+        ds, _ = _apply_selection_to_data(ds, select=select, drop=True)
 
-        if mode == 'flow_hours':
-            ds = ds * self._calculation_results.hours_per_timestep
+        if unit_type == 'flow_hours':
+            ds = ds * self._results.timestep_duration
             ds = ds.rename_vars({var: var.replace('flow_rate', 'flow_hours') for var in ds.data_vars})
 
         return ds
@@ -1118,7 +1730,7 @@ class ComponentResults(_NodeResults):
 
     @property
     def is_storage(self) -> bool:
-        return self._charge_state in self._variable_names
+        return self._charge_state in self.variable_names
 
     @property
     def _charge_state(self) -> str:
@@ -1134,75 +1746,202 @@ class ComponentResults(_NodeResults):
     def plot_charge_state(
         self,
         save: bool | pathlib.Path = False,
-        show: bool = True,
-        colors: plotting.ColorType = 'viridis',
+        show: bool | None = None,
+        colors: plotting.ColorType | None = None,
         engine: plotting.PlottingEngine = 'plotly',
-        style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
-        indexer: dict[FlowSystemDimensions, Any] | None = None,
+        mode: Literal['area', 'stacked_bar', 'line'] = 'area',
+        select: dict[FlowSystemDimensions, Any] | None = None,
+        facet_by: str | list[str] | None = 'scenario',
+        animate_by: str | None = 'period',
+        facet_cols: int | None = None,
+        **plot_kwargs: Any,
     ) -> plotly.graph_objs.Figure:
-        """Plot storage charge state over time, combined with the node balance.
+        """Plot storage charge state over time, combined with the node balance with optional faceting and animation.
 
         Args:
             save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
             show: Whether to show the plot or not.
             colors: Color scheme. Also see plotly.
             engine: Plotting engine to use. Only 'plotly' is implemented atm.
-            style: The colors to use for the plot. See `flixopt.plotting.ColorType` for options.
-            indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-                 If None, uses first value for each dimension.
-                 If empty dict {}, uses all values.
+            mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for stepped lines, or 'area' for stacked area charts.
+            select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
+                Applied BEFORE faceting/animation.
+            facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
+                or list of dimensions. Each unique value combination creates a subplot. Ignored if not found.
+            animate_by: Dimension to animate over (Plotly only). Creates animation frames that cycle through
+                dimension values. Only one dimension can be animated. Ignored if not found.
+            facet_cols: Number of columns in the facet grid layout (default: 3).
+            **plot_kwargs: Additional plotting customization options passed to underlying plotting functions.
+
+                Common options:
+
+                - **dpi** (int): Export resolution in dots per inch. Default: 300.
+
+                **For Plotly engine:**
+
+                - Any Plotly Express parameter for px.bar()/px.line()/px.area()
+                  Example: `range_y=[0, 100]`, `line_shape='linear'`
+
+                **For Matplotlib engine:**
+
+                - **plot_kwargs** (dict): Customize plot via `ax.bar()` or `ax.step()`.
+
+                See :func:`flixopt.plotting.with_plotly` and :func:`flixopt.plotting.with_matplotlib`
+                for complete parameter reference.
+
+                Note: For Plotly, you can further customize the returned figure using `fig.update_traces()`
+                and `fig.update_layout()` after calling this method.
 
         Raises:
             ValueError: If component is not a storage.
+
+        Examples:
+            Basic plot:
+
+            >>> results['Storage'].plot_charge_state()
+
+            Facet by scenario:
+
+            >>> results['Storage'].plot_charge_state(facet_by='scenario', facet_cols=2)
+
+            Animate by period:
+
+            >>> results['Storage'].plot_charge_state(animate_by='period')
+
+            Facet by scenario AND animate by period:
+
+            >>> results['Storage'].plot_charge_state(facet_by='scenario', animate_by='period')
+
+            Custom layout after creation:
+
+            >>> fig = results['Storage'].plot_charge_state(show=False)
+            >>> fig.update_layout(template='plotly_dark', height=800)
+            >>> fig.show()
+
+            High-resolution export:
+
+            >>> results['Storage'].plot_charge_state(save='storage.png', dpi=600)
         """
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+
+        # Extract charge state line color (for overlay customization)
+        overlay_color = plot_kwargs.pop('charge_state_line_color', 'black')
+
         if not self.is_storage:
             raise ValueError(f'Cant plot charge_state. "{self.label}" is not a storage')
 
-        ds = self.node_balance(with_last_timestep=True, indexer=indexer)
-        charge_state = self.charge_state
+        # Get node balance and charge state
+        ds = self.node_balance(with_last_timestep=True).fillna(0)
+        charge_state_da = self.charge_state
 
-        ds, suffix_parts = _apply_indexer_to_data(ds, indexer, drop=True)
-        charge_state, suffix_parts = _apply_indexer_to_data(charge_state, indexer, drop=True)
+        # Apply select filtering
+        ds, suffix_parts = _apply_selection_to_data(ds, select=select, drop=True)
+        charge_state_da, _ = _apply_selection_to_data(charge_state_da, select=select, drop=True)
         suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
 
         title = f'Operation Balance of {self.label}{suffix}'
 
         if engine == 'plotly':
-            fig = plotting.with_plotly(
-                ds.to_dataframe(),
-                colors=colors,
-                style=style,
+            # Plot flows (node balance) with the specified mode
+            figure_like = plotting.with_plotly(
+                ds,
+                facet_by=facet_by,
+                animate_by=animate_by,
+                colors=colors if colors is not None else self._results.colors,
+                mode=mode,
                 title=title,
+                facet_cols=facet_cols,
+                xlabel='Time in h',
+                **plot_kwargs,
             )
 
-            # TODO: Use colors for charge state?
-
-            charge_state = charge_state.to_dataframe()
-            fig.add_trace(
-                plotly.graph_objs.Scatter(
-                    x=charge_state.index, y=charge_state.values.flatten(), mode='lines', name=self._charge_state
-                )
+            # Prepare charge_state as Dataset for plotting
+            charge_state_ds = xr.Dataset({self._charge_state: charge_state_da})
+
+            # Plot charge_state with mode='line' to get Scatter traces
+            charge_state_fig = plotting.with_plotly(
+                charge_state_ds,
+                facet_by=facet_by,
+                animate_by=animate_by,
+                colors=colors if colors is not None else self._results.colors,
+                mode='line',  # Always line for charge_state
+                title='',  # No title needed for this temp figure
+                facet_cols=facet_cols,
+                xlabel='Time in h',
+                **plot_kwargs,
             )
+
+            # Add charge_state traces to the main figure
+            # This preserves subplot assignments and animation frames
+            for trace in charge_state_fig.data:
+                trace.line.width = 2  # Make charge_state line more prominent
+                trace.line.shape = 'linear'  # Smooth line for charge state (not stepped like flows)
+                trace.line.color = overlay_color
+                figure_like.add_trace(trace)
+
+            # Also add traces from animation frames if they exist
+            # Both figures use the same animate_by parameter, so they should have matching frames
+            if hasattr(charge_state_fig, 'frames') and charge_state_fig.frames:
+                # Add charge_state traces to each frame
+                for i, frame in enumerate(charge_state_fig.frames):
+                    if i < len(figure_like.frames):
+                        for trace in frame.data:
+                            trace.line.width = 2
+                            trace.line.shape = 'linear'  # Smooth line for charge state
+                            trace.line.color = overlay_color
+                            figure_like.frames[i].data = figure_like.frames[i].data + (trace,)
+
+            default_filetype = '.html'
         elif engine == 'matplotlib':
+            # Matplotlib requires only 'time' dimension; check for extras after selection
+            extra_dims = [d for d in ds.dims if d != 'time']
+            if extra_dims:
+                raise ValueError(
+                    f'Matplotlib engine only supports a single time axis, but found extra dimensions: {extra_dims}. '
+                    f'Please use select={{...}} to reduce dimensions or switch to engine="plotly" for faceting/animation.'
+                )
+            # For matplotlib, plot flows (node balance), then add charge_state as line
             fig, ax = plotting.with_matplotlib(
-                ds.to_dataframe(),
-                colors=colors,
-                style=style,
+                ds,
+                colors=colors if colors is not None else self._results.colors,
+                mode=mode,
                 title=title,
+                **plot_kwargs,
             )
 
-            charge_state = charge_state.to_dataframe()
-            ax.plot(charge_state.index, charge_state.values.flatten(), label=self._charge_state)
+            # Add charge_state as a line overlay
+            charge_state_df = charge_state_da.to_dataframe()
+            ax.plot(
+                charge_state_df.index,
+                charge_state_df.values.flatten(),
+                label=self._charge_state,
+                linewidth=2,
+                color=overlay_color,
+            )
+            # Recreate legend with the same styling as with_matplotlib
+            handles, labels = ax.get_legend_handles_labels()
+            ax.legend(
+                handles,
+                labels,
+                loc='upper center',
+                bbox_to_anchor=(0.5, -0.15),
+                ncol=5,
+                frameon=False,
+            )
             fig.tight_layout()
-            fig = fig, ax
+
+            figure_like = fig, ax
+            default_filetype = '.png'
 
         return plotting.export_figure(
-            fig,
-            default_path=self._calculation_results.folder / title,
-            default_filetype='.html',
+            figure_like=figure_like,
+            default_path=self._results.folder / title,
+            default_filetype=default_filetype,
             user_path=None if isinstance(save, bool) else pathlib.Path(save),
             show=show,
             save=True if save else False,
+            dpi=dpi,
         )
 
     def node_balance_with_charge_state(
@@ -1227,7 +1966,7 @@ class ComponentResults(_NodeResults):
         return sanitize_dataset(
             ds=self.solution[variable_names],
             threshold=threshold,
-            timesteps=self._calculation_results.timesteps_extra,
+            timesteps=self._results.timesteps_extra,
             negate=(
                 self.outputs + self.inputs
                 if negate_outputs and negate_inputs
@@ -1252,13 +1991,13 @@ class EffectResults(_ElementResults):
         Returns:
             xr.Dataset: Element shares to this effect.
         """
-        return self.solution[[name for name in self._variable_names if name.startswith(f'{element}->')]]
+        return self.solution[[name for name in self.variable_names if name.startswith(f'{element}->')]]
 
 
 class FlowResults(_ElementResults):
     def __init__(
         self,
-        calculation_results: CalculationResults,
+        results: Results,
         label: str,
         variables: list[str],
         constraints: list[str],
@@ -1266,7 +2005,7 @@ class FlowResults(_ElementResults):
         end: str,
         component: str,
     ):
-        super().__init__(calculation_results, label, variables, constraints)
+        super().__init__(results, label, variables, constraints)
         self.start = start
         self.end = end
         self.component = component
@@ -1277,7 +2016,7 @@ class FlowResults(_ElementResults):
 
     @property
     def flow_hours(self) -> xr.DataArray:
-        return (self.flow_rate * self._calculation_results.hours_per_timestep).rename(f'{self.label}|flow_hours')
+        return (self.flow_rate * self._results.timestep_duration).rename(f'{self.label}|flow_hours')
 
     @property
     def size(self) -> xr.DataArray:
@@ -1285,16 +2024,16 @@ class FlowResults(_ElementResults):
         if name in self.solution:
             return self.solution[name]
         try:
-            return self._calculation_results.flow_system.flows[self.label].size.rename(name)
+            return self._results.flow_system.flows[self.label].size.rename(name)
         except _FlowSystemRestorationError:
             logger.critical(f'Size of flow {self.label}.size not availlable. Returning NaN')
             return xr.DataArray(np.nan).rename(name)
 
 
-class SegmentedCalculationResults:
-    """Results container for segmented optimization calculations with temporal decomposition.
+class SegmentedResults:
+    """Results container for segmented optimization optimizations with temporal decomposition.
 
-    This class manages results from SegmentedCalculation runs where large optimization
+    This class manages results from SegmentedOptimization runs where large optimization
     problems are solved by dividing the time horizon into smaller, overlapping segments.
     It provides unified access to results across all segments while maintaining the
     ability to analyze individual segment behavior.
@@ -1317,8 +2056,8 @@ class SegmentedCalculationResults:
         Load and analyze segmented results:
 
         ```python
-        # Load segmented calculation results
-        results = SegmentedCalculationResults.from_file('results', 'annual_segmented')
+        # Load segmented optimization results
+        results = SegmentedResults.from_file('results', 'annual_segmented')
 
         # Access unified results across all segments
         full_timeline = results.all_timesteps
@@ -1334,20 +2073,20 @@ class SegmentedCalculationResults:
         max_discontinuity = segment_boundaries['max_storage_jump']
         ```
 
-        Create from segmented calculation:
+        Create from segmented optimization:
 
         ```python
-        # After running segmented calculation
-        segmented_calc = SegmentedCalculation(
+        # After running segmented optimization
+        segmented_opt = SegmentedOptimization(
             name='annual_system',
             flow_system=system,
             timesteps_per_segment=730,  # Monthly segments
             overlap_timesteps=48,  # 2-day overlap
         )
-        segmented_calc.do_modeling_and_solve(solver='gurobi')
+        segmented_opt.do_modeling_and_solve(solver='gurobi')
 
         # Extract unified results
-        results = SegmentedCalculationResults.from_calculation(segmented_calc)
+        results = SegmentedResults.from_optimization(segmented_opt)
 
         # Save combined results
         results.to_file(compression=5)
@@ -1388,34 +2127,50 @@ class SegmentedCalculationResults:
     """
 
     @classmethod
-    def from_calculation(cls, calculation: SegmentedCalculation):
+    def from_optimization(cls, optimization: SegmentedOptimization) -> SegmentedResults:
+        """Create SegmentedResults from a SegmentedOptimization instance.
+
+        Args:
+            optimization: The SegmentedOptimization instance to extract results from.
+
+        Returns:
+            SegmentedResults: New instance containing the optimization results.
+        """
         return cls(
-            [calc.results for calc in calculation.sub_calculations],
-            all_timesteps=calculation.all_timesteps,
-            timesteps_per_segment=calculation.timesteps_per_segment,
-            overlap_timesteps=calculation.overlap_timesteps,
-            name=calculation.name,
-            folder=calculation.folder,
+            [calc.results for calc in optimization.sub_optimizations],
+            all_timesteps=optimization.all_timesteps,
+            timesteps_per_segment=optimization.timesteps_per_segment,
+            overlap_timesteps=optimization.overlap_timesteps,
+            name=optimization.name,
+            folder=optimization.folder,
         )
 
     @classmethod
-    def from_file(cls, folder: str | pathlib.Path, name: str) -> SegmentedCalculationResults:
-        """Load SegmentedCalculationResults from saved files.
+    def from_file(cls, folder: str | pathlib.Path, name: str) -> SegmentedResults:
+        """Load SegmentedResults from saved files.
 
         Args:
             folder: Directory containing saved files.
             name: Base name of saved files.
 
         Returns:
-            SegmentedCalculationResults: Loaded instance.
+            SegmentedResults: Loaded instance.
         """
         folder = pathlib.Path(folder)
         path = folder / name
-        logger.info(f'loading calculation "{name}" from file ("{path.with_suffix(".nc4")}")')
-        with open(path.with_suffix('.json'), encoding='utf-8') as f:
-            meta_data = json.load(f)
+        meta_data_path = path.with_suffix('.json')
+        logger.info(f'loading segemented optimization meta data from file ("{meta_data_path}")')
+        meta_data = fx_io.load_json(meta_data_path)
+
+        # Handle both new 'sub_optimizations' and legacy 'sub_calculations' keys
+        sub_names = meta_data.get('sub_optimizations') or meta_data.get('sub_calculations')
+        if sub_names is None:
+            raise KeyError(
+                "Missing 'sub_optimizations' (or legacy 'sub_calculations') key in segmented results metadata."
+            )
+
         return cls(
-            [CalculationResults.from_file(folder, sub_name) for sub_name in meta_data['sub_calculations']],
+            [Results.from_file(folder, sub_name) for sub_name in sub_names],
             all_timesteps=pd.DatetimeIndex(
                 [datetime.datetime.fromisoformat(date) for date in meta_data['all_timesteps']], name='time'
             ),
@@ -1427,20 +2182,26 @@ class SegmentedCalculationResults:
 
     def __init__(
         self,
-        segment_results: list[CalculationResults],
+        segment_results: list[Results],
         all_timesteps: pd.DatetimeIndex,
         timesteps_per_segment: int,
         overlap_timesteps: int,
         name: str,
         folder: pathlib.Path | None = None,
     ):
+        warnings.warn(
+            f'SegmentedResults is deprecated and will be removed in v{DEPRECATION_REMOVAL_VERSION}. '
+            'A replacement API for segmented optimization will be provided in a future release.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
         self.segment_results = segment_results
         self.all_timesteps = all_timesteps
         self.timesteps_per_segment = timesteps_per_segment
         self.overlap_timesteps = overlap_timesteps
         self.name = name
         self.folder = pathlib.Path(folder) if folder is not None else pathlib.Path.cwd() / 'results'
-        self.hours_per_timestep = FlowSystem.calculate_hours_per_timestep(self.all_timesteps)
+        self._colors = {}
 
     @property
     def meta_data(self) -> dict[str, int | list[str]]:
@@ -1448,13 +2209,74 @@ class SegmentedCalculationResults:
             'all_timesteps': [datetime.datetime.isoformat(date) for date in self.all_timesteps],
             'timesteps_per_segment': self.timesteps_per_segment,
             'overlap_timesteps': self.overlap_timesteps,
-            'sub_calculations': [calc.name for calc in self.segment_results],
+            'sub_optimizations': [calc.name for calc in self.segment_results],
         }
 
     @property
     def segment_names(self) -> list[str]:
         return [segment.name for segment in self.segment_results]
 
+    @property
+    def colors(self) -> dict[str, str]:
+        return self._colors
+
+    @colors.setter
+    def colors(self, colors: dict[str, str]):
+        """Applies colors to all segments"""
+        self._colors = colors
+        for segment in self.segment_results:
+            segment.colors = copy.deepcopy(colors)
+
+    def setup_colors(
+        self,
+        config: dict[str, str | list[str]] | str | pathlib.Path | None = None,
+        default_colorscale: str | None = None,
+    ) -> dict[str, str]:
+        """
+        Setup colors for all variables across all segment results.
+
+        This method applies the same color configuration to all segments, ensuring
+        consistent visualization across the entire segmented optimization. The color
+        mapping is propagated to each segment's Results instance.
+
+        Args:
+            config: Configuration for color assignment. Can be:
+                - dict: Maps components to colors/colorscales:
+                    * 'component1': 'red'  # Single component to single color
+                    * 'component1': '#FF0000'  # Single component to hex color
+                    - OR maps colorscales to multiple components:
+                    * 'colorscale_name': ['component1', 'component2']  # Colorscale across components
+                - str: Path to a JSON/YAML config file or a colorscale name to apply to all
+                - Path: Path to a JSON/YAML config file
+                - None: Use default_colorscale for all components
+            default_colorscale: Default colorscale for unconfigured components (default: 'turbo')
+
+        Examples:
+            ```python
+            # Apply colors to all segments
+            segmented_results.setup_colors(
+                {
+                    'CHP': 'red',
+                    'Blues': ['Storage1', 'Storage2'],
+                    'Oranges': ['Solar1', 'Solar2'],
+                }
+            )
+
+            # Use a single colorscale for all components in all segments
+            segmented_results.setup_colors('portland')
+            ```
+
+        Returns:
+            Complete variable-to-color mapping dictionary from the first segment
+            (all segments will have the same mapping)
+        """
+        if not self.segment_results:
+            raise ValueError('No segment_results available; cannot setup colors on an empty SegmentedResults.')
+
+        self.colors = self.segment_results[0].setup_colors(config=config, default_colorscale=default_colorscale)
+
+        return self.colors
+
     def solution_without_overlap(self, variable_name: str) -> xr.DataArray:
         """Get variable solution removing segment overlaps.
 
@@ -1473,123 +2295,265 @@ class SegmentedCalculationResults:
     def plot_heatmap(
         self,
         variable_name: str,
-        heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
-        heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
-        color_map: str = 'portland',
+        reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
+        | Literal['auto']
+        | None = 'auto',
+        colors: plotting.ColorType | None = None,
         save: bool | pathlib.Path = False,
-        show: bool = True,
+        show: bool | None = None,
         engine: plotting.PlottingEngine = 'plotly',
+        facet_by: str | list[str] | None = None,
+        animate_by: str | None = None,
+        facet_cols: int | None = None,
+        fill: Literal['ffill', 'bfill'] | None = 'ffill',
+        **plot_kwargs: Any,
     ) -> plotly.graph_objs.Figure | tuple[plt.Figure, plt.Axes]:
         """Plot heatmap of variable solution across segments.
 
         Args:
             variable_name: Variable to plot.
-            heatmap_timeframes: Time aggregation level.
-            heatmap_timesteps_per_frame: Timesteps per frame.
-            color_map: Color scheme. Also see plotly.
+            reshape_time: Time reshaping configuration (default: 'auto'):
+                - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
+                - Tuple like ('D', 'h'): Explicit reshaping (days vs hours)
+                - None: Disable time reshaping
+            colors: Color scheme. See plotting.ColorType for options.
             save: Whether to save plot.
             show: Whether to display plot.
             engine: Plotting engine.
+            facet_by: Dimension(s) to create facets (subplots) for.
+            animate_by: Dimension to animate over (Plotly only).
+            facet_cols: Number of columns in the facet grid layout.
+            fill: Method to fill missing values: 'ffill' or 'bfill'.
+            **plot_kwargs: Additional plotting customization options.
+                Common options:
+
+                - **dpi** (int): Export resolution for saved plots. Default: 300.
+                - **vmin** (float): Minimum value for color scale.
+                - **vmax** (float): Maximum value for color scale.
+
+                For Matplotlib heatmaps:
+
+                - **imshow_kwargs** (dict): Additional kwargs for matplotlib's imshow.
+                - **cbar_kwargs** (dict): Additional kwargs for colorbar customization.
 
         Returns:
             Figure object.
         """
         return plot_heatmap(
-            dataarray=self.solution_without_overlap(variable_name),
+            data=self.solution_without_overlap(variable_name),
             name=variable_name,
             folder=self.folder,
-            heatmap_timeframes=heatmap_timeframes,
-            heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
-            color_map=color_map,
+            reshape_time=reshape_time,
+            colors=colors,
             save=save,
             show=show,
             engine=engine,
+            facet_by=facet_by,
+            animate_by=animate_by,
+            facet_cols=facet_cols,
+            fill=fill,
+            **plot_kwargs,
         )
 
-    def to_file(self, folder: str | pathlib.Path | None = None, name: str | None = None, compression: int = 5):
+    def to_file(
+        self,
+        folder: str | pathlib.Path | None = None,
+        name: str | None = None,
+        compression: int = 5,
+        overwrite: bool = False,
+    ):
         """Save segmented results to files.
 
         Args:
             folder: Save folder (defaults to instance folder).
             name: File name (defaults to instance name).
             compression: Compression level 0-9.
+            overwrite: If False, raise error if results files already exist. If True, overwrite existing files.
+
+        Raises:
+            FileExistsError: If overwrite=False and result files already exist.
         """
         folder = self.folder if folder is None else pathlib.Path(folder)
         name = self.name if name is None else name
         path = folder / name
-        if not folder.exists():
-            try:
-                folder.mkdir(parents=False)
-            except FileNotFoundError as e:
-                raise FileNotFoundError(
-                    f'Folder {folder} and its parent do not exist. Please create them first.'
-                ) from e
+
+        # Ensure folder exists, creating parent directories as needed
+        folder.mkdir(parents=True, exist_ok=True)
+
+        # Check if metadata file already exists (unless overwrite is True)
+        metadata_file = path.with_suffix('.json')
+        if not overwrite and metadata_file.exists():
+            raise FileExistsError(
+                f'Segmented results file already exists: {metadata_file}. '
+                f'Use overwrite=True to overwrite existing files.'
+            )
+
+        # Save segments (they will check for overwrite themselves)
         for segment in self.segment_results:
-            segment.to_file(folder=folder, name=segment.name, compression=compression)
+            segment.to_file(folder=folder, name=segment.name, compression=compression, overwrite=overwrite)
 
-        with open(path.with_suffix('.json'), 'w', encoding='utf-8') as f:
-            json.dump(self.meta_data, f, indent=4, ensure_ascii=False)
-        logger.info(f'Saved calculation "{name}" to {path}')
+        fx_io.save_json(self.meta_data, metadata_file)
+        logger.info(f'Saved optimization "{name}" to {path}')
 
 
 def plot_heatmap(
-    dataarray: xr.DataArray,
-    name: str,
-    folder: pathlib.Path,
-    heatmap_timeframes: Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'] = 'D',
-    heatmap_timesteps_per_frame: Literal['W', 'D', 'h', '15min', 'min'] = 'h',
-    color_map: str = 'portland',
+    data: xr.DataArray | xr.Dataset,
+    name: str | None = None,
+    folder: pathlib.Path | None = None,
+    colors: plotting.ColorType | None = None,
     save: bool | pathlib.Path = False,
-    show: bool = True,
+    show: bool | None = None,
     engine: plotting.PlottingEngine = 'plotly',
-    indexer: dict[str, Any] | None = None,
+    select: dict[str, Any] | None = None,
+    facet_by: str | list[str] | None = None,
+    animate_by: str | None = None,
+    facet_cols: int | None = None,
+    reshape_time: tuple[Literal['YS', 'MS', 'W', 'D', 'h', '15min', 'min'], Literal['W', 'D', 'h', '15min', 'min']]
+    | Literal['auto']
+    | None = 'auto',
+    fill: Literal['ffill', 'bfill'] | None = 'ffill',
+    **plot_kwargs: Any,
 ):
-    """Plot heatmap of time series data.
+    """Plot heatmap visualization with support for multi-variable, faceting, and animation.
+
+    This function provides a standalone interface to the heatmap plotting capabilities,
+    supporting the same modern features as Results.plot_heatmap().
 
     Args:
-        dataarray: Data to plot.
-        name: Variable name for title.
-        folder: Save folder.
-        heatmap_timeframes: Time aggregation level.
-        heatmap_timesteps_per_frame: Timesteps per frame.
-        color_map: Color scheme. Also see plotly.
-        save: Whether to save plot.
-        show: Whether to display plot.
-        engine: Plotting engine.
-        indexer: Optional selection dict, e.g., {'scenario': 'base', 'period': 2024}.
-             If None, uses first value for each dimension.
-             If empty dict {}, uses all values.
+        data: Data to plot. Can be a single DataArray or an xarray Dataset.
+            When a Dataset is provided, all data variables are combined along a new 'variable' dimension.
+        name: Optional name for the title. If not provided, uses the DataArray name or
+            generates a default title for Datasets.
+        folder: Save folder for the plot. Defaults to current directory if not provided.
+        colors: Color scheme for the heatmap. See `flixopt.plotting.ColorType` for options.
+        save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
+        show: Whether to show the plot or not.
+        engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
+        select: Optional data selection dict. Supports single values, lists, slices, and index arrays.
+        facet_by: Dimension(s) to create facets (subplots) for. Can be a single dimension name (str)
+            or list of dimensions. Each unique value combination creates a subplot.
+        animate_by: Dimension to animate over (Plotly only). Creates animation frames.
+        facet_cols: Number of columns in the facet grid layout (default: 3).
+        reshape_time: Time reshaping configuration (default: 'auto'):
+            - 'auto': Automatically applies ('D', 'h') when only 'time' dimension remains
+            - Tuple: Explicit reshaping, e.g. ('D', 'h') for days vs hours
+            - None: Disable auto-reshaping
+        fill: Method to fill missing values after reshape: 'ffill' (forward fill) or 'bfill' (backward fill).
+            Default is 'ffill'.
+
+    Examples:
+        Single DataArray with time reshaping:
+
+        >>> plot_heatmap(data, name='Temperature', folder=Path('.'), reshape_time=('D', 'h'))
+
+        Dataset with multiple variables (facet by variable):
+
+        >>> dataset = xr.Dataset({'Boiler': data1, 'CHP': data2, 'Storage': data3})
+        >>> plot_heatmap(
+        ...     dataset,
+        ...     folder=Path('.'),
+        ...     facet_by='variable',
+        ...     reshape_time=('D', 'h'),
+        ... )
+
+        Dataset with animation by variable:
+
+        >>> plot_heatmap(dataset, animate_by='variable', reshape_time=('D', 'h'))
     """
-    dataarray, suffix_parts = _apply_indexer_to_data(dataarray, indexer, drop=True)
+    # Convert Dataset to DataArray with 'variable' dimension
+    if isinstance(data, xr.Dataset):
+        # Extract all data variables from the Dataset
+        variable_names = list(data.data_vars)
+        dataarrays = [data[var] for var in variable_names]
+
+        # Combine into single DataArray with 'variable' dimension
+        data = xr.concat(dataarrays, dim='variable')
+        data = data.assign_coords(variable=variable_names)
+
+        # Use Dataset variable names for title if name not provided
+        if name is None:
+            title_name = f'Heatmap of {len(variable_names)} variables'
+        else:
+            title_name = name
+    else:
+        # Single DataArray
+        if name is None:
+            title_name = data.name if data.name else 'Heatmap'
+        else:
+            title_name = name
+
+    # Apply select filtering
+    data, suffix_parts = _apply_selection_to_data(data, select=select, drop=True)
     suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
-    name = name if not suffix_parts else name + suffix
 
-    heatmap_data = plotting.heat_map_data_from_df(
-        dataarray.to_dataframe(name), heatmap_timeframes, heatmap_timesteps_per_frame, 'ffill'
-    )
+    # Matplotlib heatmaps require at most 2D data
+    # Time dimension will be reshaped to 2D (timeframe × timestep), so can't have other dims alongside it
+    if engine == 'matplotlib':
+        dims = list(data.dims)
 
-    xlabel, ylabel = f'timeframe [{heatmap_timeframes}]', f'timesteps [{heatmap_timesteps_per_frame}]'
+        # If 'time' dimension exists and will be reshaped, we can't have any other dimensions
+        if 'time' in dims and len(dims) > 1 and reshape_time is not None:
+            extra_dims = [d for d in dims if d != 'time']
+            raise ValueError(
+                f'Matplotlib heatmaps with time reshaping cannot have additional dimensions. '
+                f'Found extra dimensions: {extra_dims}. '
+                f'Use select={{...}} to reduce to time only, use "reshape_time=None" or switch to engine="plotly" or use for multi-dimensional support.'
+            )
+        # If no 'time' dimension (already reshaped or different data), allow at most 2 dimensions
+        elif 'time' not in dims and len(dims) > 2:
+            raise ValueError(
+                f'Matplotlib heatmaps support at most 2 dimensions, but data has {len(dims)}: {dims}. '
+                f'Use select={{...}} to reduce dimensions or switch to engine="plotly".'
+            )
+
+    # Build title
+    title = f'{title_name}{suffix}'
+    if isinstance(reshape_time, tuple):
+        timeframes, timesteps_per_frame = reshape_time
+        title += f' ({timeframes} vs {timesteps_per_frame})'
 
+    # Extract dpi before passing to plotting functions
+    dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+
+    # Plot with appropriate engine
     if engine == 'plotly':
-        figure_like = plotting.heat_map_plotly(
-            heatmap_data, title=name, color_map=color_map, xlabel=xlabel, ylabel=ylabel
+        figure_like = plotting.heatmap_with_plotly(
+            data=data,
+            facet_by=facet_by,
+            animate_by=animate_by,
+            colors=colors,
+            title=title,
+            facet_cols=facet_cols,
+            reshape_time=reshape_time,
+            fill=fill,
+            **plot_kwargs,
         )
         default_filetype = '.html'
     elif engine == 'matplotlib':
-        figure_like = plotting.heat_map_matplotlib(
-            heatmap_data, title=name, color_map=color_map, xlabel=xlabel, ylabel=ylabel
+        figure_like = plotting.heatmap_with_matplotlib(
+            data=data,
+            colors=colors,
+            title=title,
+            reshape_time=reshape_time,
+            fill=fill,
+            **plot_kwargs,
         )
         default_filetype = '.png'
     else:
         raise ValueError(f'Engine "{engine}" not supported. Use "plotly" or "matplotlib"')
 
+    # Set default folder if not provided
+    if folder is None:
+        folder = pathlib.Path('.')
+
     return plotting.export_figure(
         figure_like=figure_like,
-        default_path=folder / f'{name} ({heatmap_timeframes}-{heatmap_timesteps_per_frame})',
+        default_path=folder / title,
         default_filetype=default_filetype,
         user_path=None if isinstance(save, bool) else pathlib.Path(save),
         show=show,
         save=True if save else False,
+        dpi=dpi,
     )
 
 
@@ -1787,8 +2751,13 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
         if coord_name not in array.coords:
             raise AttributeError(f"Missing required coordinate '{coord_name}'")
 
-        # Convert single value to list
-        val_list = [coord_values] if isinstance(coord_values, str) else coord_values
+        # Normalize to list for sequence-like inputs (excluding strings)
+        if isinstance(coord_values, str):
+            val_list = [coord_values]
+        elif isinstance(coord_values, (list, tuple, np.ndarray, pd.Index)):
+            val_list = list(coord_values)
+        else:
+            val_list = [coord_values]
 
         # Verify coord_values exist
         available = set(array[coord_name].values)
@@ -1798,7 +2767,7 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
 
         # Apply filter
         return array.where(
-            array[coord_name].isin(val_list) if isinstance(coord_values, list) else array[coord_name] == coord_values,
+            array[coord_name].isin(val_list) if len(val_list) > 1 else array[coord_name] == val_list[0],
             drop=True,
         )
 
@@ -1817,36 +2786,26 @@ def filter_dataarray_by_coord(da: xr.DataArray, **kwargs: str | list[str] | None
     return da
 
 
-def _apply_indexer_to_data(
-    data: xr.DataArray | xr.Dataset, indexer: dict[str, Any] | None = None, drop=False
+def _apply_selection_to_data(
+    data: xr.DataArray | xr.Dataset,
+    select: dict[str, Any] | None = None,
+    drop=False,
 ) -> tuple[xr.DataArray | xr.Dataset, list[str]]:
     """
-    Apply indexer selection or auto-select first values for non-time dimensions.
+    Apply selection to data.
 
     Args:
         data: xarray Dataset or DataArray
-        indexer: Optional selection dict
-            If None, uses first value for each dimension (except time).
-            If empty dict {}, uses all values.
+        select: Optional selection dict
+        drop: Whether to drop dimensions after selection
 
     Returns:
         Tuple of (selected_data, selection_string)
     """
     selection_string = []
 
-    if indexer is not None:
-        # User provided indexer
-        data = data.sel(indexer, drop=drop)
-        selection_string.extend(f'{v}[{k}]' for k, v in indexer.items())
-    else:
-        # Auto-select first value for each dimension except 'time'
-        selection = {}
-        for dim in data.dims:
-            if dim != 'time' and dim in data.coords:
-                first_value = data.coords[dim].values[0]
-                selection[dim] = first_value
-                selection_string.append(f'{first_value}[{dim}]')
-        if selection:
-            data = data.sel(selection, drop=drop)
+    if select:
+        data = data.sel(select, drop=drop)
+        selection_string.extend(f'{dim}={val}' for dim, val in select.items())
 
     return data, selection_string