Adds extra metadata to trial generator

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 from abc import ABC, abstractmethod
 from typing import Literal, Optional
@@ -16,7 +17,7 @@
     BiasIntervention,
     BiasInterventionParameters,
 )
-from aind_behavior_dynamic_foraging.task_logic.utils.calculate_bias import calculate_bias
+from aind_behavior_dynamic_foraging.task_logic.utils import calculate_bias, calculate_foraging_efficiency
 
 from ..trial_models import Metadata, RewardSize, Trial, TrialMetrics
 from ._base import BaseTrialGeneratorSpecModel, ITrialGenerator, TrialOutcome
@@ -27,6 +28,20 @@
 class BlockBasedTrialMetadata(BaseModel):
     """Metadata for block based trial. These fields will NOT be used by the task engine."""
 
+    time_elapsed: Optional[float] = Field(
+        default=None, description="Time elapsed in session at start of trial (in minutes)."
+    )
+    time_remaining: Optional[float] = Field(
+        default=None, description="Time remaining in session at start of trial (in minutes)."
+    )
+    current_trial: Optional[int] = Field(default=None, description="Current trial number in session.")
+    responses: Optional[int] = Field(default=None, description="Number of responses made in session.")
+    ignored: Optional[int] = Field(default=None, description="Number of ignored trials in session.")
+    earned_water: Optional[float] = Field(default=None, description="Total water earned in session (in mL).")
+    total_water: Optional[float] = Field(default=None, description="Total water delivered in session (in mL).")
+    foraging_efficiency: Optional[float] = Field(
+        default=None, description="Foraging efficiency in session (earned water / total water)."
+    )
     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."
@@ -138,6 +153,7 @@ def __init__(self, spec: BlockBasedTrialGeneratorSpec) -> None:
         """
 
         self.spec = spec
+        self.start_time = datetime.datetime.now()
         self.outcome_history: list[TrialOutcome] = []
         self.is_right_choice_history: list[bool | None] = []
         self.reward_history: list[bool] = []
@@ -228,7 +244,7 @@ def next(self) -> Trial | None:
                 % (is_auto_reward_right, lickspout_offset_delta)
             )
 
-        return Trial(
+        trial = Trial(
             p_reward_left=1 if (self.is_left_baited or is_auto_reward_right is False) else self.block.p_left_reward,
             p_reward_right=1 if (self.is_right_baited or is_auto_reward_right) else self.block.p_right_reward,
             reward_consumption_duration=self.spec.reward_consumption_duration,
@@ -243,13 +259,53 @@ 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 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,
-                ),
             ),
         )
+        extra_metadata = 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,
+        )
+        trial.metadata.extra = self._add_extra_metadata(extra_metadata)
+        return trial
+
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
+        """Adds extra metadata.
+
+        Args:
+            extra_metadata: The extra metadata to add to.
+
+        Returns:
+            Extra metadata added.
+        """
+
+        extra_metadata.time_elapsed = (datetime.datetime.now() - self.start_time).total_seconds() / 60
+        extra_metadata.current_trial = len(self.outcome_history)
+        extra_metadata.responses = sum([1 for choice in self.is_right_choice_history if choice is not None])
+        extra_metadata.ignored = sum([1 for choice in self.is_right_choice_history if choice is None])
+        extra_metadata.earned_water = sum(
+            [
+                oc.trial.reward_size.left
+                for oc in self.outcome_history
+                if oc.is_rewarded and not oc.is_right_choice and oc.trial.is_auto_reward_right is None
+            ]
+            + [
+                oc.trial.reward_size.right
+                for oc in self.outcome_history
+                if oc.is_rewarded and oc.is_right_choice and oc.trial.is_auto_reward_right is None
+            ]
+        )
+        extra_metadata.total_water = sum(
+            [oc.trial.reward_size.left for oc in self.outcome_history if oc.is_rewarded and not oc.is_right_choice]
+            + [oc.trial.reward_size.right for oc in self.outcome_history if oc.is_rewarded and oc.is_right_choice]
+        )
+        extra_metadata.foraging_efficiency = calculate_foraging_efficiency(
+            is_baiting=self.spec.is_baiting,
+            is_rewarded=self.reward_history,
+            p_left_reward=[oc.trial.metadata.p_reward_left for oc in self.outcome_history],
+            p_right_reward=[oc.trial.metadata.p_reward_right for oc in self.outcome_history],
+        )
+        return extra_metadata
 
     def get_metrics(self) -> TrialMetrics:
         """Return metrics at current state of the trial generator."""

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py

@@ -5,6 +5,9 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
+from ..block_based_trial_generator import (
+    BlockBasedTrialMetadata,
+)
 from .base_coupled_trial_generator import (
     BaseCoupledTrialGenerator,
     BaseCoupledTrialGeneratorSpec,
@@ -93,15 +96,21 @@ class CoupledTrialGenerator(BaseCoupledTrialGenerator):
 
     spec: CoupledTrialGeneratorSpec
 
-    def __init__(self, spec: CoupledTrialGeneratorSpec) -> None:
-        """Initializes the generator and records the session start time.
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
+        """Adds time remaining metadata to the trial.
 
         Args:
-            spec: The CoupledTrialGeneratorSpec defining task parameters.
+            extra_metadata: The extra metadata to which additional metadata will be added.
+
+        Returns:
+            The extra metadata with additional metadata.
         """
 
