From 0600767ecbd7be8e9aced6460262afe01a843a4f Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Mon, 11 May 2026 15:17:34 -0700 Subject: [PATCH 1/4] adds calculate_consumed_water as a data_contract util --- src/Extensions/bonsai.py | 1 + .../data_contract/utils.py | 34 +++++++++++++++++++ 2 files changed, 35 insertions(+) create mode 100644 src/aind_behavior_dynamic_foraging/data_contract/utils.py diff --git a/src/Extensions/bonsai.py b/src/Extensions/bonsai.py index 5fab0e0f..10eeb0a8 100644 --- a/src/Extensions/bonsai.py +++ b/src/Extensions/bonsai.py @@ -3,6 +3,7 @@ from pydantic import TypeAdapter from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec +from aind_behavior_dynamic_foraging.task_logic.utils import calculate_bias as calculate_bias if TYPE_CHECKING: from aind_behavior_dynamic_foraging.task_logic.trial_generators._base import ITrialGenerator 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..ab7819d8 --- /dev/null +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -0,0 +1,34 @@ +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] + + 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 + return total From 10f104cdcae8dc4497dc9b14a43e6768a5dc4368 Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Mon, 11 May 2026 15:18:33 -0700 Subject: [PATCH 2/4] removes import statement --- src/Extensions/bonsai.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/Extensions/bonsai.py b/src/Extensions/bonsai.py index 10eeb0a8..5fab0e0f 100644 --- a/src/Extensions/bonsai.py +++ b/src/Extensions/bonsai.py @@ -3,7 +3,6 @@ from pydantic import TypeAdapter from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec -from aind_behavior_dynamic_foraging.task_logic.utils import calculate_bias as calculate_bias if TYPE_CHECKING: from aind_behavior_dynamic_foraging.task_logic.trial_generators._base import ITrialGenerator From 58ad35cb5ea18a94be5391b728c1036ee41d3dda Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Fri, 29 May 2026 12:43:56 -0700 Subject: [PATCH 3/4] adds manual water subject and software event --- src/Extensions/OperationControl.bonsai | 42 +++++++++++++------ src/Extensions/Visualizers.bonsai | 4 +- .../data_contract/_dataset.py | 7 ++++ .../data_contract/utils.py | 3 ++ 4 files changed, 42 insertions(+), 14 deletions(-) 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 index ab7819d8..d801b454 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -19,6 +19,8 @@ def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]: is_right_choice = [to["is_right_choice"] for to in trial_outcomes] is_rewarded = [to["is_rewarded"] for to in trial_outcomes] + 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 @@ -31,4 +33,5 @@ def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]: total += right_reward_size * 1e-3 if choice is False: total += left_reward_size * 1e-3 + return total From 090668ab0995f4688a50b2573addfef1c3af6eb1 Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Fri, 29 May 2026 12:50:23 -0700 Subject: [PATCH 4/4] updates calculate_consumed_water function to account for manual water --- .../data_contract/utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 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 d801b454..9c5fc942 100644 --- a/src/aind_behavior_dynamic_foraging/data_contract/utils.py +++ b/src/aind_behavior_dynamic_foraging/data_contract/utils.py @@ -19,7 +19,9 @@ def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]: is_right_choice = [to["is_right_choice"] for to in trial_outcomes] is_rewarded = [to["is_rewarded"] for to in trial_outcomes] - manual_water = dataset(session_path)["Behavior"]["SoftwareEvents"]["GiveManualWaterRight"].load().data["data"] + 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) @@ -33,5 +35,11 @@ def calculate_consumed_water(session_path: os.PathLike) -> Optional[float]: 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