Feat metadata mapping

scripts/walk_through_session.py

@@ -2,7 +2,9 @@
 import os
 
 from aind_behavior_dynamic_foraging.data_contract import dataset as df_foraging_dataset
-from aind_behavior_dynamic_foraging.task_logic.trial_generators.warmup_trial_generator import WarmupTrialGeneratorSpec
+from aind_behavior_dynamic_foraging.task_logic.trial_generators import (
+    CoupledTrialGeneratorSpec,
+)
 from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome
 
 logging.basicConfig(
@@ -17,10 +19,10 @@ def walk_through_session(data_directory: os.PathLike):
     software_events.load_all()
 
     trial_outcomes = software_events["TrialOutcome"].data["data"].iloc
-    warmup_trial_generator = WarmupTrialGeneratorSpec().create_generator()
+    trial_generator = CoupledTrialGeneratorSpec().create_generator()
     for i, outcome in enumerate(trial_outcomes):
-        warmup_trial_generator.update(TrialOutcome.model_validate(outcome))
-        trial = warmup_trial_generator.next()
+        trial_generator.update(TrialOutcome.model_validate(outcome))
+        trial = trial_generator.next()
 
         if not trial:
             print(f"Session finished at trial {i}")

src/Extensions/bonsai.py

@@ -1,8 +1,12 @@
+import logging
+import sys
 from typing import TYPE_CHECKING
 
 from pydantic import TypeAdapter
 
-from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+
+from aind_behavior_dynamic_foraging.task_logic import TrialGeneratorSpec  # noqa
 
 if TYPE_CHECKING:
     from aind_behavior_dynamic_foraging.task_logic.trial_generators._base import ITrialGenerator

src/aind_behavior_dynamic_foraging/data_contract/_dataset.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from aind_behavior_curriculum import TrainerState
+from aind_behavior_curriculum import Metrics, TrainerState
 from aind_behavior_services.session import Session
 from contraqctor.contract import Dataset, DataStreamCollection
 from contraqctor.contract.camera import Camera
@@ -61,10 +61,18 @@ def make_dataset(
                 data_streams=[
                     Json(
                         name="PreviousMetrics",
-                        reader_params=Json.make_params(
+                        reader_params=PydanticModel.make_params(
+                            model=Metrics,
                             path=root_path / "behavior/previous_metrics.json",
                         ),
                     ),
+                    PydanticModel(
+                        name="Metrics",
+                        reader_params=PydanticModel.make_params(
+                            model=Metrics,
+                            path=root_path / "behavior/metrics.json",
+                        ),
+                    ),
                     PydanticModel(
                         name="TrainerState",
                         reader_params=PydanticModel.make_params(

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

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py

@@ -163,13 +163,13 @@ def next(self) -> Trial | None:
         if self.spec.is_baiting:
             random_numbers = np.random.random(2)
 
-            is_left_baited = self.block.p_left_reward > random_numbers[0] or self.is_left_baited
-            logger.debug(f"Left baited: {is_left_baited}")
-            p_reward_left = 1 if is_left_baited else p_reward_left
+            self.is_left_baited = self.block.p_left_reward > random_numbers[0] or self.is_left_baited
+            logger.debug(f"Left baited: {self.is_left_baited}")
+            p_reward_left = 1 if self.is_left_baited else p_reward_left
 
-            is_right_baited = self.block.p_right_reward > random_numbers[1] or self.is_right_baited
-            logger.debug(f"Right baited: {is_left_baited}")
-            p_reward_right = 1 if is_right_baited else p_reward_right
+            self.is_right_baited = self.block.p_right_reward > random_numbers[1] or self.is_right_baited
+            logger.debug(f"Right baited: {self.is_right_baited}")
+            p_reward_right = 1 if self.is_right_baited else p_reward_right
 
         return Trial(
             p_reward_left=p_reward_left,
@@ -195,7 +195,7 @@ def _generate_next_block(
         reward_pairs: list[list[float, float]],
         base_reward_sum: float,
         block_len: Union[UniformDistribution, ExponentialDistribution],
-        current_block: Optional[None] = None,
+        current_block: Optional[Block] = None,
     ) -> Block:
         """Generates the next block, avoiding repeating the current block's side bias.
 

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generator.py

@@ -134,6 +134,13 @@ def _are_end_conditions_met(self) -> bool:
             logger.debug("Maximum trial count exceeded.")
             return True
 
+        logger.debug(
+            "Trial generation end conditions are not met: "
+            f"total trials={len(self.is_right_choice_history)}, "
+            f"time elapsed={time_elapsed},"
+            f"ignored trial={choice_history[-win:].count(None)},"
+        )
+
         return False
 
     def update(self, outcome: TrialOutcome | str) -> None:

workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/coupled_baiting/stages.py

@@ -66,8 +66,8 @@ def make_s_stage_1_warmup():
                         CoupledTrialGeneratorSpec(
                             trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                                 max_trial=1000,
-                                max_time=75,
-                                min_time=30,
+                                max_time=4500,
+                                min_time=1800,
                                 ignore_win=20000,
                                 ignore_ratio_threshold=1,
                             ),
@@ -113,8 +113,8 @@ def make_s_stage_1():
                 trial_generator=CoupledTrialGeneratorSpec(
                     trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                         max_trial=1000,
-                        max_time=75,
-                        min_time=30,
+                        max_time=4500,
+                        min_time=1800,
                         ignore_win=20000,
                         ignore_ratio_threshold=1,
                     ),
@@ -158,8 +158,8 @@ def make_s_stage_2():
                 trial_generator=CoupledTrialGeneratorSpec(
                     trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                         max_trial=1000,
-                        max_time=75,
-                        min_time=30,
+                        max_time=4500,
+                        min_time=1800,
                         ignore_win=30,
                         ignore_ratio_threshold=0.83,
                     ),
@@ -203,8 +203,8 @@ def make_s_stage_3():
                 trial_generator=CoupledTrialGeneratorSpec(
                     trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                         max_trial=1000,
-                        max_time=75,
-                        min_time=30,
+                        max_time=4500,
+                        min_time=1800,
                         ignore_win=30,
                         ignore_ratio_threshold=0.83,
                     ),
@@ -248,8 +248,8 @@ def make_s_stage_final():
                 trial_generator=CoupledTrialGeneratorSpec(
                     trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                         max_trial=1000,
-                        max_time=75,
-                        min_time=30,
+                        max_time=4500,
+                        min_time=1800,
                         ignore_win=30,
                         ignore_ratio_threshold=0.83,
                     ),
@@ -289,8 +289,8 @@ def make_s_stage_graduated():
                 trial_generator=CoupledTrialGeneratorSpec(
                     trial_generation_end_parameters=CoupledTrialGenerationEndConditions(
                         max_trial=1000,
-                        max_time=75,
-                        min_time=30,
+                        max_time=4500,
+                        min_time=1800,
                         ignore_win=30,
                         ignore_ratio_threshold=0.83,
                     ),

workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py

@@ -71,7 +71,7 @@ def metrics_from_dataset(
     logger.debug(f"Calculated foraging efficiency as {foraging_efficiency}")
 
     try:
-        prev_metrics = DynamicForagingMetrics(**dataset["Behavior"]["PreviousMetrics"].data)
+        prev_metrics = DynamicForagingMetrics.model_validate(dataset["Behavior"]["PreviousMetrics"].data)
         prev_stage = prev_metrics.stage_name
     except FileNotFoundError:
         logger.info("No previous metrics found.")
@@ -138,7 +138,7 @@ def compute_foraging_efficiency(
 
     if not is_baiting:
         logger.debug("Calculated non baiting foraging efficiency.")
-        optimal_rewards_per_session = np.nanmean(np.max([p_right_reward], axis=0)) * len(p_left_reward)
+        optimal_rewards_per_session = np.nanmean(np.max([p_right_reward, p_left_reward], axis=0)) * len(p_left_reward)
     else:
         logger.debug("Calculated baiting foraging efficiency.")
         p_max = np.maximum(p_left_reward, p_right_reward)