Feature/plotting kwargs and streamline data conversion

[FBumann]

Description

Updating the plotting methods to be cleaner and to allow passing kwargs to poorly express and matplotlib

Type of Change

• Bug fix
• New feature
• Documentation update
• Code refactoring

Related Issues

Closes #(issue number)

Testing

• I have tested my changes
• Existing tests still pass

Checklist

• My code follows the project style
• I have updated documentation if needed
• I have added tests for new functionality (if applicable)

Summary by CodeRabbit

• New Features
  • Plots now accept xarray Datasets (DataFrame/Series still supported), support per‑bus/per‑entity visuals, and allow pass‑through plotting options (backend kwargs, px_kwargs, plot_kwargs, dpi) for finer customization and high‑resolution exports.
• Changed
  • Plot styling now respects templates; improved overlaying of original/aggregated traces, color resolution, axis labeling and selection, and handling of scalar datasets.
• Removed
  • Unused pie plotting helper removed.
• Bug Fixes
  • Clearer errors for multidimensional/non‑numeric inputs and stricter heatmap validation.
• Tests
  • Added plotting API smoke and robustness tests for backends, inputs, and error cases.

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0d857f0 and d5a3e01.

📒 Files selected for processing (7)

• CHANGELOG.md (1 hunks)
• examples/02_Complex/complex_example_results.py (1 hunks)
• examples/03_Calculation_types/example_calculation_types.py (1 hunks)
• flixopt/aggregation.py (1 hunks)
• flixopt/plotting.py (21 hunks)
• flixopt/results.py (36 hunks)
• tests/test_plotting_api.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (3)

• examples/03_Calculation_types/example_calculation_types.py
• CHANGELOG.md
• examples/02_Complex/complex_example_results.py

🧰 Additional context used

🧬 Code graph analysis (4)

tests/test_plotting_api.py (1)

flixopt/plotting.py (1)

• with_plotly (397-690)

flixopt/aggregation.py (1)

flixopt/plotting.py (1)

• with_plotly (397-690)

flixopt/plotting.py (2)

flixopt/config.py (1)

• CONFIG (61-322)

flixopt/aggregation.py (1)

• plot (144-175)

flixopt/results.py (1)

flixopt/plotting.py (1)

• with_matplotlib (693-840)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: test (3.12)
• GitHub Check: test (3.11)
• GitHub Check: test (3.13)
• GitHub Check: test (3.10)

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

flixopt/results.py (2)

1338-1345: Auto-select may KeyError if a dim exists in only one of inputs/outputs.

Select per-dataset on existing coords.

-            # Apply auto-selection
-            inputs = inputs.sel(auto_select)
-            outputs = outputs.sel(auto_select)
+            # 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})

---

1640-1645: Overlay line color not applied to animation frames.

Charge-state traces in frames lose the custom color.

-                        for trace in frame.data:
-                            trace.line.width = 2
-                            trace.line.shape = 'linear'  # Smooth line for charge state
+                        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,)

🧹 Nitpick comments (7)

tests/test_plotting_api.py (2)

133-134: Avoid brittle float equality in pie sum assertion.

Use pytest.approx to tolerate harmless FP rounding.

-                assert sum(fig.data[0].values) == 40
+                assert sum(fig.data[0].values) == pytest.approx(40)

---

12-22: Make fixtures deterministic (optional).

Random inputs aren’t used in assertions now, but a seed prevents future flakiness.

 def sample_dataset():
     """Create a sample xarray Dataset for testing."""
-    time = np.arange(10)
+    rng = np.random.default_rng(0)
+    time = np.arange(10)
     data = xr.Dataset(
         {
-            'var1': (['time'], np.random.rand(10)),
-            'var2': (['time'], np.random.rand(10)),
-            'var3': (['time'], np.random.rand(10)),
+            'var1': (['time'], rng.random(10)),
+            'var2': (['time'], rng.random(10)),
+            'var3': (['time'], rng.random(10)),
         },
         coords={'time': time},
     )
@@
 def sample_dataframe():
     """Create a sample pandas DataFrame for testing."""
-    time = np.arange(10)
-    df = pd.DataFrame({'var1': np.random.rand(10), 'var2': np.random.rand(10), 'var3': np.random.rand(10)}, index=time)
+    rng = np.random.default_rng(1)
+    time = np.arange(10)
+    df = pd.DataFrame({'var1': rng.random(10), 'var2': rng.random(10), 'var3': rng.random(10)}, index=time)
     df.index.name = 'time'
     return df

