Add better and more versitile sankey diagrams

flixopt/config.py

@@ -20,7 +20,7 @@
     COLORLOG_AVAILABLE = False
     escape_codes = None
 
-__all__ = ['CONFIG', 'MultilineFormatter', 'SUCCESS_LEVEL']
+__all__ = ['CONFIG', 'MultilineFormatter', 'SUCCESS_LEVEL', 'DEPRECATION_REMOVAL_VERSION']
 
 if COLORLOG_AVAILABLE:
     __all__.append('ColoredMultilineFormatter')

flixopt/flow_system.py

@@ -1241,13 +1241,7 @@ def plot_network(
 
         Visualizes the network structure of a FlowSystem using PyVis.
         """
-        warnings.warn(
-            f'plot_network() is deprecated and will be removed in v{DEPRECATION_REMOVAL_VERSION}. '
-            'Use flow_system.topology.plot() instead.',
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.topology.plot(path=path, controls=controls, show=show)
+        return self.topology.plot_legacy(path=path, controls=controls, show=show)
 
     def start_network_app(self) -> None:
         """

flixopt/statistics_accessor.py

@@ -417,7 +417,11 @@ def sizes(self) -> xr.Dataset:
         """All flow sizes as a Dataset with flow labels as variable names."""
         self._require_solution()
         if self._sizes is None:
-            size_vars = [v for v in self._fs.solution.data_vars if v.endswith('|size')]
+            # Get flow labels to filter only flow sizes (not storage capacity sizes)
+            flow_labels = set(self._fs.flows.keys())
+            size_vars = [
+                v for v in self._fs.solution.data_vars if v.endswith('|size') and v.replace('|size', '') in flow_labels
+            ]
             self._sizes = xr.Dataset({v.replace('|size', ''): self._fs.solution[v] for v in size_vars})
         return self._sizes
 
@@ -1094,69 +1098,186 @@ def flows(
 
         return PlotResult(data=ds, figure=fig)
 
-    def sankey(
+    def _prepare_sankey_data(
         self,
-        *,
-        timestep: int | str | None = None,
-        aggregate: Literal['sum', 'mean'] = 'sum',
-        select: SelectType | None = None,
-        colors: ColorType | None = None,
-        show: bool | None = None,
-        **plotly_kwargs: Any,
-    ) -> PlotResult:
-        """Plot Sankey diagram of energy/material flow hours.
+        mode: Literal['flow_hours', 'sizes', 'peak_flow'],
+        timestep: int | str | None,
+        aggregate: Literal['sum', 'mean'],
+        select: SelectType | None,
+    ) -> tuple[xr.Dataset, str]:
+        """Prepare data for Sankey diagram based on mode.
 
         Args:
-            timestep: Specific timestep to show, or None for aggregation.
-            aggregate: How to aggregate if timestep is None.
+            mode: What to display - flow_hours, sizes, or peak_flow.
+            timestep: Specific timestep (only for flow_hours mode).
+            aggregate: Aggregation method (only for flow_hours mode).
             select: xarray-style selection.
-            colors: Color specification for nodes (colorscale name, color list, or label-to-color dict).
-            show: Whether to display.
 
         Returns:
-            PlotResult with Sankey flow data.
+            Tuple of (prepared Dataset, title string).
         """
-        self._stats._require_solution()
-
-        ds = self._stats.flow_hours.copy()
-
-        # Apply weights
-        if 'period' in ds.dims and self._fs.period_weights is not None:
-            ds = ds * self._fs.period_weights
-        if 'scenario' in ds.dims and self._fs.scenario_weights is not None:
-            weights = self._fs.scenario_weights / self._fs.scenario_weights.sum()
-            ds = ds * weights
+        if mode == 'sizes':
+            ds = self._stats.sizes.copy()
+            title = 'Investment Sizes (Capacities)'
+        elif mode == 'peak_flow':
+            ds = self._stats.flow_rates.copy()
+            ds = _apply_selection(ds, select)
+            if 'time' in ds.dims:
+                ds = ds.max(dim='time')
+            for dim in ['period', 'scenario']:
+                if dim in ds.dims:
+                    ds = ds.max(dim=dim)
+            return ds, 'Peak Flow Rates'
+        else:  # flow_hours
+            ds = self._stats.flow_hours.copy()
+            title = 'Energy Flow'
+
+        # Apply weights for flow_hours
+        if mode == 'flow_hours':
+            if 'period' in ds.dims and self._fs.period_weights is not None:
+                ds = ds * self._fs.period_weights
+            if 'scenario' in ds.dims and self._fs.scenario_weights is not None:
+                weights = self._fs.scenario_weights / self._fs.scenario_weights.sum()
+                ds = ds * weights
 
         ds = _apply_selection(ds, select)
 
-        if timestep is not None:
-            if isinstance(timestep, int):
-                ds = ds.isel(time=timestep)
-            else:
-                ds = ds.sel(time=timestep)
-        elif 'time' in ds.dims:
-            ds = getattr(ds, aggregate)(dim='time')
+        # Time aggregation (only for flow_hours)
+        if mode == 'flow_hours':
+            if timestep is not None:
+                if isinstance(timestep, int):
+                    ds = ds.isel(time=timestep)
+                else:
+                    ds = ds.sel(time=timestep)
+            elif 'time' in ds.dims:
+                ds = getattr(ds, aggregate)(dim='time')
 
+        # Collapse remaining dimensions
         for dim in ['period', 'scenario']:
             if dim in ds.dims:
-                ds = ds.sum(dim=dim)
+                ds = ds.sum(dim=dim) if mode == 'flow_hours' else ds.max(dim=dim)
+
+        return ds, title
+
+    def _build_effects_sankey(
+        self,
+        select: SelectType | None,
+        colors: ColorType | None,
+        **plotly_kwargs: Any,
+    ) -> tuple[go.Figure, xr.Dataset]:
+        """Build Sankey diagram showing contributions from components to effects.
+
+        Creates a Sankey with:
+        - Left side: Components (grouped by type)
+        - Right side: Effects (costs, CO2, etc.)
+        - Links: Contributions from each component to each effect
+
+        Args:
+            select: xarray-style selection.
+            colors: Color specification for nodes.
+            **plotly_kwargs: Additional Plotly layout arguments.
+
+        Returns:
+            Tuple of (Plotly Figure, Dataset with link data).
+        """
+        total_effects = self._stats.total_effects
+
+        # Collect all links: component -> effect
+        nodes: set[str] = set()
+        links: dict[str, list] = {'source': [], 'target': [], 'value': [], 'label': []}
+
+        for effect_name in total_effects.data_vars:
+            effect_data = total_effects[effect_name]
+            effect_data = _apply_selection(effect_data, select)
+
+            # Sum over any remaining dimensions
+            for dim in ['period', 'scenario']:
+                if dim in effect_data.dims:
+                    effect_data = effect_data.sum(dim=dim)
+
+            contributors = effect_data.coords['contributor'].values
+            components = effect_data.coords['component'].values
+
+            for contributor, component in zip(contributors, components, strict=False):
+                value = float(effect_data.sel(contributor=contributor).values)
+                if not np.isfinite(value) or abs(value) < 1e-6:
+                    continue
+
+                # Use component as source node, effect as target
+                source = str(component)
+                target = f'[{effect_name}]'  # Bracket notation to distinguish effects
+
+                nodes.add(source)
+                nodes.add(target)
+                links['source'].append(source)
+                links['target'].append(target)
+                links['value'].append(abs(value))
+                links['label'].append(f'{contributor} → {effect_name}: {value:.2f}')
+
+        # Create figure
+        node_list = list(nodes)
+        node_indices = {n: i for i, n in enumerate(node_list)}
+
+        color_map = process_colors(colors, node_list)
+        node_colors = [color_map[node] for node in node_list]
+
+        fig = go.Figure(
+            data=[
+                go.Sankey(
+                    node=dict(
+                        pad=15, thickness=20, line=dict(color='black', width=0.5), label=node_list, color=node_colors
+                    ),
+                    link=dict(
+                        source=[node_indices[s] for s in links['source']],
+                        target=[node_indices[t] for t in links['target']],
+                        value=links['value'],
+                        label=links['label'],
+                    ),
+                )
+            ]
+        )
+        fig.update_layout(title='Effect Contributions by Component', **plotly_kwargs)
+
+        sankey_ds = xr.Dataset(
+            {'value': ('link', links['value'])},
+            coords={
+                'link': range(len(links['value'])),
+                'source': ('link', links['source']),
+                'target': ('link', links['target']),
+                'label': ('link', links['label']),
+            },
+        )
+
+        return fig, sankey_ds
+
+    def _build_sankey_links(
+        self,
+        ds: xr.Dataset,
+        min_value: float = 1e-6,
+    ) -> tuple[set[str], dict[str, list]]:
+        """Build Sankey nodes and links from flow data.
 
-        # Build Sankey
-        nodes = set()
-        links = {'source': [], 'target': [], 'value': [], 'label': []}
+        Args:
+            ds: Dataset with flow values (one variable per flow).
+            min_value: Minimum value threshold to include a link.
+
+        Returns:
+            Tuple of (nodes set, links dict with source/target/value/label).
+        """
+        nodes: set[str] = set()
+        links: dict[str, list] = {'source': [], 'target': [], 'value': [], 'label': []}
 
         for flow in self._fs.flows.values():
             label = flow.label_full
             if label not in ds:
                 continue
             value = float(ds[label].values)
-            if abs(value) < 1e-6:
+            if abs(value) < min_value:
                 continue
 
-            # Determine source/target based on flow direction
-            # is_input_in_component: True means bus -> component, False means component -> bus
+            # flow.bus and flow.component are already strings (bus label, component label_full)
             bus_label = flow.bus
-            comp_label = flow.component.label_full
+            comp_label = flow.component
 
             if flow.is_input_in_component:
                 source = bus_label
@@ -1172,6 +1293,28 @@ def sankey(
             links['value'].append(abs(value))
             links['label'].append(label)
 
+        return nodes, links
+
+    def _create_sankey_figure(
+        self,
+        nodes: set[str],
+        links: dict[str, list],
+        colors: ColorType | None,
+        title: str,
+        **plotly_kwargs: Any,
+    ) -> go.Figure:
+        """Create Plotly Sankey figure.
+
+        Args:
+            nodes: Set of node labels.
+            links: Dict with source, target, value, label lists.
+            colors: Color specification for nodes.
+            title: Figure title.
+            **plotly_kwargs: Additional Plotly layout arguments.
+
+        Returns:
+            Plotly Figure with Sankey diagram.
+        """
         node_list = list(nodes)
         node_indices = {n: i for i, n in enumerate(node_list)}
 
@@ -1193,12 +1336,75 @@ def sankey(
                 )
             ]
         )
-        fig.update_layout(title='Energy Flow Sankey', **plotly_kwargs)
+        fig.update_layout(title=title, **plotly_kwargs)
+        return fig
 
-        sankey_ds = xr.Dataset(
-            {'value': ('link', links['value'])},
-            coords={'link': links['label'], 'source': ('link', links['source']), 'target': ('link', links['target'])},
-        )
+    def sankey(
+        self,
+        *,
+        mode: Literal['flow_hours', 'sizes', 'peak_flow', 'effects'] = 'flow_hours',
+        timestep: int | str | None = None,
+        aggregate: Literal['sum', 'mean'] = 'sum',
+        select: SelectType | None = None,
+        max_size: float | None = None,
+        colors: ColorType | None = None,
+        show: bool | None = None,
+        **plotly_kwargs: Any,
+    ) -> PlotResult:
+        """Plot Sankey diagram of the flow system.
+
+        Args:
+            mode: What to display:
+                - 'flow_hours': Energy/material amounts (default)
+                - 'sizes': Investment capacities
+                - 'peak_flow': Maximum flow rates
+                - 'effects': Component contributions to all effects (costs, CO2, etc.)
+            timestep: Specific timestep to show, or None for aggregation (flow_hours only).
+            aggregate: How to aggregate if timestep is None ('sum' or 'mean', flow_hours only).
+            select: xarray-style selection.
+            max_size: Filter flows with sizes exceeding this value (sizes mode only).
+            colors: Color specification for nodes (colorscale name, color list, or label-to-color dict).
+            show: Whether to display.
+            **plotly_kwargs: Additional arguments passed to Plotly layout.
+
+        Returns:
+            PlotResult with Sankey flow data and figure.
+
+        Examples:
+            >>> # Show energy flows (default)
+            >>> flow_system.statistics.plot.sankey()
+            >>> # Show investment sizes/capacities
+            >>> flow_system.statistics.plot.sankey(mode='sizes')
+            >>> # Show peak flow rates
+            >>> flow_system.statistics.plot.sankey(mode='peak_flow')
+            >>> # Show effect contributions (components -> effects like costs, CO2)
+            >>> flow_system.statistics.plot.sankey(mode='effects')
+        """
+        self._stats._require_solution()
+
+        if mode == 'effects':
+            fig, sankey_ds = self._build_effects_sankey(select, colors, **plotly_kwargs)
+        else:
+            ds, title = self._prepare_sankey_data(mode, timestep, aggregate, select)
+
+            # Apply max_size filter for sizes mode
+            if max_size is not None and mode == 'sizes' and ds.data_vars:
+                valid_labels = [lbl for lbl in ds.data_vars if float(ds[lbl].max()) < max_size]
+                ds = ds[valid_labels]
+
+            nodes, links = self._build_sankey_links(ds)
+            fig = self._create_sankey_figure(nodes, links, colors, title, **plotly_kwargs)
+
+            n_links = len(links['value'])
+            sankey_ds = xr.Dataset(
+                {'value': ('link', links['value'])},
+                coords={
+                    'link': range(n_links),
+                    'source': ('link', links['source']),
+                    'target': ('link', links['target']),
+                    'label': ('link', links['label']),
+                },
+            )
 
         if show is None:
             show = CONFIG.Plotting.default_show