diff --git a/src/aind_behavior_dynamic_foraging/data_contract/utils.py b/src/aind_behavior_dynamic_foraging/data_contract/utils.py index 8a95b027..83687894 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/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