diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 8158162..fd6a892 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -339,17 +339,32 @@ def _side_bias(payload: t.Any) -> t.Optional[float]: # ------------------------------------------------------------------ # @staticmethod def _rewarded_history( - is_rewarded: bool, is_right_choice: t.Optional[bool], *, is_right: bool + trial: Trial, + is_rewarded: bool, + is_right_choice: t.Optional[bool], + *, + is_right: bool, ) -> bool: - """Return if mouse was rewarded based on choice. + """Return whether the mouse *earned* reward on the requested side. + + ``rewarded_history`` records earned reward only, i.e. water the animal + worked for. ``TrialOutcome.is_rewarded`` is ``True`` for any water + delivered on the trial, autowater included, so an autowater trial + (``trial.is_auto_reward_right is not None``) is ``False`` on *both* + sides here — its water is reported by ``auto_waterL``/``auto_waterR`` + instead. This matches the ``earned``/``automatic`` split in + :func:`~dynamic_foraging_processing.utils.rewards.get_annotated_rewards`. - A trial with no reward or an ignored trial (no choice) counts as not - rewarded on either side (``False``). + A trial with no reward or an ignored trial (no choice) likewise counts + as not rewarded on either side (``False``). Parameters ---------- + trial : Trial + The per-trial task-logic model; ``is_auto_reward_right`` being set + (to either side) marks the trial as an autowater trial. is_rewarded : bool - Whether the trial delivered reward. + Whether the trial delivered reward (earned *or* auto). is_right_choice : bool or None ``True`` for a right choice, ``False`` for left, ``None`` for ignored. is_right : bool @@ -358,10 +373,10 @@ def _rewarded_history( Returns ------- bool - ``True`` only when the trial was rewarded and the choice was on the - requested side; ``False`` otherwise. + ``True`` only when the trial delivered reward with no autowater and + the choice was on the requested side; ``False`` otherwise. """ - if is_right_choice is None: + if is_right_choice is None or trial.is_auto_reward_right is not None: return False return is_rewarded and (is_right_choice is is_right) @@ -862,8 +877,12 @@ def _build_row( **periods, delay_start_time=start, animal_response=self._animal_response(response), - rewarded_historyL=self._rewarded_history(is_rewarded, is_right_choice, is_right=False), - rewarded_historyR=self._rewarded_history(is_rewarded, is_right_choice, is_right=True), + rewarded_historyL=self._rewarded_history( + trial, is_rewarded, is_right_choice, is_right=False + ), + rewarded_historyR=self._rewarded_history( + trial, is_rewarded, is_right_choice, is_right=True + ), goCue_start_time=self._closest_time_in_window(go_cue_times, start, stop), left_valve_open_time=left_valve_open_time, right_valve_open_time=right_valve_open_time, diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index f34a57a..46af23f 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -69,10 +69,16 @@ class TrialConfig(BaseModel): description="The response of the animal. 0, left choice; 1, right choice; 2, no response", ) rewarded_historyL: bool = Field( - default=False, description="The reward history of left lick port" + default=False, + description=( + "The earned reward history of the left lick port; False on autowater trials, whose water is reported by auto_waterL" + ), ) rewarded_historyR: bool = Field( - default=False, description="The reward history of right lick port" + default=False, + description=( + "The earned reward history of the right lick port; False on autowater trials, whose water is reported by auto_waterR" + ), ) delay_start_time: Optional[float] = Field( default=None, diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index ebadf13..90de39a 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -662,6 +662,30 @@ def test_is_baited_forfeited_by_auto_response_on_same_side(): assert TrialTableBuilder._is_baited(trial, is_right=True) is False +def test_rewarded_history_is_earned_reward_only(): + """Rewarded history is the choice side on earned trials and False otherwise.""" + earned = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=None) + ).trial + assert TrialTableBuilder._rewarded_history(earned, True, True, is_right=True) is True + assert TrialTableBuilder._rewarded_history(earned, True, True, is_right=False) is False + # An unrewarded trial is False on both sides. + assert TrialTableBuilder._rewarded_history(earned, False, True, is_right=True) is False + # An ignored trial (no choice) is False on both sides. + assert TrialTableBuilder._rewarded_history(earned, True, None, is_right=True) is False + assert TrialTableBuilder._rewarded_history(earned, True, None, is_right=False) is False + + +def test_rewarded_history_false_on_every_auto_reward_trial(): + """Autowater is not earned: an auto-reward trial is False on both sides.""" + for auto in (True, False): + trial = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=auto, is_rewarded=True, auto=auto) + ).trial + assert TrialTableBuilder._rewarded_history(trial, True, auto, is_right=True) is False + assert TrialTableBuilder._rewarded_history(trial, True, auto, is_right=False) is False + + def test_auto_water_encodes_side_from_auto_response(): """A non-null auto response encodes ``1`` on its side and ``0`` on the other.""" trial = TrialOutcome.model_validate(