Also applies to: 27-33

flixopt/plotting.py (3)

380-395: Type annotation of resolve_colors is too narrow for Matplotlib.

Return type can include RGBA tuples for Matplotlib; narrow typing to str is misleading.

-) -> dict[str, str]:
+) -> dict[str, Any]:
     """Resolve colors parameter to a dict mapping variable names to colors."""

---

386-394: Plotly + dict colors: normalize non-string color values.

If users provide Matplotlib RGBA tuples in a dict while using Plotly, traces may fail. Normalize tuples to 'rgba(r,g,b,a)' strings when engine='plotly'.

-    if isinstance(colors, dict):
-        return colors
+    if isinstance(colors, dict):
+        if engine == 'plotly':
+            def _mpl_tuple_to_rgba(c):
+                if isinstance(c, tuple) and (len(c) == 3 or len(c) == 4):
+                    r, g, b, *rest = c
+                    a = rest[0] if rest else 1.0
+                    # Matplotlib tuples are 0..1 floats; convert to 0..255
+                    if all(isinstance(x, (int, float)) and 0.0 <= x <= 1.0 for x in (r, g, b)):
+                        r, g, b = int(round(r*255)), int(round(g*255)), int(round(b*255))
+                    return f'rgba({r},{g},{b},{a})'
+                return c
+            return {k: _mpl_tuple_to_rgba(v) for k, v in colors.items()}
+        return colors

---

1758-1766: export_figure docstring mismatches behavior (show/dpi).

Doc says show defaults to None and dpi uses CONFIG; code uses show=True and passes dpi directly (None -> Matplotlib defaults). Align docs.

-        show: Whether to display the figure. If None, uses CONFIG.Plotting.default_show (default: None).
+        show: Whether to display the figure (default: True).
@@
-        dpi: DPI (dots per inch) for saving Matplotlib figures. If None, uses CONFIG.Plotting.default_dpi.
+        dpi: DPI (dots per inch) for saving Matplotlib figures. If None, Matplotlib rcParams are used.

flixopt/results.py (2)

1665-1675: Matplotlib legend style is overridden by a second default legend call.

The new ax.legend() replaces the nicely formatted legend from with_matplotlib. Recreate with the same styling so the overlay integrates cleanly.

-            ax.legend()
+            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,
+            )

---

1309-1315: Optional: title suffix may omit input-side selections.

suffix_parts from inputs is overwritten by outputs selection. Concatenate both for clarity.

-        inputs, suffix_parts = _apply_selection_to_data(inputs, select=select, drop=True)
-        outputs, suffix_parts = _apply_selection_to_data(outputs, select=select, drop=True)
+        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
@@
-        suffix = '--' + '-'.join(suffix_parts) if suffix_parts else ''
+        suffix = '--' + '-'.join(sorted(set(suffix_parts))) if suffix_parts else ''

Also applies to: 1346-1348

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 71db28f and 5ca2c51.

📒 Files selected for processing (3)

• flixopt/plotting.py (20 hunks)
• flixopt/results.py (36 hunks)
• tests/test_plotting_api.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (3)

tests/test_plotting_api.py (1)

flixopt/plotting.py (4)

• with_plotly (397-702)
• with_matplotlib (705-852)
• dual_pie_with_plotly (1188-1280)
• dual_pie_with_matplotlib (1283-1387)

flixopt/results.py (1)

flixopt/plotting.py (1)

• with_matplotlib (705-852)

flixopt/plotting.py (2)

flixopt/config.py (2)

• CONFIG (61-322)
• to_dict (289-322)

flixopt/aggregation.py (1)

• plot (144-175)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: test (3.13)
• GitHub Check: test (3.10)
• GitHub Check: test (3.12)
• GitHub Check: test (3.11)

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (5)

flixopt/results.py (4)

1159-1192: Avoid leaking Matplotlib-only kwargs to Plotly; partition kwargs by engine.

Passing plot_kwargs to Plotly will raise “unexpected keyword” errors; likewise, Plotly-only kwargs shouldn’t reach Matplotlib. Partition before dispatch.

Apply:

