From 5a14ee53faebf6e9f2801bf238e41b4ff03ecd94 Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Wed, 29 Jul 2026 10:49:18 -0700 Subject: [PATCH 1/3] update utils function to calculate water delivery based on open and delivery times --- .../data_contract/utils.py | 97 +++++++++++++------ 1 file changed, 65 insertions(+), 32 deletions(-) diff --git a/src/aind_behavior_dynamic_foraging/data_contract/utils.py b/src/aind_behavior_dynamic_foraging/data_contract/utils.py index 8a95b027..bf2ccb7a 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -1,42 +1,75 @@ 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 pandas as pd +import numpy as np +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 (mL) for one side from pulse durations and calibration.""" + + delivery_times = delivery_times.reset_index(names="Time") + if delivery_times.empty: + return 0.0 + + pulse_series_ms = pd.to_numeric(set_open_time_ms, errors="coerce").dropna().sort_index() + if pulse_series_ms.empty: + return 0.0 + + setpoints = pulse_series_ms.rename("open_time_ms").to_frame().reset_index(names="Time") + delivery_times = delivery_times[["Time"]].sort_values("Time") + + matched = pd.merge_asof(delivery_times, setpoints, on="Time", direction="backward") + open_times_s = (matched["open_time_ms"].dropna() / 1000.0).to_numpy() + if len(open_times_s) == 0: + return 0.0 + print(open_times_s, slope_g_per_s, offset_g) + delivered_g = np.round((slope_g_per_s * open_times_s) + offset_g, 4) + print(delivered_g) + # For water, 1 g is approximately 1 mL. + return float(delivered_g.sum()) + + +def calculate_consumed_water(session_path: str | os.PathLike[str]) -> float: + """Calculate the delivered water volume for left/right valves and total session consumption. Args: - session_path (os.PathLike): Path to the session directory. + session_path (str | os.PathLike[str]): Path to the session directory. Returns: - Optional[float]: Total volume of water consumed in milliliters, or None if unavailable. + float: Total water delivered in mL for the session. """ - 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 + 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 From 859d7698b194b1dee01433fa45f2aa77d8adbe9d Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Wed, 29 Jul 2026 11:09:41 -0700 Subject: [PATCH 2/3] removes prints --- src/aind_behavior_dynamic_foraging/data_contract/utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aind_behavior_dynamic_foraging/data_contract/utils.py b/src/aind_behavior_dynamic_foraging/data_contract/utils.py index bf2ccb7a..6c78939d 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -30,9 +30,7 @@ def _calculate_side_volume_ml( open_times_s = (matched["open_time_ms"].dropna() / 1000.0).to_numpy() if len(open_times_s) == 0: return 0.0 - print(open_times_s, slope_g_per_s, offset_g) delivered_g = np.round((slope_g_per_s * open_times_s) + offset_g, 4) - print(delivered_g) # For water, 1 g is approximately 1 mL. return float(delivered_g.sum()) From 06359be61ef2ac686a4d22510681aab3cb8a1cad Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Fri, 31 Jul 2026 10:49:35 -0700 Subject: [PATCH 3/3] lints and comments --- .../data_contract/utils.py | 39 +++++++++++++------ 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/src/aind_behavior_dynamic_foraging/data_contract/utils.py b/src/aind_behavior_dynamic_foraging/data_contract/utils.py index 6c78939d..83687894 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -1,8 +1,9 @@ import os from pathlib import Path -import pandas as pd 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 @@ -13,30 +14,46 @@ def _calculate_side_volume_ml( slope_g_per_s: float, offset_g: float, ) -> float: - """Estimate delivered volume (mL) for one side from pulse durations and calibration.""" + """Estimate delivered volume for one side from set open times and valve-open events. - delivery_times = delivery_times.reset_index(names="Time") + Args: + 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: + float: Total delivered volume in mL for the side. + """ + + delivery_times = delivery_times.reset_index(names="Time")[["Time"]].sort_values("Time") if delivery_times.empty: return 0.0 - pulse_series_ms = pd.to_numeric(set_open_time_ms, errors="coerce").dropna().sort_index() - if pulse_series_ms.empty: + # 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 - setpoints = pulse_series_ms.rename("open_time_ms").to_frame().reset_index(names="Time") - delivery_times = delivery_times[["Time"]].sort_values("Time") - + # 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["open_time_ms"].dropna() / 1000.0).to_numpy() + 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) - # For water, 1 g is approximately 1 mL. return float(delivered_g.sum()) def calculate_consumed_water(session_path: str | os.PathLike[str]) -> float: - """Calculate the delivered water volume for left/right valves and total session consumption. + """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.