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 417a944..934d90a 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 @@ -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.""" diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py index 897e5d2..11d75d3 100644 --- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py +++ b/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. diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py index 61ebf6e..189448a 100644 --- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py +++ b/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. diff --git a/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py b/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py index ed9492e..7fea047 100644 --- a/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py +++ b/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"] diff --git a/src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py b/src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py new file mode 100644 index 0000000..99ac12a --- /dev/null +++ b/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) diff --git a/tests/trial_generators/test_coupled_trial_generator.py b/tests/trial_generators/test_coupled_trial_generator.py index dca8efd..2d259e7 100644 --- a/tests/trial_generators/test_coupled_trial_generator.py +++ b/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: diff --git a/tests/trial_generators/test_uncoupled_trial_generator.py b/tests/trial_generators/test_uncoupled_trial_generator.py index d952172..b590588 100644 --- a/tests/trial_generators/test_uncoupled_trial_generator.py +++ b/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: diff --git a/tests/trial_generators/test_warmup_trial_generator.py b/tests/trial_generators/test_warmup_trial_generator.py index 80b4fc5..291d03b 100644 --- a/tests/trial_generators/test_warmup_trial_generator.py +++ b/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: diff --git a/tests/trial_generators/util.py b/tests/trial_generators/util.py index 9ae5625..55ecae3 100644 --- a/tests/trial_generators/util.py +++ b/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) diff --git a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py index d49731c..9010a4a 100644 --- a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py +++ b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py @@ -2,9 +2,9 @@ import os from typing import Annotated, List, Literal, Optional -import numpy as np from aind_behavior_curriculum import Metrics from aind_behavior_dynamic_foraging.data_contract import dataset as df_foraging_dataset +from aind_behavior_dynamic_foraging.task_logic.utils.calculate_foraging_efficiency import calculate_foraging_efficiency from pydantic import BeforeValidator, Field STAGE_NAMES = Literal["stage_1_warmup", "stage_1", "stage_2", "stage_3", "final", "graduated"] @@ -72,9 +72,9 @@ def metrics_from_dataset( ] is_right_choice = [to["is_right_choice"] for to in filtered] is_rewarded = [to["is_rewarded"] for to in filtered] - p_right_reward = [to["trial"]["p_reward_right"] for to in filtered] - p_left_reward = [to["trial"]["p_reward_left"] for to in filtered] - foraging_efficiency = compute_foraging_efficiency( + p_right_reward = [to["trial"]["metadata"]["p_reward_right"] for to in filtered] + p_left_reward = [to["trial"]["metadata"]["p_reward_left"] for to in filtered] + foraging_efficiency = calculate_foraging_efficiency( is_baiting=is_baiting, is_rewarded=is_rewarded, p_left_reward=p_left_reward, p_right_reward=p_right_reward ) logger.debug(f"Calculated foraging efficiency as {foraging_efficiency}") @@ -104,64 +104,5 @@ def metrics_from_dataset( ) -def compute_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) +if __name__ == "__main__": + print(metrics_from_dataset(r"C:\Users\micah.woodard\Downloads\864253_2026-07-24T194251Z").model_dump_json(indent=4))