@@
-        # Extract dpi for export_figure
-        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)  # None handled in export_figure
+        # Partition kwargs by engine to avoid passing unsupported params
+        px_kwargs = {k: v for k, v in plot_kwargs.items() if k != 'plot_kwargs'}
+        mpl_kwargs = {'plot_kwargs': plot_kwargs.get('plot_kwargs', {})}
@@
-            figure_like = plotting.with_plotly(
+            figure_like = plotting.with_plotly(
                 ds,
@@
-                xlabel='Time in h',
-                **plot_kwargs,
+                xlabel='Time in h',
+                **px_kwargs,
             )
@@
-            figure_like = plotting.with_matplotlib(
+            figure_like = plotting.with_matplotlib(
                 ds,
                 colors=colors,
                 mode=mode,
                 title=title,
-                **plot_kwargs,
+                **mpl_kwargs,
             )

---

1291-1376: Matplotlib pie doesn’t accept arbitrary kwargs; only pass supported keys.

dual_pie_with_matplotlib has no **kwargs, so forwarding plot_kwargs can explode (e.g., dpi, hover_template). Extract only supported params; keep Plotly extras separate.

Apply:

@@
-        # Extract dpi for export_figure
-        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)  # None handled in export_figure
+        # Partition kwargs and extract matplotlib-supported keys
+        pie_px_kwargs = {k: v for k, v in plot_kwargs.items() if k != 'plot_kwargs'}
+        hole_kw = plot_kwargs.get('hole')
+        figsize_kw = plot_kwargs.get('figsize')
@@
-            figure_like = plotting.dual_pie_with_plotly(
+            figure_like = plotting.dual_pie_with_plotly(
                 data_left=inputs,
                 data_right=outputs,
                 colors=colors,
                 title=title,
                 text_info=text_info,
                 subtitles=('Inputs', 'Outputs'),
                 legend_title='Flows',
                 lower_percentage_group=lower_percentage_group,
-                **plot_kwargs,
+                **pie_px_kwargs,
             )
@@
-            figure_like = plotting.dual_pie_with_matplotlib(
+            figure_like = plotting.dual_pie_with_matplotlib(
                 data_left=inputs.to_pandas(),
                 data_right=outputs.to_pandas(),
                 colors=colors,
                 title=title,
                 subtitles=('Inputs', 'Outputs'),
                 legend_title='Flows',
                 lower_percentage_group=lower_percentage_group,
-                **plot_kwargs,
+                **({k: v for k, v in {'hole': hole_kw, 'figsize': figsize_kw}.items() if v is not None}),
             )

---

838-859: Don’t forward Matplotlib-style plot_kwargs into Plotly heatmaps.

plot_kwargs may include a plot_kwargs dict intended for Matplotlib; remove it when engine='plotly' to prevent px.imshow argument errors.

Apply:

@@
-        return plot_heatmap(
+        return plot_heatmap(
             data=self.solution[variable_name],
@@
-            **plot_kwargs,
+            **plot_kwargs,
         )

And in the module-level plot_heatmap() below (lines ~2274–2316), pop plot_kwargs for Plotly:

@@
-    # Extract dpi before passing to plotting functions
-    dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+    # Extract dpi before passing to plotting functions
+    dpi = plot_kwargs.pop('dpi', None)
+    # Avoid leaking Matplotlib-only container to Plotly
+    if engine == 'plotly':
+        plot_kwargs.pop('plot_kwargs', None)

---

1578-1664: Partition kwargs in plot_charge_state; fix overlay path for both engines.

Same leakage issue as node balance; ensure Plotly doesn’t receive plot_kwargs and Matplotlib only receives plot_kwargs dict.

Apply:

@@
-        # Extract dpi for export_figure
-        dpi = plot_kwargs.pop('dpi', None)  # None uses CONFIG.Plotting.default_dpi
+        # Extract dpi for export_figure
+        dpi = plot_kwargs.pop('dpi', None)
@@
-        # Extract charge state line color (for overlay customization)
-        overlay_color = plot_kwargs.pop('charge_state_line_color', 'black')
+        # Extract charge state line color (for overlay customization)
+        overlay_color = plot_kwargs.pop('charge_state_line_color', 'black')
+        # Partition kwargs by engine
+        px_kwargs = {k: v for k, v in plot_kwargs.items() if k != 'plot_kwargs'}
+        mpl_kwargs = {'plot_kwargs': plot_kwargs.get('plot_kwargs', {})}
@@
-            figure_like = plotting.with_plotly(
+            figure_like = plotting.with_plotly(
                 ds,
@@
-                xlabel='Time in h',
-                **plot_kwargs,
+                xlabel='Time in h',
+                **px_kwargs,
             )
@@
-            charge_state_fig = plotting.with_plotly(
+            charge_state_fig = plotting.with_plotly(
                 charge_state_ds,
@@
-                xlabel='Time in h',
-                **plot_kwargs,
+                xlabel='Time in h',
+                **px_kwargs,
             )
@@
-            fig, ax = plotting.with_matplotlib(
+            fig, ax = plotting.with_matplotlib(
                 ds,
                 colors=colors,
                 mode=mode,
                 title=title,
-                **plot_kwargs,
+                **mpl_kwargs,
             )

flixopt/plotting.py (1)

1764-1837: Honor “dpi=None uses defaults” and avoid passing None to savefig.

export_figure should treat dpi=None as “use rcParams/savefig.dpi”. Current code forwards None, which may error and contradicts docs.

Apply:

@@
-def export_figure(
+def export_figure(
@@
-    dpi: int = 300,
+    dpi: int | None = 300,
 ) -> go.Figure | tuple[plt.Figure, plt.Axes]:
@@
-        if save:
-            fig.savefig(str(filename), dpi=dpi)
+        if save:
+            eff_dpi = matplotlib.rcParams.get('savefig.dpi', 300) if dpi is None else dpi
+            fig.savefig(str(filename), dpi=eff_dpi)
             plt.close(fig)  # Close figure to free memory

🧹 Nitpick comments (2)

tests/test_plotting_api.py (1)

37-46: Good passthrough check for px kwargs. Consider adding heatmap cases.

Add a small test to assert that vmin/vmax work with Plotly heatmaps (mapped to zmin/zmax) and that Matplotlib-only keys (plot_kwargs) don’t break Plotly calls.

flixopt/plotting.py (1)

1402-1462: Optional: accept list colorscales directly for Plotly heatmaps.

color_continuous_scale can be a list; enable colors list passthrough instead of forcing 'viridis'.

Apply:

@@
-    common_args = {
+    common_args = {
         'img': data,
-        'color_continuous_scale': colors if isinstance(colors, str) else 'viridis',
+        'color_continuous_scale': colors if isinstance(colors, (str, list)) else 'viridis',
         'title': title,
     }

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5ca2c51 and a68dc04.

📒 Files selected for processing (3)

• flixopt/plotting.py (20 hunks)
• flixopt/results.py (38 hunks)
• tests/test_plotting_api.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (3)

tests/test_plotting_api.py (1)

flixopt/plotting.py (4)

• with_plotly (397-702)
• with_matplotlib (705-863)
• dual_pie_with_plotly (1199-1291)
• dual_pie_with_matplotlib (1294-1398)

flixopt/results.py (1)

flixopt/plotting.py (2)

• with_matplotlib (705-863)
• dual_pie_with_plotly (1199-1291)

flixopt/plotting.py (2)

flixopt/config.py (1)

• CONFIG (61-322)

flixopt/aggregation.py (1)

• plot (144-175)

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Map vmin/vmax to zmin/zmax for Plotly heatmaps.

Docs say vmin/vmax are supported “for both engines.” px.imshow expects zmin/zmax; without mapping, users hit TypeError.

Apply:

@@
-    # Merge in additional imshow kwargs
-    common_args.update(imshow_kwargs)
+    # Map matplotlib-style vmin/vmax to Plotly's zmin/zmax if provided
+    if 'vmin' in imshow_kwargs and 'zmin' not in imshow_kwargs:
+        imshow_kwargs['zmin'] = imshow_kwargs.pop('vmin')
+    if 'vmax' in imshow_kwargs and 'zmax' not in imshow_kwargs:
+        imshow_kwargs['zmax'] = imshow_kwargs.pop('vmax')
+    # Merge in additional imshow kwargs
+    common_args.update(imshow_kwargs)

🤖 Prompt for AI Agents

In flixopt/plotting.py around lines 1594 to 1609, px.imshow expects zmin/zmax
but the code passes vmin/vmax which causes a TypeError; before calling px.imshow
(both for the main common_args and the fallback_args) map any vmin -> zmin and
vmax -> zmax (and remove the original vmin/vmax keys) so the arguments passed to
px.imshow use zmin/zmax instead; do this mapping defensively (only if vmin/vmax
present and not None) for both the merged common_args and the fallback_args
prior to calling px.imshow.