diff --git a/src/Extensions/OperationControl.bonsai b/src/Extensions/OperationControl.bonsai
index 9735b8ab..c30ba2ff 100644
--- a/src/Extensions/OperationControl.bonsai
+++ b/src/Extensions/OperationControl.bonsai
@@ -343,6 +343,9 @@
GiveRewardRight
+
+ GiveManualWaterRight
+
SetRewardAmount
@@ -458,6 +461,18 @@
+
+
+
+
+ GiveManualWaterRight
+
+
+ GiveManualWaterRight
+
+
+ GiveRewardRight
+
GiveRewardRight
@@ -607,27 +622,30 @@
-
-
-
-
+
+
-
+
-
-
+
+
+
-
-
+
-
+
-
-
+
+
+
+
+
+
+
diff --git a/src/Extensions/Visualizers.bonsai b/src/Extensions/Visualizers.bonsai
index 9a7b6985..11705ea6 100644
--- a/src/Extensions/Visualizers.bonsai
+++ b/src/Extensions/Visualizers.bonsai
@@ -416,7 +416,7 @@
- GiveRewardRight
+ GiveManualWaterRight
@@ -447,7 +447,7 @@
- GiveRewardRight
+ GiveManualWaterRight
diff --git a/src/aind_behavior_dynamic_foraging/data_contract/_dataset.py b/src/aind_behavior_dynamic_foraging/data_contract/_dataset.py
index 3f2e2381..14339d53 100644
--- a/src/aind_behavior_dynamic_foraging/data_contract/_dataset.py
+++ b/src/aind_behavior_dynamic_foraging/data_contract/_dataset.py
@@ -194,6 +194,13 @@ def make_dataset(
name="SoftwareEvents",
description="Software events generated by the workflow. The timestamps of these events are low precision and should not be used to align to physiology data.",
data_streams=[
+ SoftwareEvents(
+ name="GiveManualWaterRight",
+ description="An event emitted when manual water is given through visualizer.",
+ reader_params=SoftwareEvents.make_params(
+ root_path / "behavior/SoftwareEvents/GiveManualWaterRigt.json"
+ ),
+ ),
SoftwareEvents(
name="TrialGeneratorSpec",
description="An event emitted with the specification for the trial generator.",
diff --git a/src/aind_behavior_dynamic_foraging/data_contract/utils.py b/src/aind_behavior_dynamic_foraging/data_contract/utils.py
new file mode 100644
index 00000000..9c5fc942
--- /dev/null
+++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py
@@ -0,0 +1,45 @@
+import os
+from typing import Optional
+
+from aind_behavior_dynamic_foraging.data_contract import dataset
+from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic
+
+
+def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]:
+ """Calculate the total volume of water consumed during a session.
+
+ Args:
+ session_path (os.PathLike): Path to the session directory.
+
+ Returns:
+ Optional[float]: Total volume of water consumed in milliliters, or None if unavailable.
+ """
+
+ 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]
+
+ is_right_manual_water = (
+ dataset(session_path)["Behavior"]["SoftwareEvents"]["GiveManualWaterRight"].load().data["data"]
+ )
+
+ 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
+
+ for is_right in is_right_manual_water:
+ if is_right:
+ total += right_reward_size * 1e-3
+ else:
+ total += left_reward_size * 1e-3
+
+ return total