fix: missing production and consumption column

Signed-off-by: Mohammad Tayyab <[email protected]>

Author: Mohammad-Tayyab-Frequenz

src/frequenz/lib/notebooks/reporting/utils/helpers.py

@@ -23,6 +23,8 @@
 import pandas as pd
 
 from frequenz.lib.notebooks.reporting.metrics.reporting_metrics import (
+    asset_production,
+    consumption,
     grid_feed_in,
     production_excess,
     production_excess_in_bat,
@@ -31,7 +33,9 @@
 )
 
 
-def _get_numeric_series(df: pd.DataFrame, col: str | None) -> pd.Series:
+def _get_numeric_series(
+    df: pd.DataFrame, col: str | None, *, clip_non_negative: bool = False
+) -> pd.Series:
     """Safely extract a numeric Series or return zeros if missing.
 
     Ensures consistent numeric handling even when the requested column
@@ -41,16 +45,25 @@ def _get_numeric_series(df: pd.DataFrame, col: str | None) -> pd.Series:
     Args:
         df: Input DataFrame from which to extract the column.
         col: Column name to retrieve. If None or missing, zeros are returned.
+        clip_non_negative: Clip the resulting series at zero (lower bound).
 
     Returns:
-        A float64 Series with non-negative values, matching the input index.
+        A float64 Series aligned to the input index, optionally clipped to be
+        non-negative.
     """
     if col is None:
-        return pd.Series(0.0, index=df.index, dtype="float64")
-    return df.reindex(columns=[col], fill_value=0)[col].astype("float64").clip(lower=0)
+        series = pd.Series(0.0, index=df.index, dtype="float64")
+    else:
+        series = df.reindex(columns=[col], fill_value=0)[col].astype("float64")
+
+    if clip_non_negative:
+        series = series.clip(lower=0)
+    return series
 
 
-def _sum_cols(df: pd.DataFrame, cols: list[str] | None) -> pd.Series:
+def _sum_cols(
+    df: pd.DataFrame, cols: list[str] | None, *, clip_non_negative: bool = False
+) -> pd.Series:
     """Safely sum multiple numeric columns into a single Series.
 
     Ensures robust aggregation even when some columns are missing or None.
@@ -59,24 +72,30 @@ def _sum_cols(df: pd.DataFrame, cols: list[str] | None) -> pd.Series:
     Args:
         df: Input DataFrame containing the columns to be summed.
         cols: list of column names to sum. If empty, returns a zero-filled Series.
+        clip_non_negative: Clip individual series at zero before summing.
 
     Returns:
         A float64 Series representing the elementwise sum of all specified columns.
-        Missing or invalid columns are treated as zeros.
+        Missing or invalid columns are treated as zeros. The result is
+        non-negative when ``clip_non_negative`` is True.
     """
     if not cols:
         return pd.Series(0.0, index=df.index, dtype="float64")
 
-    # Safely extract each column as a numeric, non-negative Series, then sum row-wise
-    series_list = [_get_numeric_series(df, c) for c in cols]
+    # Safely extract each column as a numeric Series
+    # (clipped to non-negative if clip_non_negative=True), then sum row-wise
+    series_list = [
+        _get_numeric_series(df, c, clip_non_negative=clip_non_negative) for c in cols
+    ]
     return pd.concat(series_list, axis=1).sum(axis=1).astype("float64")
 
 
-# pylint: disable=too-many-arguments, too-many-locals
+# pylint: disable=too-many-arguments, too-many-locals, too-many-positional-arguments
 def add_energy_flows(
     df: pd.DataFrame,
     production_cols: list[str] | None = None,
     consumption_cols: list[str] | None = None,
+    grid_cols: list[str] | None = None,
     battery_charge_col: str | None = None,
     production_is_positive: bool = False,
 ) -> pd.DataFrame:
