From 0446b48279c6c40dfb62bea599f6713e8dcc03cc Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Tue, 7 Jul 2026 12:40:06 -0700 Subject: [PATCH] adds explicit metadata for bias intervention --- .../block_based_trial_generator.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py index 273ac14..f601aec 100644 --- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py +++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py @@ -28,6 +28,12 @@ class BlockBasedTrialMetadata(BaseModel): """Metadata for block based trial. These fields will NOT be used by the task engine.""" is_autowater: bool = Field(default=False, description="Flag indicating if autowater is given for trial.") + is_bias_water_intervention: bool = Field( + default=False, description="Flag indicating if bias water intervention is given for trial." + ) + is_bias_stage_intervention: bool = Field( + default=False, description="Flag indicating if bias stage intervention is given for trial." + ) class AutoWaterParameters(BaseModel): @@ -203,7 +209,7 @@ def next(self) -> Trial | None: # determine bias correction. Overrides autowater lickspout_offset_delta = 0 - if self.bias_intervention.are_antibias_conditions_met(self.bias): + if is_bias_intervention := self.bias_intervention.are_antibias_conditions_met(self.bias): is_auto_reward_right, lickspout_offset_delta = self.bias_intervention.determine_antibias_intervention( self.bias ) @@ -224,7 +230,11 @@ def next(self) -> Trial | None: metadata=Metadata( p_reward_left=self.block.p_left_reward, p_reward_right=self.block.p_right_reward, - extra=BlockBasedTrialMetadata(is_autowater=is_autowater), + extra=BlockBasedTrialMetadata( + is_autowater=is_autowater and not is_bias_intervention, + is_bias_water_intervention=is_bias_intervention and is_auto_reward_right is not None, + is_bias_stage_intervention=is_bias_intervention and lickspout_offset_delta != 0, + ), ), )