fix: move filtering and rounding before early return in case data_only=True

Author: FBumann

flixopt/statistics_accessor.py

@@ -1552,6 +1552,13 @@ def balance(
 
         ds = _apply_selection(ds, select)
 
+        # Round to avoid numerical noise (tiny negative values from solver precision)
+        if round_decimals is not None:
+            ds = ds.round(round_decimals)
+
+        # Filter out variables below threshold
+        ds = _filter_small_variables(ds, threshold)
+
         # Build color kwargs: bus balance → component colors, component balance → carrier colors
         color_by: Literal['component', 'carrier'] = 'component' if is_bus else 'carrier'
         color_kwargs = self._build_color_kwargs(colors, list(ds.data_vars), color_by)
@@ -1560,13 +1567,6 @@ def balance(
         if data_only:
             return PlotResult(data=ds, figure=go.Figure())
 
-        # Round to avoid numerical noise (tiny negative values from solver precision)
-        if round_decimals is not None:
-            ds = ds.round(round_decimals)
-
-        # Filter out variables below threshold
-        ds = _filter_small_variables(ds, threshold)
-
         # Sort for consistent plotting order
         ds = _sort_dataset(ds)
 
@@ -1708,20 +1708,20 @@ def carrier_balance(
 
         ds = _apply_selection(ds, select)
 
-        # Build color kwargs with component colors (flows colored by their parent component)
-        color_kwargs = self._build_color_kwargs(colors, list(ds.data_vars), color_by='component')
-
-        # Early return for data_only mode (skip figure creation for performance)
-        if data_only:
-            return PlotResult(data=ds, figure=go.Figure())
-
         # Round to avoid numerical noise (tiny negative values from solver precision)
         if round_decimals is not None:
             ds = ds.round(round_decimals)
 
         # Filter out variables below threshold
         ds = _filter_small_variables(ds, threshold)
 
+        # Build color kwargs with component colors (flows colored by their parent component)
+        color_kwargs = self._build_color_kwargs(colors, list(ds.data_vars), color_by='component')
+
+        # Early return for data_only mode (skip figure creation for performance)
+        if data_only:
+            return PlotResult(data=ds, figure=go.Figure())
+
         # Sort for consistent plotting order
         ds = _sort_dataset(ds)
 
@@ -1887,13 +1887,13 @@ def flows(
 
         ds = _apply_selection(ds, select)
 
+        # Filter out variables below threshold
+        ds = _filter_small_variables(ds, threshold)
+
         # Early return for data_only mode (skip figure creation for performance)
         if data_only:
             return PlotResult(data=ds, figure=go.Figure())
 
-        # Filter out variables below threshold
-        ds = _filter_small_variables(ds, threshold)
-
         # Sort for consistent plotting order
         ds = _sort_dataset(ds)
 
@@ -1956,13 +1956,13 @@ def sizes(
             valid_labels = [lbl for lbl in ds.data_vars if float(ds[lbl].max()) < max_size]
             ds = ds[valid_labels]
 
+        # Filter out variables below threshold
+        ds = _filter_small_variables(ds, threshold)
+
         # Early return for data_only mode (skip figure creation for performance)
         if data_only:
             return PlotResult(data=ds, figure=go.Figure())
 
-        # Filter out variables below threshold
-        ds = _filter_small_variables(ds, threshold)
-
         if not ds.data_vars:
             fig = go.Figure()
         else:
@@ -2051,13 +2051,13 @@ def duration_curve(
 
         result_ds = ds.fxstats.to_duration_curve(normalize=normalize)
 
+        # Filter out variables below threshold
+        result_ds = _filter_small_variables(result_ds, threshold)
+
         # Early return for data_only mode (skip figure creation for performance)
         if data_only:
             return PlotResult(data=result_ds, figure=go.Figure())
 
-        # Filter out variables below threshold
-        result_ds = _filter_small_variables(result_ds, threshold)
-
         # Sort for consistent plotting order
         result_ds = _sort_dataset(result_ds)
 
@@ -2180,13 +2180,13 @@ def effects(
         else:
             raise ValueError(f"'by' must be one of 'component', 'contributor', 'time', or None, got {by!r}")
 
+        # Filter out variables below threshold
+        ds = _filter_small_variables(ds, threshold)
+
         # Early return for data_only mode (skip figure creation for performance)
         if data_only:
             return PlotResult(data=ds, figure=go.Figure())
 
-        # Filter out variables below threshold
-        ds = _filter_small_variables(ds, threshold)
-
         # Sort for consistent plotting order
         ds = _sort_dataset(ds)
 
@@ -2262,13 +2262,13 @@ def charge_states(
 
         ds = _apply_selection(ds, select)
 
+        # Filter out variables below threshold
+        ds = _filter_small_variables(ds, threshold)
+
         # Early return for data_only mode (skip figure creation for performance)
         if data_only:
             return PlotResult(data=ds, figure=go.Figure())
 
-        # Filter out variables below threshold
-        ds = _filter_small_variables(ds, threshold)
-
         # Sort for consistent plotting order
         ds = _sort_dataset(ds)
 
@@ -2367,10 +2367,6 @@ def storage(
         # Apply selection
         ds = _apply_selection(ds, select)
 
-        # Early return for data_only mode (skip figure creation for performance)
-        if data_only:
-            return PlotResult(data=ds, figure=go.Figure())
-
         # Separate flow data from charge_state
         flow_labels = [lbl for lbl in ds.data_vars if lbl != 'charge_state']
         flow_ds = ds[flow_labels]
@@ -2383,6 +2379,12 @@ def storage(
         # Filter out flow variables below threshold
         flow_ds = _filter_small_variables(flow_ds, threshold)
 
+        # Early return for data_only mode (skip figure creation for performance)
+        if data_only:
+            result_ds = flow_ds.copy()
+            result_ds['charge_state'] = charge_da
+            return PlotResult(data=result_ds, figure=go.Figure())
+
         # Sort for consistent plotting order
         flow_ds = _sort_dataset(flow_ds)
 

tests/test_comparison.py

@@ -359,12 +359,14 @@ def test_balance_returns_plot_result(self, optimized_base, optimized_with_chp):
         assert isinstance(result.data, xr.Dataset)
 
     def test_balance_includes_all_flows(self, optimized_base, optimized_with_chp):
-        """balance() includes flows from both systems."""
+        """balance() includes flows from both systems (with non-zero values)."""
         comp = fx.Comparison([optimized_base, optimized_with_chp])
         result = comp.statistics.plot.balance('Heat', show=False)
 
-        # Should include CHP flow even though it's only in one system
-        assert 'CHP(Q_th_chp)' in result.data
+        # Should include flows that have non-zero values in at least one system
+        # Note: CHP is not used (all zeros) in this test, so it's correctly filtered out
+        # The Boiler flow is present in both systems
+        assert 'Boiler(Q_th)' in result.data
 
     def test_balance_data_has_case_dimension(self, optimized_base, optimized_with_chp):
         """balance() data has 'case' dimension."""