Fixes calculate consumed water

src/aind_behavior_dynamic_foraging/data_contract/utils.py

@@ -1,42 +1,90 @@
 import os
-from typing import Optional
+from pathlib import Path
 
-from aind_behavior_dynamic_foraging.data_contract import dataset
-from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic
+import numpy as np
+import pandas as pd
 
+from aind_behavior_dynamic_foraging.data_contract import dataset as df_dataset
+from aind_behavior_dynamic_foraging.rig import AindDynamicForagingRig
 
-def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]:
-    """Calculate the total volume of water consumed during a session.
+
+def _calculate_side_volume_ml(
+    set_open_time_ms: pd.Series,
+    delivery_times: pd.DataFrame,
+    slope_g_per_s: float,
+    offset_g: float,
+) -> float:
+    """Estimate delivered volume for one side from set open times and valve-open events.
 
     Args:
-        session_path (os.PathLike): Path to the session directory.
+        set_open_time_ms (pd.Series): Time-indexed set open-time values in milliseconds.
+        delivery_times (pd.DataFrame): Event rows where the side valve was commanded open.
+        slope_g_per_s (float): Calibration slope converting open duration (s) to delivered (g).
+        offset_g (float): Calibration offset in grams applied per delivered event.
 
     Returns:
-        Optional[float]: Total volume of water consumed in milliliters, or None if unavailable.
+        float: Total delivered volume in mL for the side.
     """
 
-    trial_outcomes = dataset(session_path)["Behavior"]["SoftwareEvents"]["TrialOutcome"].load().data["data"]
-    is_right_choice = [to["is_right_choice"] for to in trial_outcomes]
-    is_rewarded = [to["is_rewarded"] for to in trial_outcomes]
-
-    task_logic_data = dataset(session_path)["Behavior"]["InputSchemas"]["TaskLogic"].load().data
-    task_logic = AindDynamicForagingTaskLogic.model_validate(task_logic_data)
-    right_reward_size = task_logic.task_parameters.reward_size.right_value_volume
-    left_reward_size = task_logic.task_parameters.reward_size.left_value_volume
-
-    total = 0
-    for choice, rewarded in zip(is_right_choice, is_rewarded):
-        if rewarded:
-            if choice is True:
-                total += right_reward_size * 1e-3
-            if choice is False:
-                total += left_reward_size * 1e-3
-
-    is_right_manual_water = dataset(session_path)["Behavior"]["SoftwareEvents"]["GiveManualWaterRight"].load()
-    if is_right_manual_water.has_data:
-        for is_right in is_right_manual_water.data["data"]:
-            if is_right:
-                total += right_reward_size * 1e-3
-            else:
-                total += left_reward_size * 1e-3
-    return total
+    delivery_times = delivery_times.reset_index(names="Time")[["Time"]].sort_values("Time")
+    if delivery_times.empty:
+        return 0.0
+
+    # normalize setpoints to numeric values and reshape into a Time-keyed frame.
+    setpoints = (
+        pd.to_numeric(set_open_time_ms, errors="coerce")
+        .dropna()
+        .sort_index()
+        .rename("set_open_time_ms")
+        .to_frame()
+        .reset_index(names="Time")
+    )
+    if setpoints.empty:
+        return 0.0
+
+    # Each valve-open event uses the most recent set open-time configured at or before that event.
+    matched = pd.merge_asof(delivery_times, setpoints, on="Time", direction="backward")
+    open_times_s = (matched["set_open_time_ms"].dropna() / 1000.0).to_numpy()
+    if len(open_times_s) == 0:
+        return 0.0
+
+    delivered_g = np.round((slope_g_per_s * open_times_s) + offset_g, 4)
+    return float(delivered_g.sum())
+
+
+def calculate_consumed_water(session_path: str | os.PathLike[str]) -> float:
+    """Calculate total delivered water volume across left and right valves for a session.
+
+    Args:
+        session_path (str | os.PathLike[str]): Path to the session directory.
+
+    Returns:
+        float: Total water delivered in mL for the session.
+    """
+
+    dataset = df_dataset(Path(session_path))["Behavior"]
+
+    rig = AindDynamicForagingRig.model_validate(dataset["InputSchemas"]["Rig"].data)
+    left_calibration = rig.calibration.water_valve_left
+    right_calibration = rig.calibration.water_valve_right
+
+    left_set_open_time_ms = dataset["HarpBehavior"]["PulseSupplyPort0"].load().data
+    right_set_open_time_ms = dataset["HarpBehavior"]["PulseSupplyPort1"].load().data
+    output_set_stream = dataset["HarpBehavior"]["OutputSet"].load().data
+    writes = output_set_stream[output_set_stream["MessageType"] == "WRITE"]
+
+    left_ml = _calculate_side_volume_ml(
+        set_open_time_ms=left_set_open_time_ms["PulseSupplyPort0"],
+        delivery_times=writes[writes["SupplyPort0"].fillna(False).astype(bool)],
+        slope_g_per_s=float(left_calibration.slope),
+        offset_g=float(left_calibration.offset),
+    )
+
+    right_ml = _calculate_side_volume_ml(
+        set_open_time_ms=right_set_open_time_ms["PulseSupplyPort1"],
+        delivery_times=writes[writes["SupplyPort1"].fillna(False).astype(bool)],
+        slope_g_per_s=float(right_calibration.slope),
+        offset_g=float(right_calibration.offset),
+    )
+
+    return left_ml + right_ml