@@ -91,6 +110,7 @@ def add_energy_flows(
             battery charge data.
         production_cols: list of column names representing production sources.
         consumption_cols: list of column names representing consumption sources.
+        grid_cols: list of column names representing grid import/export.
         battery_charge_col: optional column name for battery charging power. If None,
             battery-related flows are set to zero.
         production_is_positive: Whether production values are already positive.
@@ -106,15 +126,71 @@ def add_energy_flows(
     """
     df_flows = df.copy()
 
-    # Total production and consumption (returns pandas series with 0.0 for missing cols)
-    df_flows["production_total"] = _sum_cols(df_flows, production_cols)
-    df_flows["consumption_total"] = _sum_cols(df_flows, consumption_cols)
-
-    # Surplus vs. consumption
+    # Normalize production, grid and consumption columns by removing None entries
+    resolved_production_cols = [
+        col for col in (production_cols or []) if col is not None
+    ]
+    resolved_consumption_cols = [
+        col for col in (consumption_cols or []) if col is not None
+    ]
+    resolved_grid_cols = [col for col in (grid_cols or []) if col is not None]
+
+    # Compute total asset production
+    asset_production_cols: list[str] = []
+    for col in resolved_production_cols:
+        series = _get_numeric_series(
+            df_flows,
+            col,
+        )
+        if col not in df_flows:
+            df_flows[col] = series
+        asset_series = asset_production(
+            series,
+            production_is_positive=production_is_positive,
+        )
+        asset_col_name = f"{col}_asset_production"
+        df_flows[asset_col_name] = asset_series
+        asset_production_cols.append(asset_col_name)
+
+    df_flows["production_total"] = _sum_cols(df_flows, asset_production_cols)
+
+    # Compute total consumption
+    if not resolved_consumption_cols:
+        # Use existing 'consumption' column if present
+        if "consumption" in df_flows.columns:
+            resolved_consumption_cols = ["consumption"]
+
+        # Infer from grid/production if grid columns exist
+        elif resolved_grid_cols:
+            for col in resolved_grid_cols:
+                if col not in df_flows:
+                    df_flows[col] = _get_numeric_series(
+                        df_flows,
+                        col,
+                        clip_non_negative=True,
+                    )
+            inferred_consumption = consumption(
+                df_flows, asset_production_cols, resolved_grid_cols
+            ).astype("float64")
+            inferred_name = str(inferred_consumption.name or "consumption")
+            df_flows[inferred_name] = inferred_consumption
+            resolved_consumption_cols = [inferred_name]
+
+        # Fallback — create a zero-filled consumption column
+        else:
+            inferred_consumption = pd.Series(
+                0.0, index=df_flows.index, dtype="float64", name="consumption"
+            )
+            df_flows[inferred_consumption.name] = inferred_consumption
+            resolved_consumption_cols = [str(inferred_consumption.name)]
+
+    df_flows["consumption_total"] = _sum_cols(df_flows, resolved_consumption_cols)
+
+    # Surplus vs. consumption (production is already positive because of the above cleaning)
     df_flows["production_excess"] = production_excess(
         df_flows["production_total"],
         df_flows["consumption_total"],
-        production_is_positive=production_is_positive,
+        production_is_positive=True,
     )
 
     # Battery charging power (optional)
@@ -123,15 +199,15 @@ def add_energy_flows(
         df_flows["production_total"],
         df_flows["consumption_total"],
         bat_in,
-        production_is_positive=production_is_positive,
+        production_is_positive=True,
     )
 
     # Split excess into battery vs. grid
     df_flows["grid_feed_in"] = grid_feed_in(
         df_flows["production_total"],
         df_flows["consumption_total"],
         bat_in,
-        production_is_positive=production_is_positive,
+        production_is_positive=True,
     )
 
     # If no production columns exist, set self-consumption metrics to zero
@@ -141,12 +217,12 @@ def add_energy_flows(
         df_flows["production_self_use"] = production_self_consumption(
             df_flows["production_total"],
             df_flows["consumption_total"],
-            production_is_positive=production_is_positive,
+            production_is_positive=True,
         )
         df_flows["production_self_share"] = production_self_share(
             df_flows["production_total"],
             df_flows["consumption_total"],
-            production_is_positive=production_is_positive,
+            production_is_positive=True,
         )
     else:
         df_flows["production_self_use"] = 0.0