-        super().__init__(spec)
-        self.start_time = datetime.now()
+        extra_metadata = super()._add_extra_metadata(extra_metadata)
+        extra_metadata.time_remaining = (
+            timedelta(seconds=self.spec.trial_generation_end_parameters.max_time) - (datetime.now() - self.start_time)
+        ).total_seconds() / 60
+        return extra_metadata
 
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.

src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py

@@ -15,6 +15,7 @@
     Block,
     BlockBasedTrialGenerator,
     BlockBasedTrialGeneratorSpec,
+    BlockBasedTrialMetadata,
 )
 
 logger = logging.getLogger(__name__)
@@ -112,7 +113,6 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
         """
 
         super().__init__(spec)
-        self.start_time = datetime.now()
 
         block_length_min = spec.block_length.distribution_parameters.min
         block_length_max = spec.block_length.distribution_parameters.max
@@ -126,6 +126,22 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
         self.trials_in_left_block = 0
         self.left_dominance_streak = 0
 
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
+        """Adds time remaining metadata to the trial.
+
+        Args:
+            extra_metadata: The extra metadata to which additional metadata will be added.
+
+        Returns:
+            The extra metadata with additional metadata.
+        """
+
+        extra_metadata = super()._add_extra_metadata(extra_metadata)
+        extra_metadata.time_remaining = (
+            timedelta(seconds=self.spec.trial_generation_end_parameters.max_time) - (datetime.now() - self.start_time)
+        ).total_seconds() / 60
+        return extra_metadata
+
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
 

src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py

@@ -1,3 +1,4 @@
 from .calculate_bias import calculate_bias
+from .calculate_foraging_efficiency import calculate_foraging_efficiency
 
-__all__ = ["calculate_bias"]
+__all__ = ["calculate_bias", "calculate_foraging_efficiency"]

src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py

@@ -0,0 +1,69 @@
+import logging
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_foraging_efficiency(
+    is_baiting: bool, is_rewarded: list[bool], p_right_reward: list[float], p_left_reward: list[float]
+) -> Optional[float]:
+    """
+    Compute foraging efficiency for a two-arm bandit task.
+
+    This function calculates the ratio of actual rewards obtained to the
+    optimal expected rewards for a session. The implementation is adapted from the Allen Institute dynamic foraging
+    analysis codebase.
+
+    Args:
+        is_baiting (bool):
+            Whether the task uses a baiting schedule. If True, rewards can
+            accumulate on unchosen options; if False, rewards are independent
+            per trial.
+
+        is_rewarded (list[bool | None]):
+            List indicating whether each trial resulted in a reward. `True`
+            indicates a rewarded trial, `False` indicates no reward.
+
+        p_right_reward (list[float]):
+            Probability of reward for the right option on each trial.
+
+        p_left_reward (list[float]):
+            Probability of reward for the left option on each trial.
+
+    Returns:
+        float:
+            Foraging efficiency, defined as the ratio of the number of
+            rewarded trials to the optimal expected number of rewards for
+            the session.
+
+    Raises:
+        ValueError:
+            If input lists have mismatched lengths.
+
+    Notes:
+        Adapted from:
+        https://github.com/AllenNeuralDynamics/aind-dynamic-foraging-basic-analysis/blob/main/src/aind_dynamic_foraging_basic_analysis/metrics/foraging_efficiency.py
+    """
+
+    if not is_baiting:
+        logger.debug("Calculated non baiting foraging efficiency.")
+        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)
+        p_min = np.minimum(p_left_reward, p_right_reward)
+
+        with np.errstate(divide="ignore", invalid="ignore"):
+            optimal_visit_ratio = np.floor(np.log(1 - p_max) / np.log(1 - p_min))
+            optimal_general_reward_rates = p_max + (1 - (1 - p_min) ** (optimal_visit_ratio + 1) - p_max**2) / (
+                optimal_visit_ratio + 1
+            )
+
+        simple_case = (p_min == 0) | (p_max >= 1)
+        optimal_reward_per_trial = np.where(simple_case, p_max, optimal_general_reward_rates)
+
+        optimal_rewards_per_session = np.nanmean(optimal_reward_per_trial) * len(p_left_reward)
+    foraging_efficiency = float(is_rewarded.count(True) / optimal_rewards_per_session)
+    return round(foraging_efficiency, 3)

tests/trial_generators/test_coupled_trial_generator.py

@@ -19,7 +19,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -33,6 +33,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:

tests/trial_generators/test_uncoupled_trial_generator.py

@@ -26,7 +26,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -40,6 +40,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:

tests/trial_generators/test_warmup_trial_generator.py

@@ -21,7 +21,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -35,6 +35,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:

tests/trial_generators/util.py

@@ -1,10 +1,16 @@
+from typing import Optional
+
 import numpy as np
 
 from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome
 
 
 def simulate_response(
-    previous_reward: bool, previous_choice: bool | None, previous_left_bait: bool, previous_right_bait: bool
+    previous_reward: bool,
+    previous_choice: bool | None,
+    previous_left_bait: bool,
+    previous_right_bait: bool,
+    trial: Optional[Trial] = None,
 ) -> TrialOutcome:
 
     np.random.seed(42)
@@ -23,4 +29,4 @@ def simulate_response(
     else:
         is_rewarded = previous_right_bait if is_right_choice else previous_left_bait
 
-    return TrialOutcome(trial=Trial(), is_right_choice=is_right_choice, is_rewarded=is_rewarded)
+    return TrialOutcome(trial=trial or Trial(), is_right_choice=is_right_choice, is_rewarded=is_rewarded)