From 120f68df96209e0dee7890f10e88caad809feebf Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 13 Aug 2026 19:47:09 -0700 Subject: [PATCH 01/13] fix: try to use response times when correlating with hardware reward delivery times --- .../nwb/acquisition/acquisition_builder.py | 28 +++++++++++++++ .../utils/rewards.py | 34 ++++++++++++++++--- 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index ebd5270..3c11462 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -76,6 +76,26 @@ def get_trial_outcomes(self) -> pd.DataFrame: self.loader.dataset.at("Behavior").at("SoftwareEvents").at("TrialOutcome").load().data ) + def get_response_times(self) -> np.ndarray: + """Get the per-trial ``Response`` software-event timestamps. + + The event fires when the animal's choice is registered, within + milliseconds of the valve opening, so it anchors a reward delivery to + its trial. Only the event timestamp is used; the payload's ``Item1`` + field nominally carries a response time but is unreliable (it can lag + the event by thousands of seconds), so it is ignored. + + Returns + ------- + numpy.ndarray + The ``Response`` event timestamps, positionally aligned with the + ``TrialOutcome`` stream. + """ + responses = ( + self.loader.dataset.at("Behavior").at("SoftwareEvents").at("Response").load().data + ) + return responses.index.to_numpy() + def get_manual_water_times(self) -> pd.DataFrame: """Get the manual-water software-event stream. @@ -172,6 +192,7 @@ def _reward_delivery_series( writes: pd.DataFrame, trial_outcomes: pd.DataFrame, manual_water: pd.DataFrame, + response_times: np.ndarray, *, port_column: str, is_right: bool, @@ -193,6 +214,9 @@ def _reward_delivery_series( manual_water : pandas.DataFrame The ``GiveManualWaterRight`` stream; the ``data`` column selects the side (``True`` right, ``False`` left). + response_times : numpy.ndarray + ``Response`` event timestamps, one per trial, used to match each + delivery to its trial. port_column : str Supply-port column for this side (``"SupplyPort0"`` left, ``"SupplyPort1"`` right). @@ -215,6 +239,7 @@ def _reward_delivery_series( delivery_times, trial_outcomes, manual_water_times, + response_times, ) return AcquisitionSeries( name=name, @@ -251,6 +276,7 @@ def build_acquisition( rewards = self.get_reward_delivery() trial_outcomes = self.get_trial_outcomes() manual_water = self.get_manual_water_times() + response_times = self.get_response_times() acquisition_streams = self.loader.get_all_raw_data() acqusition_streams_descriptions = self.loader.raw_data_stream_descriptions @@ -274,6 +300,7 @@ def build_acquisition( rewards, trial_outcomes, manual_water, + response_times, port_column="SupplyPort0", is_right=False, name="left_reward_delivery_time", @@ -285,6 +312,7 @@ def build_acquisition( rewards, trial_outcomes, manual_water, + response_times, port_column="SupplyPort1", is_right=True, name="right_reward_delivery_time", diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 02234a8..fe91004 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -34,6 +34,7 @@ def get_annotated_rewards( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, + response_times: np.ndarray, ) -> np.ndarray: """Annotate each reward delivery as ``earned``, ``auto``, or ``manual``. @@ -49,8 +50,17 @@ def get_annotated_rewards( (``is_auto_reward_right is not None``). - ``earned`` -- otherwise (no matching trial, or no auto-response). - Deliveries are matched to trials by the ``TrialOutcome`` software-event - timestamp: each delivery takes the annotation of the closest trial. + Deliveries are matched to trials by the ``Response`` software-event + timestamp: each delivery takes the annotation of the trial whose response is + closest. The response is used rather than the ``TrialOutcome`` timestamp + because ``TrialOutcome`` fires at the *end* of a trial, after the + reward-consumption and ITI periods, while the valve opens within + milliseconds of the response. Matching on trial end lets a delivery land + nearer the *previous* trial's outcome and inherit its + ``is_auto_reward_right``, flipping ``earned`` and ``auto``. + + ``response_times`` is aligned to ``trial_outcome_df`` positionally: entry + ``i`` is the response of the trial in row ``i``. Parameters ---------- @@ -62,22 +72,36 @@ def get_annotated_rewards( manual_water_times : numpy.ndarray Software-event timestamps of this port's manual water deliveries (``GiveManualWaterLeft`` / ``GiveManualWaterRight``). + response_times : numpy.ndarray + ``Response`` software-event timestamps, one per trial, positionally + aligned with the rows of ``trial_outcome_df``. Returns ------- numpy.ndarray Array of the same shape as ``reward_delivery_times`` whose entries are ``"earned"``, ``"auto"``, or ``"manual"``. + + Raises + ------ + ValueError + If ``response_times`` has a different length than ``trial_outcome_df``, + since the two are paired by position. """ + response_times = np.asarray(response_times) + if response_times.size != len(trial_outcome_df): + raise ValueError( + f"response_times has {response_times.size} entries but there are " + f"{len(trial_outcome_df)} trials; the two are paired by position." + ) + reward_times = np.asarray(reward_delivery_times) if reward_times.size == 0: return np.array([], dtype=object) # Annotate each delivery from its originating trial: query with reward_times so we # get one trial position per reward delivery. - trial_indices_in_reward_times = find_closest_timestamps( - reward_times, trial_outcome_df.index.to_numpy() - ) + trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) annotated_rewards = [] for trial_index in trial_indices_in_reward_times: From ccf9926726c8d4ff6bac0367664765fbf5eeb110 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 13 Aug 2026 20:36:43 -0700 Subject: [PATCH 02/13] fix: exclude unrewarded trials --- .../nwb/acquisition/acquisition_builder.py | 5 ++- .../utils/rewards.py | 37 ++++++++++++++----- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index 3c11462..779b730 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -203,7 +203,8 @@ def _reward_delivery_series( Only valve-open events (``port_column`` is truthy) are reward deliveries; the ``data`` field annotates each as earned, manual, or - auto via :func:`get_annotated_rewards`. + auto via :func:`get_annotated_rewards`, which also drops the autowater + deliveries the animal never collected. Parameters ---------- @@ -235,7 +236,7 @@ def _reward_delivery_series( open_writes = writes[writes[port_column].fillna(False).astype(bool)] delivery_times = open_writes.index.to_numpy() manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy() - annotations = get_annotated_rewards( + delivery_times, annotations = get_annotated_rewards( delivery_times, trial_outcomes, manual_water_times, diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index fe91004..07c2d7e 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -35,7 +35,7 @@ def get_annotated_rewards( trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, response_times: np.ndarray, -) -> np.ndarray: +) -> t.Tuple[np.ndarray, np.ndarray]: """Annotate each reward delivery as ``earned``, ``auto``, or ``manual``. Annotates the deliveries of a single lick port. Each delivery is classified @@ -50,6 +50,13 @@ def get_annotated_rewards( (``is_auto_reward_right is not None``). - ``earned`` -- otherwise (no matching trial, or no auto-response). + Deliveries the animal never received are dropped rather than annotated: + autowater is delivered *before* the response, so on an auto trial the animal + answered the other way the water goes uncollected and the trial reports + ``is_rewarded=False``. Manual water keeps its precedence and is never + dropped. The surviving timestamps are returned alongside their annotations + so the two stay aligned. + Deliveries are matched to trials by the ``Response`` software-event timestamp: each delivery takes the annotation of the trial whose response is closest. The response is used rather than the ``TrialOutcome`` timestamp @@ -79,7 +86,10 @@ def get_annotated_rewards( Returns ------- numpy.ndarray - Array of the same shape as ``reward_delivery_times`` whose entries are + The retained reward-delivery timestamps: ``reward_delivery_times`` less + the uncollected autowater deliveries. + numpy.ndarray + The matching annotations, one per retained timestamp, each ``"earned"``, ``"auto"``, or ``"manual"``. Raises @@ -97,20 +107,20 @@ def get_annotated_rewards( reward_times = np.asarray(reward_delivery_times) if reward_times.size == 0: - return np.array([], dtype=object) + return reward_times, np.array([], dtype=object) # Annotate each delivery from its originating trial: query with reward_times so we # get one trial position per reward delivery. trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) annotated_rewards = [] + is_unrewarded_auto = [] for trial_index in trial_indices_in_reward_times: outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"]) trial = outcome.trial if outcome is not None else None - if trial is None or trial.is_auto_reward_right is None: - annotated_rewards.append("earned") - else: - annotated_rewards.append("auto") + is_auto = trial is not None and trial.is_auto_reward_right is not None + annotated_rewards.append("auto" if is_auto else "earned") + is_unrewarded_auto.append(is_auto and outcome is not None and not outcome.is_rewarded) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" # and "earned" entries would be too narrow to hold "manual" and would truncate it. @@ -121,8 +131,15 @@ def get_annotated_rewards( # manual-water software event to its closest reward delivery; the returned # positions index into reward_times, i.e. the deliveries that are manual. manual_water_times = np.asarray(manual_water_times) + manual_mask = np.zeros(reward_times.size, dtype=bool) if manual_water_times.size: manual_indices_in_reward_times = find_closest_timestamps(manual_water_times, reward_times) - annotated_rewards[manual_indices_in_reward_times] = "manual" - - return annotated_rewards + manual_mask[manual_indices_in_reward_times] = True + annotated_rewards[manual_mask] = "manual" + + # Autowater is delivered before the response, so on a trial the animal answered the + # other way the water is never collected and the trial reports is_rewarded=False. + # Those deliveries are dropped: they are not reward the animal received. Manual + # water keeps its precedence and is never dropped. + keep = ~(np.array(is_unrewarded_auto, dtype=bool) & ~manual_mask) + return reward_times[keep], annotated_rewards[keep] From 4b2c6f18fbdebd7091f716f901d4ae93b4353e53 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 13 Aug 2026 21:02:06 -0700 Subject: [PATCH 03/13] refactor: overhaul function names to match behavior --- .../nwb/acquisition/acquisition_builder.py | 12 ++++++------ .../processing/_trial_table.py | 2 +- src/dynamic_foraging_processing/utils/__init__.py | 4 ++-- src/dynamic_foraging_processing/utils/rewards.py | 4 ++-- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index 779b730..be45f0e 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -11,7 +11,7 @@ ) from dynamic_foraging_processing.nwb.utils import clean_for_nwb from dynamic_foraging_processing.raw_data_loader import RawDataLoader -from dynamic_foraging_processing.utils.rewards import get_annotated_rewards +from dynamic_foraging_processing.utils.rewards import get_reward_deliveries class LickSource(t.NamedTuple): @@ -49,8 +49,8 @@ def __init__(self, loader: RawDataLoader): """ self.loader = loader - def get_reward_delivery(self) -> pd.DataFrame: - """Get the reward delivery stream from the dataset. + def get_valve_writes(self) -> pd.DataFrame: + """Get the raw valve command stream. Returns ------- @@ -203,7 +203,7 @@ def _reward_delivery_series( Only valve-open events (``port_column`` is truthy) are reward deliveries; the ``data`` field annotates each as earned, manual, or - auto via :func:`get_annotated_rewards`, which also drops the autowater + auto via :func:`get_reward_deliveries`, which also drops the autowater deliveries the animal never collected. Parameters @@ -236,7 +236,7 @@ def _reward_delivery_series( open_writes = writes[writes[port_column].fillna(False).astype(bool)] delivery_times = open_writes.index.to_numpy() manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy() - delivery_times, annotations = get_annotated_rewards( + delivery_times, annotations = get_reward_deliveries( delivery_times, trial_outcomes, manual_water_times, @@ -274,7 +274,7 @@ def build_acquisition( list of AcquisitionSeries or AcquisitionTable Acquisition entries to write to the NWB acquisition module. """ - rewards = self.get_reward_delivery() + rewards = self.get_valve_writes() trial_outcomes = self.get_trial_outcomes() manual_water = self.get_manual_water_times() response_times = self.get_response_times() diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 724828b..5760fed 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -353,7 +353,7 @@ def _rewarded_history( (``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`. + :func:`~dynamic_foraging_processing.utils.rewards.get_reward_deliveries`. A trial with no reward or an ignored trial (no choice) likewise counts as not rewarded on either side (``False``). diff --git a/src/dynamic_foraging_processing/utils/__init__.py b/src/dynamic_foraging_processing/utils/__init__.py index 1c77e9b..cecf2d5 100644 --- a/src/dynamic_foraging_processing/utils/__init__.py +++ b/src/dynamic_foraging_processing/utils/__init__.py @@ -1,6 +1,6 @@ """Utility helpers for dynamic foraging processing.""" -from dynamic_foraging_processing.utils.rewards import get_annotated_rewards +from dynamic_foraging_processing.utils.rewards import get_reward_deliveries from dynamic_foraging_processing.utils.timestamps import find_closest_timestamps -__all__ = ["find_closest_timestamps", "get_annotated_rewards"] +__all__ = ["find_closest_timestamps", "get_reward_deliveries"] diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 07c2d7e..ed26504 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -30,13 +30,13 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: return TrialOutcome.model_validate(payload) -def get_annotated_rewards( +def get_reward_deliveries( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, response_times: np.ndarray, ) -> t.Tuple[np.ndarray, np.ndarray]: - """Annotate each reward delivery as ``earned``, ``auto``, or ``manual``. + """Get one lick port's reward deliveries, each ``earned``, ``auto``, or ``manual``. Annotates the deliveries of a single lick port. Each delivery is classified as follows, with ``manual`` taking precedence because manual water is not From 64e46e358ce98a5f8e6abb97a29f6a7c7f6191dc Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 13 Aug 2026 21:02:21 -0700 Subject: [PATCH 04/13] test: update tests --- .../test_acquisition_builder.py | 15 +- tests/test_utils/test_rewards.py | 170 ++++++++++++++---- 2 files changed, 146 insertions(+), 39 deletions(-) diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index f16915e..73d6dd3 100644 --- a/tests/test_nwb/test_acquisition/test_acquisition_builder.py +++ b/tests/test_nwb/test_acquisition/test_acquisition_builder.py @@ -82,6 +82,14 @@ def _make_trial_outcome_frame() -> pd.DataFrame: ) +def _make_response_frame() -> pd.DataFrame: + """One ``Response`` event per trial, just before each trial's outcome.""" + return pd.DataFrame( + {"data": [{"Item1": 0.05, "Item2": False}, {"Item1": 0.35, "Item2": True}]}, + index=pd.Index([0.05, 0.35], name="time"), + ) + + def _empty_manual_water_frame() -> pd.DataFrame: """Build an empty manual-water stream with the ``data`` side column.""" return pd.DataFrame({"data": []}, index=pd.Index([], name="time")) @@ -120,6 +128,7 @@ def _make_dataset(manual_water=None): "SoftwareEvents": _FakeNode( { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), + "Response": _FakeStream(_make_response_frame()), "GiveManualWaterRight": _FakeStream(manual_water), } ), @@ -147,11 +156,11 @@ def test_init_stores_loader(): assert builder.loader is loader -def test_get_reward_delivery_filters_to_write_messages(): +def test_get_valve_writes_filters_to_write_messages(): """Only ``MessageType == 'WRITE'`` rows are returned.""" builder = AcquisitionBuilder(loader=_make_loader()) - result = builder.get_reward_delivery() + result = builder.get_valve_writes() assert list(result["MessageType"]) == ["WRITE", "WRITE", "WRITE"] assert list(result.index) == [0.1, 0.3, 0.5] @@ -177,6 +186,7 @@ def test_get_manual_water_times_returns_empty_when_absent(): "SoftwareEvents": _FakeNode( { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), + "Response": _FakeStream(_make_response_frame()), } ), } @@ -224,6 +234,7 @@ def test_get_lick_times_returns_empty_when_absent(): "SoftwareEvents": _FakeNode( { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), + "Response": _FakeStream(_make_response_frame()), "GiveManualWaterRight": _FakeStream(_empty_manual_water_frame()), } ), diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index 2be8269..f2b94a0 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -4,12 +4,13 @@ import numpy as np import pandas as pd +import pytest from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome -from dynamic_foraging_processing.utils.rewards import get_annotated_rewards +from dynamic_foraging_processing.utils.rewards import get_reward_deliveries -def _outcome_payload(auto=None) -> dict: +def _outcome_payload(auto=None, is_rewarded: bool = True) -> dict: """Return a serialized ``TrialOutcome`` payload with the given auto-response.""" return { "trial": { @@ -22,103 +23,198 @@ def _outcome_payload(auto=None) -> dict: "is_auto_reward_right": auto, }, "is_right_choice": True, - "is_rewarded": True, + "is_rewarded": is_rewarded, } -def _trial_outcome_df(trial_times: np.ndarray, autos=None) -> pd.DataFrame: +def _trial_outcome_df(trial_times: np.ndarray, autos=None, rewarded=None) -> pd.DataFrame: """Build a trial outcome DataFrame with one row per entry of ``trial_times``.""" autos = autos if autos is not None else [None] * len(trial_times) + rewarded = rewarded if rewarded is not None else [True] * len(trial_times) return pd.DataFrame( - {"data": [_outcome_payload(auto) for auto in autos]}, + {"data": [_outcome_payload(a, r) for a, r in zip(autos, rewarded)]}, index=pd.Index(trial_times, name="time"), ) -def test_get_annotated_rewards_marks_default_trials_as_earned(): +def test_get_reward_deliveries_marks_default_trials_as_earned(): """Trials with no auto-response setting and no manual water are ``earned``.""" reward_times = np.array([0.15, 0.42, 0.95]) - trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4, 0.9])) + response_times = np.array([0.1, 0.4, 0.9]) + trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4, 1.9])) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["earned", "earned", "earned"])) -def test_get_annotated_rewards_marks_auto_response_trials_as_auto(): +def test_get_reward_deliveries_marks_auto_response_trials_as_auto(): """Trials with ``is_auto_reward_right`` set (either side) are ``auto``.""" reward_times = np.array([0.15, 0.42]) - trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4]), autos=[True, False]) + response_times = np.array([0.1, 0.4]) + trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4]), autos=[True, False]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -def test_get_annotated_rewards_matches_closest_trial_outcome_time(): - """Each delivery takes the annotation of the closest ``TrialOutcome`` event.""" - # Both deliveries sit nearest the second (auto) trial, so both are auto even - # though the first trial is earned. +def test_get_reward_deliveries_matches_closest_response_time(): + """Each delivery takes the annotation of the trial whose response is closest. + + The trial-outcome timestamps deliberately disagree with the response times: + matching on the outcome would pick the first (earned) trial, so this pins the + match to the ``Response`` stream. + """ reward_times = np.array([0.95, 1.05]) - trial_outcome_df = _trial_outcome_df(np.array([0.1, 1.0]), autos=[None, True]) + response_times = np.array([0.1, 1.0]) + # Outcome events fire at the end of each trial, far from the deliveries. + trial_outcome_df = _trial_outcome_df(np.array([0.9, 5.0]), autos=[None, True]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -def test_get_annotated_rewards_marks_manual_water_as_manual(): +def test_get_reward_deliveries_drops_uncollected_auto_water(): + """Autowater on a trial reporting ``is_rewarded=False`` is dropped, not annotated. + + The water is delivered before the response, so a trial the animal answered + the other way leaves it uncollected; it is not reward the animal received. + """ + reward_times = np.array([0.15, 0.42, 0.95]) + response_times = np.array([0.1, 0.4, 0.9]) + trial_outcome_df = _trial_outcome_df( + np.array([1.1, 1.4, 1.9]), + autos=[None, True, True], + rewarded=[True, False, True], + ) + + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + + np.testing.assert_array_equal(times, np.array([0.15, 0.95])) + np.testing.assert_array_equal(annotations, np.array(["earned", "auto"])) + + +def test_get_reward_deliveries_keeps_unrewarded_non_auto_deliveries(): + """A delivery on an unrewarded trial that is not autowater is kept.""" + reward_times = np.array([0.15]) + response_times = np.array([0.1]) + trial_outcome_df = _trial_outcome_df(np.array([1.1]), autos=[None], rewarded=[False]) + + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + + np.testing.assert_array_equal(times, reward_times) + np.testing.assert_array_equal(annotations, np.array(["earned"])) + + +def test_get_reward_deliveries_marks_manual_water_as_manual(): """Deliveries closest to a manual-water event are annotated as ``manual``.""" reward_times = np.array([0.15, 0.42, 0.95]) - trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4, 0.9])) + response_times = np.array([0.1, 0.4, 0.9]) + trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4, 1.9])) # Software event near the second delivery (0.42). manual_water_times = np.array([0.43]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, manual_water_times) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water_times, response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["earned", "manual", "earned"])) -def test_get_annotated_rewards_manual_takes_precedence_over_auto(): +def test_get_reward_deliveries_manual_takes_precedence_over_auto(): """A manual delivery is ``manual`` even when the trial has auto-response set.""" reward_times = np.array([0.15, 0.42]) - trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4]), autos=[None, True]) + response_times = np.array([0.1, 0.4]) + trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4]), autos=[None, True]) manual_water_times = np.array([0.42]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, manual_water_times) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water_times, response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["earned", "manual"])) -def test_get_annotated_rewards_empty_deliveries_returns_empty(): - """No reward deliveries yields an empty annotation array.""" +def test_get_reward_deliveries_manual_water_survives_the_auto_drop(): + """Manual water on an unrewarded auto trial is kept, not dropped.""" + reward_times = np.array([0.15, 0.42]) + response_times = np.array([0.1, 0.4]) + trial_outcome_df = _trial_outcome_df( + np.array([1.1, 1.4]), autos=[None, True], rewarded=[True, False] + ) + manual_water_times = np.array([0.42]) + + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water_times, response_times + ) + + np.testing.assert_array_equal(times, reward_times) + np.testing.assert_array_equal(annotations, np.array(["earned", "manual"])) + + +def test_get_reward_deliveries_empty_deliveries_returns_empty(): + """No reward deliveries yields empty timestamp and annotation arrays.""" trial_outcome_df = _trial_outcome_df(np.array([0.0])) - result = get_annotated_rewards(np.array([]), trial_outcome_df, np.array([])) + times, annotations = get_reward_deliveries( + np.array([]), trial_outcome_df, np.array([]), np.array([0.0]) + ) - assert isinstance(result, np.ndarray) - assert result.size == 0 + assert isinstance(annotations, np.ndarray) + assert times.size == 0 + assert annotations.size == 0 -def test_get_annotated_rewards_accepts_json_and_model_payloads(): +def test_get_reward_deliveries_accepts_json_and_model_payloads(): """``data`` payloads may be JSON strings or already-parsed ``TrialOutcome``.""" reward_times = np.array([0.15, 0.42]) + response_times = np.array([0.1, 0.4]) payload = _outcome_payload(True) trial_outcome_df = pd.DataFrame( {"data": [json.dumps(payload), TrialOutcome.model_validate(payload)]}, - index=pd.Index([0.1, 0.4], name="time"), + index=pd.Index([1.1, 1.4], name="time"), ) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + np.testing.assert_array_equal(times, reward_times) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -def test_get_annotated_rewards_returns_ndarray(): - """The return value is a ``numpy.ndarray``.""" - reward_times = np.array([0.1]) - trial_outcome_df = _trial_outcome_df(np.array([0.0])) +def test_get_reward_deliveries_rejects_misaligned_response_times(): + """``response_times`` must have one entry per trial; they pair by position.""" + trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4])) - result = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + with pytest.raises(ValueError, match="paired by position"): + get_reward_deliveries(np.array([0.15]), trial_outcome_df, np.array([]), np.array([0.1])) + + +def test_get_reward_deliveries_returns_ndarray(): + """Both return values are :class:`numpy.ndarray`.""" + trial_outcome_df = _trial_outcome_df(np.array([1.0])) + + times, annotations = get_reward_deliveries( + np.array([0.1]), trial_outcome_df, np.array([]), np.array([0.0]) + ) - assert isinstance(result, np.ndarray) + assert isinstance(times, np.ndarray) + assert isinstance(annotations, np.ndarray) From c631c89db8d6f6aea990d23e07d52d0f295545b0 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Fri, 14 Aug 2026 11:06:25 -0700 Subject: [PATCH 05/13] refactor: make clear that non-rewwarded trials are excluded --- .../utils/rewards.py | 26 +++++++++---------- tests/test_utils/test_rewards.py | 12 ++++++--- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index ed26504..7734d29 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -50,12 +50,12 @@ def get_reward_deliveries( (``is_auto_reward_right is not None``). - ``earned`` -- otherwise (no matching trial, or no auto-response). - Deliveries the animal never received are dropped rather than annotated: - autowater is delivered *before* the response, so on an auto trial the animal - answered the other way the water goes uncollected and the trial reports - ``is_rewarded=False``. Manual water keeps its precedence and is never - dropped. The surviving timestamps are returned alongside their annotations - so the two stay aligned. + Deliveries matched to a trial reporting ``is_rewarded=False`` are dropped + rather than annotated: the water was not reward the animal received. In + practice these are all autowater, which is delivered *before* the response, + so on a trial the animal answered the other way it goes uncollected. Manual + water keeps its precedence and is never dropped. The surviving timestamps + are returned alongside their annotations so the two stay aligned. Deliveries are matched to trials by the ``Response`` software-event timestamp: each delivery takes the annotation of the trial whose response is @@ -114,13 +114,13 @@ def get_reward_deliveries( trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) annotated_rewards = [] - is_unrewarded_auto = [] + is_unrewarded = [] for trial_index in trial_indices_in_reward_times: outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"]) trial = outcome.trial if outcome is not None else None is_auto = trial is not None and trial.is_auto_reward_right is not None annotated_rewards.append("auto" if is_auto else "earned") - is_unrewarded_auto.append(is_auto and outcome is not None and not outcome.is_rewarded) + is_unrewarded.append(outcome is not None and not outcome.is_rewarded) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" # and "earned" entries would be too narrow to hold "manual" and would truncate it. @@ -137,9 +137,9 @@ def get_reward_deliveries( manual_mask[manual_indices_in_reward_times] = True annotated_rewards[manual_mask] = "manual" - # Autowater is delivered before the response, so on a trial the animal answered the - # other way the water is never collected and the trial reports is_rewarded=False. - # Those deliveries are dropped: they are not reward the animal received. Manual - # water keeps its precedence and is never dropped. - keep = ~(np.array(is_unrewarded_auto, dtype=bool) & ~manual_mask) + # A delivery on a trial reporting is_rewarded=False is not reward the animal + # received, so it is dropped. In practice these are all autowater: it is delivered + # before the response, so on a trial the animal answered the other way the water is + # never collected. Manual water keeps its precedence and is never dropped. + keep = ~(np.array(is_unrewarded, dtype=bool) & ~manual_mask) return reward_times[keep], annotated_rewards[keep] diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index f2b94a0..38bda02 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -107,8 +107,12 @@ def test_get_reward_deliveries_drops_uncollected_auto_water(): np.testing.assert_array_equal(annotations, np.array(["earned", "auto"])) -def test_get_reward_deliveries_keeps_unrewarded_non_auto_deliveries(): - """A delivery on an unrewarded trial that is not autowater is kept.""" +def test_get_reward_deliveries_drops_any_delivery_on_an_unrewarded_trial(): + """The drop rule is ``is_rewarded=False``, not autowater specifically. + + Autowater is the only case seen in practice, but a delivery on any trial + reporting no reward is water the animal did not receive. + """ reward_times = np.array([0.15]) response_times = np.array([0.1]) trial_outcome_df = _trial_outcome_df(np.array([1.1]), autos=[None], rewarded=[False]) @@ -117,8 +121,8 @@ def test_get_reward_deliveries_keeps_unrewarded_non_auto_deliveries(): reward_times, trial_outcome_df, np.array([]), response_times ) - np.testing.assert_array_equal(times, reward_times) - np.testing.assert_array_equal(annotations, np.array(["earned"])) + assert times.size == 0 + assert annotations.size == 0 def test_get_reward_deliveries_marks_manual_water_as_manual(): From 71290e72f78e7282936b01a719e8867b055fb566 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Fri, 14 Aug 2026 14:38:52 -0700 Subject: [PATCH 06/13] fix: try temp fix using is rewarded from trial outcome for auto water --- .../processing/_trial_table.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 5760fed..5f978c2 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -442,14 +442,14 @@ def _is_baited(trial: Trial, *, is_right: bool) -> bool: return trial.p_reward_left == 1 and auto in (None, True) @staticmethod - def _auto_water(trial: Trial, *, is_right: bool) -> int: + def _auto_water(trial: Trial, outcome: TrialOutcome, *, is_right: bool) -> int: """Encode autowater for a side from ``is_auto_reward_right``. Returns ``1`` if the auto response was to the requested side, else ``0``. No auto-response (``is_auto_reward_right`` is ``None``) counts as no autowater (``0``). ``is_right`` is ``True`` for right. """ - if trial.is_auto_reward_right is None: + if trial.is_auto_reward_right is None or not outcome.is_rewarded: return 0 return int(trial.is_auto_reward_right is is_right) @@ -897,8 +897,8 @@ def _build_row( reward_consumption_duration=trial.reward_consumption_duration, ITI_duration=trial.inter_trial_interval_duration, delay_duration=trial.quiescence_period_duration, - auto_waterL=self._auto_water(trial, is_right=False), - auto_waterR=self._auto_water(trial, is_right=True), + auto_waterL=self._auto_water(trial, outcome, is_right=False), + auto_waterR=self._auto_water(trial, outcome, is_right=True), anti_bias_left_water=self._anti_bias_water(trial, bias_metadata, is_right=False), anti_bias_right_water=self._anti_bias_water(trial, bias_metadata, is_right=True), anti_bias_lickspout_movement=self._anti_bias_lickspout_movement(trial, bias_metadata), From 2f8e5276a25b1528ef0306a8146e0d6c3060e620 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Fri, 14 Aug 2026 18:03:30 -0700 Subject: [PATCH 07/13] fix: test to use trial metadata for auto water to distinguish from anti bias --- .../nwb/acquisition/acquisition_builder.py | 11 ++- .../processing/_trial_table.py | 55 ++++++----- .../utils/rewards.py | 93 ++++++++++++------- .../utils/trial_metadata.py | 43 +++++++++ 4 files changed, 139 insertions(+), 63 deletions(-) create mode 100644 src/dynamic_foraging_processing/utils/trial_metadata.py diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index be45f0e..c4a7035 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -202,9 +202,10 @@ def _reward_delivery_series( """Build one lick port's reward-delivery series with reward annotations. Only valve-open events (``port_column`` is truthy) are reward - deliveries; the ``data`` field annotates each as earned, manual, or - auto via :func:`get_reward_deliveries`, which also drops the autowater - deliveries the animal never collected. + deliveries; the ``data`` field annotates each as earned, manual, auto, or + anti-bias via :func:`get_reward_deliveries`. Every valve opening is + reported, so the series stays a complete record of the water delivered at + this port. Parameters ---------- @@ -236,7 +237,7 @@ def _reward_delivery_series( open_writes = writes[writes[port_column].fillna(False).astype(bool)] delivery_times = open_writes.index.to_numpy() manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy() - delivery_times, annotations = get_reward_deliveries( + annotations = get_reward_deliveries( delivery_times, trial_outcomes, manual_water_times, @@ -249,7 +250,7 @@ def _reward_delivery_series( unit="second", description=( f"The reward delivery time of the {side_label} lick port. The data field " - "annotates whether the reward was earned, manual, or auto" + "annotates whether the reward was earned, manual, auto, or anti_bias" ), ) diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 5f978c2..123aece 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -25,6 +25,7 @@ from contraqctor.contract import Dataset from dynamic_foraging_processing.processing.models import TrialConfig +from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata logger = logging.getLogger(__name__) @@ -442,28 +443,42 @@ def _is_baited(trial: Trial, *, is_right: bool) -> bool: return trial.p_reward_left == 1 and auto in (None, True) @staticmethod - def _auto_water(trial: Trial, outcome: TrialOutcome, *, is_right: bool) -> int: - """Encode autowater for a side from ``is_auto_reward_right``. + def _auto_water(trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool) -> int: + """Return whether scheduled autowater was delivered to the requested side. + + ``is_auto_reward_right`` is the *delivery channel*: it says free water was + triggered and to which side (``True`` right, ``False`` left, ``None`` no + free water), but not what kind. Autowater and the anti-bias water + intervention share that channel and are told apart only by the metadata + flags, so this reads ``is_autowater`` and takes the side from the channel + -- the mirror of :meth:`_anti_bias_water`. A trial whose free water came + from the anti-bias algorithm is therefore ``0`` here, and is reported by + ``anti_bias_left_water``/``anti_bias_right_water`` instead. - Returns ``1`` if the auto response was to the requested side, else ``0``. - No auto-response (``is_auto_reward_right`` is ``None``) counts as no - autowater (``0``). ``is_right`` is ``True`` for right. + Parameters + ---------- + trial : Trial + The per-trial task-logic model. + bias_metadata : BlockBasedTrialMetadata + The trial's extra metadata (see ``_bias_metadata``). + is_right : bool + ``True`` for the right port, ``False`` for the left port. + + Returns + ------- + int + ``1`` when scheduled autowater targeted the requested side, else ``0``. """ - if trial.is_auto_reward_right is None or not outcome.is_rewarded: + if not bias_metadata.is_autowater: return 0 return int(trial.is_auto_reward_right is is_right) @staticmethod def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: - """Return the block-based extra metadata carrying the anti-bias flags. + """Return the block-based extra metadata carrying the free-water flags. - The anti-bias flags (``is_bias_water_intervention``, - ``is_bias_stage_intervention``) live on ``trial.metadata.extra``. That - field is schema-typed ``Any``, so it deserializes off the stream as a - plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` - instance is also accepted. When metadata or extra is missing (e.g. an - older session, or a non-block-based generator), the model's all-``False`` - default is returned so the anti-bias columns are simply inert. + Thin wrapper over :func:`get_bias_metadata`, shared with the reward + annotation so both read the autowater and anti-bias flags identically. Parameters ---------- @@ -476,13 +491,7 @@ def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: The parsed extra metadata, or an all-``False`` default when absent or unrecognized. """ - metadata = trial.metadata - extra = metadata.extra if metadata is not None else None - if isinstance(extra, BlockBasedTrialMetadata): - return extra - if isinstance(extra, dict): - return BlockBasedTrialMetadata.model_validate(extra) - return BlockBasedTrialMetadata() + return get_bias_metadata(trial) @staticmethod def _anti_bias_water( @@ -897,8 +906,8 @@ def _build_row( reward_consumption_duration=trial.reward_consumption_duration, ITI_duration=trial.inter_trial_interval_duration, delay_duration=trial.quiescence_period_duration, - auto_waterL=self._auto_water(trial, outcome, is_right=False), - auto_waterR=self._auto_water(trial, outcome, is_right=True), + auto_waterL=self._auto_water(trial, bias_metadata, is_right=False), + auto_waterR=self._auto_water(trial, bias_metadata, is_right=True), anti_bias_left_water=self._anti_bias_water(trial, bias_metadata, is_right=False), anti_bias_right_water=self._anti_bias_water(trial, bias_metadata, is_right=True), anti_bias_lickspout_movement=self._anti_bias_lickspout_movement(trial, bias_metadata), diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 7734d29..7047ac4 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -4,9 +4,10 @@ import numpy as np import pandas as pd -from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome +from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome from dynamic_foraging_processing.utils.timestamps import find_closest_timestamps +from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: @@ -30,13 +31,42 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: return TrialOutcome.model_validate(payload) +def _free_water_label(trial: t.Optional[Trial]) -> str: + """Classify a delivery's trial as ``auto``, ``anti_bias``, or ``earned``. + + ``is_auto_reward_right`` marks only that free water was triggered and on + which side; scheduled autowater and the anti-bias intervention share that + channel and are told apart by the block-based metadata flags. A trial with no + free water, or free water flagged as neither mechanism, is ``earned``. + + Parameters + ---------- + trial : Trial or None + The per-trial task-logic model, or ``None`` when the outcome payload was + missing. + + Returns + ------- + str + ``"auto"``, ``"anti_bias"``, or ``"earned"``. + """ + if trial is None or trial.is_auto_reward_right is None: + return "earned" + metadata = get_bias_metadata(trial) + if metadata.is_bias_water_intervention: + return "anti_bias" + if metadata.is_autowater: + return "auto" + return "earned" + + def get_reward_deliveries( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, response_times: np.ndarray, -) -> t.Tuple[np.ndarray, np.ndarray]: - """Get one lick port's reward deliveries, each ``earned``, ``auto``, or ``manual``. +) -> np.ndarray: + """Classify one lick port's reward deliveries by how the water was given. Annotates the deliveries of a single lick port. Each delivery is classified as follows, with ``manual`` taking precedence because manual water is not @@ -46,16 +76,23 @@ def get_reward_deliveries( ``GiveManualWater`` software event for this port. The software-event timestamps are correlated to the reward-delivery timestamps with :func:`find_closest_timestamps`. - - ``auto`` -- otherwise, when the matching trial auto-responded - (``is_auto_reward_right is not None``). - - ``earned`` -- otherwise (no matching trial, or no auto-response). - - Deliveries matched to a trial reporting ``is_rewarded=False`` are dropped - rather than annotated: the water was not reward the animal received. In - practice these are all autowater, which is delivered *before* the response, - so on a trial the animal answered the other way it goes uncollected. Manual - water keeps its precedence and is never dropped. The surviving timestamps - are returned alongside their annotations so the two stay aligned. + - ``anti_bias`` -- otherwise, when the trial's free water came from the + anti-bias algorithm (``is_bias_water_intervention``). + - ``auto`` -- otherwise, when the trial's free water was scheduled autowater + (``is_autowater``). + - ``earned`` -- otherwise: water the animal worked for. + + ``is_auto_reward_right`` is only the delivery *channel* -- it says free water + was triggered and on which side, not what kind -- so the ``auto`` versus + ``anti_bias`` split comes from the block-based metadata flags (see + :func:`get_bias_metadata`), mirroring ``auto_waterL``/``auto_waterR`` and + ``anti_bias_left_water``/``anti_bias_right_water`` in the trials table. + + Every delivery is annotated and none is filtered out. A trial reporting + ``is_rewarded=False`` keeps its delivery: free water is triggered immediately + at the go cue and the trial then continues normally, so ``is_rewarded`` + reports the outcome of the animal's own choice -- a separate event from the + water being classified here. Deliveries are matched to trials by the ``Response`` software-event timestamp: each delivery takes the annotation of the trial whose response is @@ -86,11 +123,8 @@ def get_reward_deliveries( Returns ------- numpy.ndarray - The retained reward-delivery timestamps: ``reward_delivery_times`` less - the uncollected autowater deliveries. - numpy.ndarray - The matching annotations, one per retained timestamp, each - ``"earned"``, ``"auto"``, or ``"manual"``. + Array of the same shape as ``reward_delivery_times`` whose entries are + ``"earned"``, ``"auto"``, ``"anti_bias"``, or ``"manual"``. Raises ------ @@ -107,23 +141,19 @@ def get_reward_deliveries( reward_times = np.asarray(reward_delivery_times) if reward_times.size == 0: - return reward_times, np.array([], dtype=object) + return np.array([], dtype=object) # Annotate each delivery from its originating trial: query with reward_times so we # get one trial position per reward delivery. trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) annotated_rewards = [] - is_unrewarded = [] for trial_index in trial_indices_in_reward_times: outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"]) - trial = outcome.trial if outcome is not None else None - is_auto = trial is not None and trial.is_auto_reward_right is not None - annotated_rewards.append("auto" if is_auto else "earned") - is_unrewarded.append(outcome is not None and not outcome.is_rewarded) + annotated_rewards.append(_free_water_label(outcome.trial if outcome is not None else None)) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" - # and "earned" entries would be too narrow to hold "manual" and would truncate it. + # and "earned" entries would be too narrow to hold "anti_bias" and would truncate it. annotated_rewards = np.array(annotated_rewards, dtype=object) # Manual water is independent of trials (multiple can occur within a trial) and @@ -131,15 +161,8 @@ def get_reward_deliveries( # manual-water software event to its closest reward delivery; the returned # positions index into reward_times, i.e. the deliveries that are manual. manual_water_times = np.asarray(manual_water_times) - manual_mask = np.zeros(reward_times.size, dtype=bool) if manual_water_times.size: manual_indices_in_reward_times = find_closest_timestamps(manual_water_times, reward_times) - manual_mask[manual_indices_in_reward_times] = True - annotated_rewards[manual_mask] = "manual" - - # A delivery on a trial reporting is_rewarded=False is not reward the animal - # received, so it is dropped. In practice these are all autowater: it is delivered - # before the response, so on a trial the animal answered the other way the water is - # never collected. Manual water keeps its precedence and is never dropped. - keep = ~(np.array(is_unrewarded, dtype=bool) & ~manual_mask) - return reward_times[keep], annotated_rewards[keep] + annotated_rewards[manual_indices_in_reward_times] = "manual" + + return annotated_rewards diff --git a/src/dynamic_foraging_processing/utils/trial_metadata.py b/src/dynamic_foraging_processing/utils/trial_metadata.py new file mode 100644 index 0000000..329e663 --- /dev/null +++ b/src/dynamic_foraging_processing/utils/trial_metadata.py @@ -0,0 +1,43 @@ +"""Helpers for reading a trial's block-based extra metadata.""" + +from aind_behavior_dynamic_foraging.task_logic.trial_generators.block_based_trial_generator import ( + BlockBasedTrialMetadata, +) +from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial + + +def get_bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: + """Return the block-based extra metadata naming a trial's free-water mechanism. + + ``trial.is_auto_reward_right`` is only the delivery *channel*: it says free + water was triggered and on which side, not what kind. Scheduled autowater and + the anti-bias water intervention share that channel and are told apart here, + by ``is_autowater`` and ``is_bias_water_intervention``. + (``is_bias_stage_intervention`` marks the anti-bias algorithm's other lever, + moving the lickspouts.) + + The field is schema-typed ``Any``, so it deserializes off the stream as a + plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` instance is + also accepted. When metadata or extra is missing (e.g. a non-block-based + generator), the model's all-``False`` default is returned, so a trial whose + mechanism the data does not record is reported as neither kind rather than + guessed at. + + Parameters + ---------- + trial : Trial + The per-trial task-logic model. + + Returns + ------- + BlockBasedTrialMetadata + The parsed extra metadata, or an all-``False`` default when absent + or unrecognized. + """ + metadata = trial.metadata + extra = metadata.extra if metadata is not None else None + if isinstance(extra, BlockBasedTrialMetadata): + return extra + if isinstance(extra, dict): + return BlockBasedTrialMetadata.model_validate(extra) + return BlockBasedTrialMetadata() From 1e03019aff29a7336e0b6d61050116d787fccc83 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 11:34:10 -0700 Subject: [PATCH 08/13] refactor: filter on whether trial was rewarded --- .../nwb/acquisition/acquisition_builder.py | 8 +-- .../processing/_trial_table.py | 55 +++++++++++++++---- .../utils/rewards.py | 43 ++++++++++----- 3 files changed, 77 insertions(+), 29 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index c4a7035..0a4d380 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -203,9 +203,9 @@ def _reward_delivery_series( Only valve-open events (``port_column`` is truthy) are reward deliveries; the ``data`` field annotates each as earned, manual, auto, or - anti-bias via :func:`get_reward_deliveries`. Every valve opening is - reported, so the series stays a complete record of the water delivered at - this port. + anti-bias via :func:`get_reward_deliveries`, which also drops deliveries + on trials that did not pay out, so the series reports reward rather than + every valve opening. Parameters ---------- @@ -237,7 +237,7 @@ def _reward_delivery_series( open_writes = writes[writes[port_column].fillna(False).astype(bool)] delivery_times = open_writes.index.to_numpy() manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy() - annotations = get_reward_deliveries( + delivery_times, annotations = get_reward_deliveries( delivery_times, trial_outcomes, manual_water_times, diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 123aece..5ec9e60 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -443,8 +443,14 @@ def _is_baited(trial: Trial, *, is_right: bool) -> bool: return trial.p_reward_left == 1 and auto in (None, True) @staticmethod - def _auto_water(trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool) -> int: - """Return whether scheduled autowater was delivered to the requested side. + def _auto_water( + trial: Trial, + outcome: TrialOutcome, + bias_metadata: BlockBasedTrialMetadata, + *, + is_right: bool, + ) -> int: + """Return whether scheduled autowater was rewarded on the requested side. ``is_auto_reward_right`` is the *delivery channel*: it says free water was triggered and to which side (``True`` right, ``False`` left, ``None`` no @@ -455,10 +461,17 @@ def _auto_water(trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_righ from the anti-bias algorithm is therefore ``0`` here, and is reported by ``anti_bias_left_water``/``anti_bias_right_water`` instead. + Free water on a trial that did not pay out (``is_rewarded`` is ``False``) + is ``0``, matching the reward-keyed reward-delivery series: downstream + analysis counts water that was reward, and free water fires at the go cue + whether or not the animal's own choice later paid out. + Parameters ---------- trial : Trial The per-trial task-logic model. + outcome : TrialOutcome + The trial's outcome, read for ``is_rewarded``. bias_metadata : BlockBasedTrialMetadata The trial's extra metadata (see ``_bias_metadata``). is_right : bool @@ -467,9 +480,10 @@ def _auto_water(trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_righ Returns ------- int - ``1`` when scheduled autowater targeted the requested side, else ``0``. + ``1`` when rewarded scheduled autowater targeted the requested side, + else ``0``. """ - if not bias_metadata.is_autowater: + if not bias_metadata.is_autowater or not outcome.is_rewarded: return 0 return int(trial.is_auto_reward_right is is_right) @@ -495,9 +509,13 @@ def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: @staticmethod def _anti_bias_water( - trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool + trial: Trial, + outcome: TrialOutcome, + bias_metadata: BlockBasedTrialMetadata, + *, + is_right: bool, ) -> bool: - """Return whether the anti-bias algorithm watered the requested side. + """Return whether the anti-bias algorithm's water was rewarded on this side. The anti-bias algorithm delivers its water intervention through the same auto-response channel as ordinary autowater (``is_auto_reward_right``: @@ -506,10 +524,18 @@ def _anti_bias_water( was a bias-water intervention *and* the auto-response was to the requested side. + Free water on a trial that did not pay out (``is_rewarded`` is ``False``) + is ``False``, matching :meth:`_auto_water` and the reward-keyed + reward-delivery series. The intervention still fired on those trials -- + it is triggered at the go cue regardless of the animal's later choice -- + so this column counts rewarded interventions, not every intervention. + Parameters ---------- trial : Trial The per-trial task-logic model. + outcome : TrialOutcome + The trial's outcome, read for ``is_rewarded``. bias_metadata : BlockBasedTrialMetadata The trial's extra metadata (see ``_bias_metadata``). is_right : bool @@ -518,9 +544,10 @@ def _anti_bias_water( Returns ------- bool - Whether an anti-bias water intervention targeted the requested side. + Whether a rewarded anti-bias water intervention targeted the + requested side. """ - if not bias_metadata.is_bias_water_intervention: + if not bias_metadata.is_bias_water_intervention or not outcome.is_rewarded: return False return trial.is_auto_reward_right is is_right @@ -906,10 +933,14 @@ def _build_row( reward_consumption_duration=trial.reward_consumption_duration, ITI_duration=trial.inter_trial_interval_duration, delay_duration=trial.quiescence_period_duration, - auto_waterL=self._auto_water(trial, bias_metadata, is_right=False), - auto_waterR=self._auto_water(trial, bias_metadata, is_right=True), - anti_bias_left_water=self._anti_bias_water(trial, bias_metadata, is_right=False), - anti_bias_right_water=self._anti_bias_water(trial, bias_metadata, is_right=True), + auto_waterL=self._auto_water(trial, outcome, bias_metadata, is_right=False), + auto_waterR=self._auto_water(trial, outcome, bias_metadata, is_right=True), + anti_bias_left_water=self._anti_bias_water( + trial, outcome, bias_metadata, is_right=False + ), + anti_bias_right_water=self._anti_bias_water( + trial, outcome, bias_metadata, is_right=True + ), anti_bias_lickspout_movement=self._anti_bias_lickspout_movement(trial, bias_metadata), **session, **lickspout, diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 7047ac4..a160053 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -65,8 +65,8 @@ def get_reward_deliveries( trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, response_times: np.ndarray, -) -> np.ndarray: - """Classify one lick port's reward deliveries by how the water was given. +) -> t.Tuple[np.ndarray, np.ndarray]: + """Get one lick port's reward deliveries, classified by how the water was given. Annotates the deliveries of a single lick port. Each delivery is classified as follows, with ``manual`` taking precedence because manual water is not @@ -88,11 +88,17 @@ def get_reward_deliveries( :func:`get_bias_metadata`), mirroring ``auto_waterL``/``auto_waterR`` and ``anti_bias_left_water``/``anti_bias_right_water`` in the trials table. - Every delivery is annotated and none is filtered out. A trial reporting - ``is_rewarded=False`` keeps its delivery: free water is triggered immediately - at the go cue and the trial then continues normally, so ``is_rewarded`` - reports the outcome of the animal's own choice -- a separate event from the - water being classified here. + Deliveries on a trial reporting ``is_rewarded=False`` are dropped rather than + annotated, so the series reports only water that counted as reward. In + practice these are all free water: it is triggered immediately at the go cue + and the trial then continues normally, so a trial whose own choice did not + pay out still carries the delivery. Manual water is experimenter-driven and + is never dropped. The surviving timestamps are returned alongside their + annotations so the two stay aligned. + + Note this makes the series reward-keyed rather than a complete record of the + hardware's valve openings: free water delivered on an unrewarded trial is + real water the animal received, and it is excluded here. Deliveries are matched to trials by the ``Response`` software-event timestamp: each delivery takes the annotation of the trial whose response is @@ -123,8 +129,11 @@ def get_reward_deliveries( Returns ------- numpy.ndarray - Array of the same shape as ``reward_delivery_times`` whose entries are - ``"earned"``, ``"auto"``, ``"anti_bias"``, or ``"manual"``. + The retained reward-delivery timestamps: ``reward_delivery_times`` less + the deliveries on unrewarded trials. + numpy.ndarray + The matching annotations, one per retained timestamp, each ``"earned"``, + ``"auto"``, ``"anti_bias"``, or ``"manual"``. Raises ------ @@ -141,16 +150,18 @@ def get_reward_deliveries( reward_times = np.asarray(reward_delivery_times) if reward_times.size == 0: - return np.array([], dtype=object) + return reward_times, np.array([], dtype=object) # Annotate each delivery from its originating trial: query with reward_times so we # get one trial position per reward delivery. trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) annotated_rewards = [] + is_unrewarded = [] for trial_index in trial_indices_in_reward_times: outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"]) annotated_rewards.append(_free_water_label(outcome.trial if outcome is not None else None)) + is_unrewarded.append(outcome is not None and not outcome.is_rewarded) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" # and "earned" entries would be too narrow to hold "anti_bias" and would truncate it. @@ -161,8 +172,14 @@ def get_reward_deliveries( # manual-water software event to its closest reward delivery; the returned # positions index into reward_times, i.e. the deliveries that are manual. manual_water_times = np.asarray(manual_water_times) + manual_mask = np.zeros(reward_times.size, dtype=bool) if manual_water_times.size: manual_indices_in_reward_times = find_closest_timestamps(manual_water_times, reward_times) - annotated_rewards[manual_indices_in_reward_times] = "manual" - - return annotated_rewards + manual_mask[manual_indices_in_reward_times] = True + annotated_rewards[manual_mask] = "manual" + + # Downstream analysis is keyed on reward, so a delivery whose trial did not pay out + # is excluded. Manual water is experimenter-driven, unrelated to the trial's + # outcome, and keeps its delivery. + keep = ~(np.array(is_unrewarded, dtype=bool) & ~manual_mask) + return reward_times[keep], annotated_rewards[keep] From 9b03efd68d788fd6db24423363bb47b3ea61a129 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 14:52:08 -0700 Subject: [PATCH 09/13] refactor: use same logic across building both acquisition and trial table for auto --- .../nwb/acquisition/acquisition_builder.py | 10 +-- .../processing/_trial_table.py | 64 +++++++++---------- .../utils/rewards.py | 42 ++++++------ .../utils/trial_metadata.py | 43 ------------- 4 files changed, 54 insertions(+), 105 deletions(-) delete mode 100644 src/dynamic_foraging_processing/utils/trial_metadata.py diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index 0a4d380..c46acef 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -202,10 +202,10 @@ def _reward_delivery_series( """Build one lick port's reward-delivery series with reward annotations. Only valve-open events (``port_column`` is truthy) are reward - deliveries; the ``data`` field annotates each as earned, manual, auto, or - anti-bias via :func:`get_reward_deliveries`, which also drops deliveries - on trials that did not pay out, so the series reports reward rather than - every valve opening. + deliveries; the ``data`` field annotates each as earned, manual, or auto + via :func:`get_reward_deliveries`, which also drops deliveries on trials + that did not pay out, so the series reports reward rather than every + valve opening. Parameters ---------- @@ -250,7 +250,7 @@ def _reward_delivery_series( unit="second", description=( f"The reward delivery time of the {side_label} lick port. The data field " - "annotates whether the reward was earned, manual, auto, or anti_bias" + "annotates whether the reward was earned, manual, or auto" ), ) diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 5ec9e60..15a4547 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -25,7 +25,6 @@ from contraqctor.contract import Dataset from dynamic_foraging_processing.processing.models import TrialConfig -from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata logger = logging.getLogger(__name__) @@ -443,28 +442,19 @@ def _is_baited(trial: Trial, *, is_right: bool) -> bool: return trial.p_reward_left == 1 and auto in (None, True) @staticmethod - def _auto_water( - trial: Trial, - outcome: TrialOutcome, - bias_metadata: BlockBasedTrialMetadata, - *, - is_right: bool, - ) -> int: - """Return whether scheduled autowater was rewarded on the requested side. - - ``is_auto_reward_right`` is the *delivery channel*: it says free water was - triggered and to which side (``True`` right, ``False`` left, ``None`` no - free water), but not what kind. Autowater and the anti-bias water - intervention share that channel and are told apart only by the metadata - flags, so this reads ``is_autowater`` and takes the side from the channel - -- the mirror of :meth:`_anti_bias_water`. A trial whose free water came - from the anti-bias algorithm is therefore ``0`` here, and is reported by - ``anti_bias_left_water``/``anti_bias_right_water`` instead. + def _auto_water(trial: Trial, outcome: TrialOutcome, *, is_right: bool) -> int: + """Return whether autowater was rewarded on the requested side. - Free water on a trial that did not pay out (``is_rewarded`` is ``False``) - is ``0``, matching the reward-keyed reward-delivery series: downstream - analysis counts water that was reward, and free water fires at the go cue - whether or not the animal's own choice later paid out. + ``is_auto_reward_right`` triggers an immediate reward to one side + (``True`` right, ``False`` left, ``None`` no autowater). The anti-bias + water intervention is itself autowater delivered through this channel, so + every auto-triggered reward counts here; ``anti_bias_left_water`` / + ``anti_bias_right_water`` mark the subset the anti-bias algorithm drove. + + Autowater on a trial that did not pay out (``is_rewarded`` is ``False``) + is ``0``, so this column matches the reward-keyed reward-delivery series: + the water fires at the go cue whether or not the animal's own choice + later pays out, and the series reports only water that was reward. Parameters ---------- @@ -472,27 +462,29 @@ def _auto_water( The per-trial task-logic model. outcome : TrialOutcome The trial's outcome, read for ``is_rewarded``. - bias_metadata : BlockBasedTrialMetadata - The trial's extra metadata (see ``_bias_metadata``). is_right : bool ``True`` for the right port, ``False`` for the left port. Returns ------- int - ``1`` when rewarded scheduled autowater targeted the requested side, - else ``0``. + ``1`` when rewarded autowater targeted the requested side, else ``0``. """ - if not bias_metadata.is_autowater or not outcome.is_rewarded: + if trial.is_auto_reward_right is None or not outcome.is_rewarded: return 0 return int(trial.is_auto_reward_right is is_right) @staticmethod def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: - """Return the block-based extra metadata carrying the free-water flags. + """Return the block-based extra metadata carrying the anti-bias flags. - Thin wrapper over :func:`get_bias_metadata`, shared with the reward - annotation so both read the autowater and anti-bias flags identically. + The anti-bias flags (``is_bias_water_intervention``, + ``is_bias_stage_intervention``) live on ``trial.metadata.extra``. That + field is schema-typed ``Any``, so it deserializes off the stream as a + plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` + instance is also accepted. When metadata or extra is missing (e.g. a + non-block-based generator), the model's all-``False`` default is returned + so the anti-bias columns are simply inert. Parameters ---------- @@ -505,7 +497,13 @@ def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: The parsed extra metadata, or an all-``False`` default when absent or unrecognized. """ - return get_bias_metadata(trial) + metadata = trial.metadata + extra = metadata.extra if metadata is not None else None + if isinstance(extra, BlockBasedTrialMetadata): + return extra + if isinstance(extra, dict): + return BlockBasedTrialMetadata.model_validate(extra) + return BlockBasedTrialMetadata() @staticmethod def _anti_bias_water( @@ -933,8 +931,8 @@ def _build_row( reward_consumption_duration=trial.reward_consumption_duration, ITI_duration=trial.inter_trial_interval_duration, delay_duration=trial.quiescence_period_duration, - auto_waterL=self._auto_water(trial, outcome, bias_metadata, is_right=False), - auto_waterR=self._auto_water(trial, outcome, bias_metadata, is_right=True), + auto_waterL=self._auto_water(trial, outcome, is_right=False), + auto_waterR=self._auto_water(trial, outcome, is_right=True), anti_bias_left_water=self._anti_bias_water( trial, outcome, bias_metadata, is_right=False ), diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index a160053..ee3b75d 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -7,7 +7,6 @@ from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome from dynamic_foraging_processing.utils.timestamps import find_closest_timestamps -from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: @@ -32,12 +31,13 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: def _free_water_label(trial: t.Optional[Trial]) -> str: - """Classify a delivery's trial as ``auto``, ``anti_bias``, or ``earned``. + """Classify a delivery's trial as ``auto`` or ``earned``. - ``is_auto_reward_right`` marks only that free water was triggered and on - which side; scheduled autowater and the anti-bias intervention share that - channel and are told apart by the block-based metadata flags. A trial with no - free water, or free water flagged as neither mechanism, is ``earned``. + ``is_auto_reward_right`` triggers an immediate reward to one side; the + anti-bias water intervention is itself autowater delivered through that + channel, so any auto-triggered reward is ``auto``. This is the same condition + ``auto_waterL``/``auto_waterR`` use in the trials table, so the two agree + trial for trial. Parameters ---------- @@ -48,16 +48,11 @@ def _free_water_label(trial: t.Optional[Trial]) -> str: Returns ------- str - ``"auto"``, ``"anti_bias"``, or ``"earned"``. + ``"auto"`` when the trial auto-triggered a reward, else ``"earned"``. """ if trial is None or trial.is_auto_reward_right is None: return "earned" - metadata = get_bias_metadata(trial) - if metadata.is_bias_water_intervention: - return "anti_bias" - if metadata.is_autowater: - return "auto" - return "earned" + return "auto" def get_reward_deliveries( @@ -76,17 +71,16 @@ def get_reward_deliveries( ``GiveManualWater`` software event for this port. The software-event timestamps are correlated to the reward-delivery timestamps with :func:`find_closest_timestamps`. - - ``anti_bias`` -- otherwise, when the trial's free water came from the - anti-bias algorithm (``is_bias_water_intervention``). - - ``auto`` -- otherwise, when the trial's free water was scheduled autowater - (``is_autowater``). + - ``auto`` -- otherwise, when the trial auto-triggered a reward + (``is_auto_reward_right is not None``). The anti-bias water intervention is + autowater delivered through that same channel, so it is ``auto`` too; the + trials table's ``anti_bias_left_water``/``anti_bias_right_water`` mark + which of these the anti-bias algorithm drove. - ``earned`` -- otherwise: water the animal worked for. - ``is_auto_reward_right`` is only the delivery *channel* -- it says free water - was triggered and on which side, not what kind -- so the ``auto`` versus - ``anti_bias`` split comes from the block-based metadata flags (see - :func:`get_bias_metadata`), mirroring ``auto_waterL``/``auto_waterR`` and - ``anti_bias_left_water``/``anti_bias_right_water`` in the trials table. + The ``auto`` condition is the one ``auto_waterL``/``auto_waterR`` use in the + trials table, and the drop below matches those columns' ``is_rewarded`` gate, + so the ``auto`` count here equals the trials table's autowater count. Deliveries on a trial reporting ``is_rewarded=False`` are dropped rather than annotated, so the series reports only water that counted as reward. In @@ -133,7 +127,7 @@ def get_reward_deliveries( the deliveries on unrewarded trials. numpy.ndarray The matching annotations, one per retained timestamp, each ``"earned"``, - ``"auto"``, ``"anti_bias"``, or ``"manual"``. + ``"auto"``, or ``"manual"``. Raises ------ @@ -164,7 +158,7 @@ def get_reward_deliveries( is_unrewarded.append(outcome is not None and not outcome.is_rewarded) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" - # and "earned" entries would be too narrow to hold "anti_bias" and would truncate it. + # and "earned" entries would be too narrow to hold "manual" and would truncate it. annotated_rewards = np.array(annotated_rewards, dtype=object) # Manual water is independent of trials (multiple can occur within a trial) and diff --git a/src/dynamic_foraging_processing/utils/trial_metadata.py b/src/dynamic_foraging_processing/utils/trial_metadata.py deleted file mode 100644 index 329e663..0000000 --- a/src/dynamic_foraging_processing/utils/trial_metadata.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Helpers for reading a trial's block-based extra metadata.""" - -from aind_behavior_dynamic_foraging.task_logic.trial_generators.block_based_trial_generator import ( - BlockBasedTrialMetadata, -) -from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial - - -def get_bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: - """Return the block-based extra metadata naming a trial's free-water mechanism. - - ``trial.is_auto_reward_right`` is only the delivery *channel*: it says free - water was triggered and on which side, not what kind. Scheduled autowater and - the anti-bias water intervention share that channel and are told apart here, - by ``is_autowater`` and ``is_bias_water_intervention``. - (``is_bias_stage_intervention`` marks the anti-bias algorithm's other lever, - moving the lickspouts.) - - The field is schema-typed ``Any``, so it deserializes off the stream as a - plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` instance is - also accepted. When metadata or extra is missing (e.g. a non-block-based - generator), the model's all-``False`` default is returned, so a trial whose - mechanism the data does not record is reported as neither kind rather than - guessed at. - - Parameters - ---------- - trial : Trial - The per-trial task-logic model. - - Returns - ------- - BlockBasedTrialMetadata - The parsed extra metadata, or an all-``False`` default when absent - or unrecognized. - """ - metadata = trial.metadata - extra = metadata.extra if metadata is not None else None - if isinstance(extra, BlockBasedTrialMetadata): - return extra - if isinstance(extra, dict): - return BlockBasedTrialMetadata.model_validate(extra) - return BlockBasedTrialMetadata() From a28cf53996c280229e635fe4e75791679402ac27 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 14:52:18 -0700 Subject: [PATCH 10/13] test: update tests --- tests/test_processing/test_trial_table.py | 85 +++++++++++++++++++---- tests/test_utils/test_rewards.py | 34 +++++++-- 2 files changed, 101 insertions(+), 18 deletions(-) diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index b213ba1..13a0619 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -690,16 +690,49 @@ def test_rewarded_history_false_on_every_auto_reward_trial(): 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( + outcome = TrialOutcome.model_validate( _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=True) - ).trial - assert TrialTableBuilder._auto_water(trial, is_right=True) == 1 - assert TrialTableBuilder._auto_water(trial, is_right=False) == 0 + ) + assert TrialTableBuilder._auto_water(outcome.trial, outcome, is_right=True) == 1 + assert TrialTableBuilder._auto_water(outcome.trial, outcome, is_right=False) == 0 # No auto-response counts as no autowater (0). no_auto = TrialOutcome.model_validate( _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=None) - ).trial - assert TrialTableBuilder._auto_water(no_auto, is_right=True) == 0 + ) + assert TrialTableBuilder._auto_water(no_auto.trial, no_auto, is_right=True) == 0 + + +def test_auto_water_excludes_unrewarded_trials(): + """Autowater on a trial that did not pay out is ``0``. + + Free water fires at the go cue whether or not the animal's own choice later + pays out, so this keeps the column matching the reward-keyed reward-delivery + series, which drops those deliveries. + """ + unrewarded = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=False, is_rewarded=False, auto=True) + ) + assert TrialTableBuilder._auto_water(unrewarded.trial, unrewarded, is_right=True) == 0 + + +def test_auto_water_counts_anti_bias_water(): + """Anti-bias water is autowater delivered through the auto-response channel. + + The anti-bias intervention shares ``is_auto_reward_right`` with scheduled + autowater, so it counts here as well; ``anti_bias_right_water`` marks which + of these the algorithm drove. + """ + bias_water = TrialOutcome.model_validate( + _outcome( + 1.0, + 1.0, + is_right_choice=True, + is_rewarded=True, + auto=True, + is_bias_water_intervention=True, + ) + ) + assert TrialTableBuilder._auto_water(bias_water.trial, bias_water, is_right=True) == 1 def test_bias_metadata_parses_dict_model_and_default(): @@ -750,18 +783,44 @@ def test_anti_bias_water_gated_on_intervention_flag_and_side(): auto=True, is_bias_water_intervention=True, ) - ).trial - meta = TrialTableBuilder._bias_metadata(right) - assert TrialTableBuilder._anti_bias_water(right, meta, is_right=True) is True - assert TrialTableBuilder._anti_bias_water(right, meta, is_right=False) is False + ) + meta = TrialTableBuilder._bias_metadata(right.trial) + assert TrialTableBuilder._anti_bias_water(right.trial, right, meta, is_right=True) is True + assert TrialTableBuilder._anti_bias_water(right.trial, right, meta, is_right=False) is False # Auto-response to the left without the bias flag is ordinary autowater, not # an anti-bias intervention. autowater = TrialOutcome.model_validate( _outcome(1.0, 1.0, is_right_choice=False, is_rewarded=True, auto=False) - ).trial - auto_meta = TrialTableBuilder._bias_metadata(autowater) - assert TrialTableBuilder._anti_bias_water(autowater, auto_meta, is_right=False) is False + ) + auto_meta = TrialTableBuilder._bias_metadata(autowater.trial) + assert ( + TrialTableBuilder._anti_bias_water(autowater.trial, autowater, auto_meta, is_right=False) + is False + ) + + +def test_anti_bias_water_excludes_unrewarded_trials(): + """An intervention on a trial that did not pay out is ``False``. + + Keeps the column consistent with ``auto_waterL``/``auto_waterR``, of which it + marks a subset; the intervention still fired, but the trial reports no reward. + """ + unrewarded = TrialOutcome.model_validate( + _outcome( + 1.0, + 1.0, + is_right_choice=False, + is_rewarded=False, + auto=True, + is_bias_water_intervention=True, + ) + ) + meta = TrialTableBuilder._bias_metadata(unrewarded.trial) + assert ( + TrialTableBuilder._anti_bias_water(unrewarded.trial, unrewarded, meta, is_right=True) + is False + ) def test_anti_bias_lickspout_movement_gated_on_stage_flag(): diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index 38bda02..545d6c6 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -65,6 +65,28 @@ def test_get_reward_deliveries_marks_auto_response_trials_as_auto(): np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) +def test_get_reward_deliveries_marks_anti_bias_water_as_auto(): + """Anti-bias water is autowater, so it annotates as ``auto``. + + The anti-bias intervention is delivered through the same + ``is_auto_reward_right`` channel as scheduled autowater, so the annotation + does not distinguish them; ``anti_bias_left_water``/``anti_bias_right_water`` + in the trials table mark which deliveries the algorithm drove. + """ + reward_times = np.array([0.15]) + response_times = np.array([0.1]) + payload = _outcome_payload(True) + payload["trial"]["metadata"] = {"extra": {"is_bias_water_intervention": True}} + trial_outcome_df = pd.DataFrame({"data": [payload]}, index=pd.Index([1.1], name="time")) + + times, annotations = get_reward_deliveries( + reward_times, trial_outcome_df, np.array([]), response_times + ) + + np.testing.assert_array_equal(times, reward_times) + np.testing.assert_array_equal(annotations, np.array(["auto"])) + + def test_get_reward_deliveries_matches_closest_response_time(): """Each delivery takes the annotation of the trial whose response is closest. @@ -85,11 +107,12 @@ def test_get_reward_deliveries_matches_closest_response_time(): np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -def test_get_reward_deliveries_drops_uncollected_auto_water(): +def test_get_reward_deliveries_drops_auto_water_on_unrewarded_trials(): """Autowater on a trial reporting ``is_rewarded=False`` is dropped, not annotated. - The water is delivered before the response, so a trial the animal answered - the other way leaves it uncollected; it is not reward the animal received. + The water is delivered at the go cue and the trial then continues normally, + so a trial whose own choice did not pay out still carries the delivery. The + series is reward-keyed, so those deliveries are excluded. """ reward_times = np.array([0.15, 0.42, 0.95]) response_times = np.array([0.1, 0.4, 0.9]) @@ -110,8 +133,9 @@ def test_get_reward_deliveries_drops_uncollected_auto_water(): def test_get_reward_deliveries_drops_any_delivery_on_an_unrewarded_trial(): """The drop rule is ``is_rewarded=False``, not autowater specifically. - Autowater is the only case seen in practice, but a delivery on any trial - reporting no reward is water the animal did not receive. + Autowater is the only case seen in practice, but the condition is the trial's + reward outcome, so any delivery on a trial that did not pay out is excluded + regardless of what triggered it. """ reward_times = np.array([0.15]) response_times = np.array([0.1]) From 038a553516a0372eff288d09236295a7191fdf7b Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 18:20:55 -0700 Subject: [PATCH 11/13] feat: add column for auto water offered to distinguish from auto water for rewards --- .../processing/_trial_table.py | 64 ++++++++++++------- .../processing/models/trial_config.py | 31 +++++++-- tests/test_processing/test_trial_table.py | 47 ++++++++++---- 3 files changed, 102 insertions(+), 40 deletions(-) diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 15a4547..98c8fce 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -474,6 +474,34 @@ def _auto_water(trial: Trial, outcome: TrialOutcome, *, is_right: bool) -> int: return 0 return int(trial.is_auto_reward_right is is_right) + @staticmethod + def _auto_water_offered(trial: Trial, *, is_right: bool) -> int: + """Return whether autowater was *offered* to the requested side. + + The ungated counterpart of :meth:`_auto_water`: ``1`` whenever the trial + auto-triggered a reward to this side, whether or not the trial went on to + pay out. This is the legacy ``dynamic-foraging-task`` meaning of + ``auto_waterL``/``auto_waterR`` (its ``B_AutoWaterTrial``, which that GUI + also passes to foraging efficiency as ``autowater_offered``), kept so the + trial table still records every autowater delivery while ``auto_water*`` + stays reward-keyed to match the reward-delivery series. + + Parameters + ---------- + trial : Trial + The per-trial task-logic model. + is_right : bool + ``True`` for the right port, ``False`` for the left port. + + Returns + ------- + int + ``1`` when autowater was offered to the requested side, else ``0``. + """ + if trial.is_auto_reward_right is None: + return 0 + return int(trial.is_auto_reward_right is is_right) + @staticmethod def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: """Return the block-based extra metadata carrying the anti-bias flags. @@ -507,13 +535,9 @@ def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: @staticmethod def _anti_bias_water( - trial: Trial, - outcome: TrialOutcome, - bias_metadata: BlockBasedTrialMetadata, - *, - is_right: bool, + trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool ) -> bool: - """Return whether the anti-bias algorithm's water was rewarded on this side. + """Return whether the anti-bias algorithm watered the requested side. The anti-bias algorithm delivers its water intervention through the same auto-response channel as ordinary autowater (``is_auto_reward_right``: @@ -522,18 +546,17 @@ def _anti_bias_water( was a bias-water intervention *and* the auto-response was to the requested side. - Free water on a trial that did not pay out (``is_rewarded`` is ``False``) - is ``False``, matching :meth:`_auto_water` and the reward-keyed - reward-delivery series. The intervention still fired on those trials -- - it is triggered at the go cue regardless of the animal's later choice -- - so this column counts rewarded interventions, not every intervention. + This records what the *algorithm* did, so it is not gated on + ``is_rewarded``: the intervention fires at the go cue regardless of how + the animal's own choice later resolves. It is therefore not a subset of + ``auto_waterL``/``auto_waterR``, which count only rewarded autowater -- + an intervention on a trial that did not pay out appears here and not + there. Parameters ---------- trial : Trial The per-trial task-logic model. - outcome : TrialOutcome - The trial's outcome, read for ``is_rewarded``. bias_metadata : BlockBasedTrialMetadata The trial's extra metadata (see ``_bias_metadata``). is_right : bool @@ -542,10 +565,9 @@ def _anti_bias_water( Returns ------- bool - Whether a rewarded anti-bias water intervention targeted the - requested side. + Whether an anti-bias water intervention targeted the requested side. """ - if not bias_metadata.is_bias_water_intervention or not outcome.is_rewarded: + if not bias_metadata.is_bias_water_intervention: return False return trial.is_auto_reward_right is is_right @@ -933,12 +955,10 @@ def _build_row( delay_duration=trial.quiescence_period_duration, auto_waterL=self._auto_water(trial, outcome, is_right=False), auto_waterR=self._auto_water(trial, outcome, is_right=True), - anti_bias_left_water=self._anti_bias_water( - trial, outcome, bias_metadata, is_right=False - ), - anti_bias_right_water=self._anti_bias_water( - trial, outcome, bias_metadata, is_right=True - ), + auto_water_offeredL=self._auto_water_offered(trial, is_right=False), + auto_water_offeredR=self._auto_water_offered(trial, is_right=True), + anti_bias_left_water=self._anti_bias_water(trial, bias_metadata, is_right=False), + anti_bias_right_water=self._anti_bias_water(trial, bias_metadata, is_right=True), anti_bias_lickspout_movement=self._anti_bias_lickspout_movement(trial, bias_metadata), **session, **lickspout, diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index 22e0e96..57a4211 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -189,20 +189,43 @@ class TrialConfig(BaseModel): ) # --- auto_waterL/R (autowater per-side; autoTrain curriculum fields out of scope) --- - auto_waterL: int = Field(default=0, description="Autowater given at Left") - auto_waterR: int = Field(default=0, description="Autowater given at Right") + auto_waterL: int = Field( + default=0, + description=( + "Rewarded autowater given at Left: 1 only when the trial auto-triggered a reward to the left (is_auto_reward_right is False) AND the trial was rewarded (is_rewarded). Autowater on a trial that did not pay out is 0, so this column counts autowater that was reward and matches the left reward-delivery series. The anti-bias water intervention uses the same autowater channel and is included here; anti_bias_left_water reports it ungated." + ), + ) + auto_waterR: int = Field( + default=0, + description=( + "Rewarded autowater given at Right: 1 only when the trial auto-triggered a reward to the right (is_auto_reward_right is True) AND the trial was rewarded (is_rewarded). Autowater on a trial that did not pay out is 0, so this column counts autowater that was reward and matches the right reward-delivery series. The anti-bias water intervention uses the same autowater channel and is included here; anti_bias_right_water reports it ungated." + ), + ) + + auto_water_offeredL: int = Field( + default=0, + description=( + "Autowater offered at Left: 1 whenever the trial auto-triggered a reward to the left (is_auto_reward_right is False), whether or not the trial paid out. This is the legacy dynamic-foraging-task meaning of auto_waterL, kept so no delivery is lost from the trial table; auto_waterL is the reward-keyed subset that matches the reward-delivery series." + ), + ) + auto_water_offeredR: int = Field( + default=0, + description=( + "Autowater offered at Right: 1 whenever the trial auto-triggered a reward to the right (is_auto_reward_right is True), whether or not the trial paid out. This is the legacy dynamic-foraging-task meaning of auto_waterR, kept so no delivery is lost from the trial table; auto_waterR is the reward-keyed subset that matches the reward-delivery series." + ), + ) # --- anti_bias (interventions the anti-bias algorithm applies) --- anti_bias_left_water: bool = Field( default=False, description=( - "Whether the anti-bias algorithm delivered a water intervention to the left lickport on this trial." + "Whether the anti-bias algorithm delivered a water intervention to the left lickport on this trial. Records what the algorithm did, so unlike auto_waterL this is NOT conditioned on is_rewarded: the intervention fires at the go cue regardless of how the animal's own choice resolves. It can therefore be True where auto_waterL is 0." ), ) anti_bias_right_water: bool = Field( default=False, description=( - "Whether the anti-bias algorithm delivered a water intervention to the right lickport on this trial." + "Whether the anti-bias algorithm delivered a water intervention to the right lickport on this trial. Records what the algorithm did, so unlike auto_waterR this is NOT conditioned on is_rewarded: the intervention fires at the go cue regardless of how the animal's own choice resolves. It can therefore be True where auto_waterR is 0." ), ) anti_bias_lickspout_movement: float = Field( diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index 13a0619..f422b8b 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -715,6 +715,27 @@ def test_auto_water_excludes_unrewarded_trials(): assert TrialTableBuilder._auto_water(unrewarded.trial, unrewarded, is_right=True) == 0 +def test_auto_water_offered_is_ungated(): + """``auto_water_offered*`` records every autowater, paid out or not. + + This is the legacy ``dynamic-foraging-task`` meaning of ``auto_water*``, kept + so the trial table loses no delivery now that ``auto_water*`` is reward-keyed. + """ + unrewarded = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=False, is_rewarded=False, auto=True) + ) + assert TrialTableBuilder._auto_water_offered(unrewarded.trial, is_right=True) == 1 + assert TrialTableBuilder._auto_water_offered(unrewarded.trial, is_right=False) == 0 + # The reward-keyed column is 0 on the very same trial. + assert TrialTableBuilder._auto_water(unrewarded.trial, unrewarded, is_right=True) == 0 + + # No autowater at all is 0 on both sides. + no_auto = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=None) + ) + assert TrialTableBuilder._auto_water_offered(no_auto.trial, is_right=True) == 0 + + def test_auto_water_counts_anti_bias_water(): """Anti-bias water is autowater delivered through the auto-response channel. @@ -785,8 +806,8 @@ def test_anti_bias_water_gated_on_intervention_flag_and_side(): ) ) meta = TrialTableBuilder._bias_metadata(right.trial) - assert TrialTableBuilder._anti_bias_water(right.trial, right, meta, is_right=True) is True - assert TrialTableBuilder._anti_bias_water(right.trial, right, meta, is_right=False) is False + assert TrialTableBuilder._anti_bias_water(right.trial, meta, is_right=True) is True + assert TrialTableBuilder._anti_bias_water(right.trial, meta, is_right=False) is False # Auto-response to the left without the bias flag is ordinary autowater, not # an anti-bias intervention. @@ -794,17 +815,16 @@ def test_anti_bias_water_gated_on_intervention_flag_and_side(): _outcome(1.0, 1.0, is_right_choice=False, is_rewarded=True, auto=False) ) auto_meta = TrialTableBuilder._bias_metadata(autowater.trial) - assert ( - TrialTableBuilder._anti_bias_water(autowater.trial, autowater, auto_meta, is_right=False) - is False - ) + assert TrialTableBuilder._anti_bias_water(autowater.trial, auto_meta, is_right=False) is False -def test_anti_bias_water_excludes_unrewarded_trials(): - """An intervention on a trial that did not pay out is ``False``. +def test_anti_bias_water_includes_unrewarded_trials(): + """An intervention on a trial that did not pay out still counts. - Keeps the column consistent with ``auto_waterL``/``auto_waterR``, of which it - marks a subset; the intervention still fired, but the trial reports no reward. + The column records what the anti-bias algorithm did, and the intervention + fires at the go cue regardless of how the animal's own choice resolves. This + is deliberately *not* gated on ``is_rewarded``, so it is not a subset of + ``auto_waterL``/``auto_waterR``. """ unrewarded = TrialOutcome.model_validate( _outcome( @@ -817,10 +837,9 @@ def test_anti_bias_water_excludes_unrewarded_trials(): ) ) meta = TrialTableBuilder._bias_metadata(unrewarded.trial) - assert ( - TrialTableBuilder._anti_bias_water(unrewarded.trial, unrewarded, meta, is_right=True) - is False - ) + assert TrialTableBuilder._anti_bias_water(unrewarded.trial, meta, is_right=True) is True + # The same trial contributes no autowater, which *is* reward-gated. + assert TrialTableBuilder._auto_water(unrewarded.trial, unrewarded, is_right=True) == 0 def test_anti_bias_lickspout_movement_gated_on_stage_flag(): From 53131c773869699d0a555e881754eb0620a704e3 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 18:21:05 -0700 Subject: [PATCH 12/13] docs: update docs --- docs/trials_table_mapping.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index c69f3db..afe3fd7 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -61,8 +61,9 @@ Columns are grouped by the raw source they map from. | Trials column | Mapping | | --- | --- | -| `auto_waterL` / `auto_waterR` | From `is_auto_reward_right`. `1` on the auto-responded side; `0` on the other side, when there was no auto-response (`None`), or when the trial is missing. | -| `anti_bias_left_water` / `anti_bias_right_water` | Boolean. `True` when the anti-bias algorithm delivered a water intervention to that side — i.e. `trial.metadata.extra.is_bias_water_intervention` is `True` **and** `is_auto_reward_right` points to that side (`False` → left, `True` → right). The anti-bias water uses the same auto-response channel as ordinary autowater, so the `is_bias_water_intervention` flag is what distinguishes it. `False` otherwise. | +| `auto_waterL` / `auto_waterR` | **Rewarded** autowater only: `1` when `is_auto_reward_right` points to that side **and** `is_rewarded` is `True`. `0` on the other side, when there was no auto-response (`None`), when the trial delivered autowater but did not pay out, or when the trial is missing. `is_auto_reward_right` is the delivery channel — it says free water was triggered and to which side, not what kind — and the anti-bias water intervention is autowater delivered through it, so anti-bias water counts here too. Note this is narrower than the legacy `dynamic-foraging-task` column of the same name, which was ungated ("Autowater given at Left", straight from `B_AutoWaterTrial`). | +| `auto_water_offeredL` / `auto_water_offeredR` | Ungated autowater: `1` whenever `is_auto_reward_right` points to that side, whether or not the trial paid out. This is the legacy `dynamic-foraging-task` meaning of `auto_waterL` / `auto_waterR` (`B_AutoWaterTrial`, which that GUI also passes to foraging efficiency as `autowater_offered`), kept so the trial table records every autowater delivery. `auto_waterL` / `auto_waterR` are the reward-keyed subset. | +| `anti_bias_left_water` / `anti_bias_right_water` | Boolean. `True` when the anti-bias algorithm delivered a water intervention to that side — i.e. `trial.metadata.extra.is_bias_water_intervention` is `True` **and** `is_auto_reward_right` points to that side (`False` → left, `True` → right). The anti-bias water uses the same auto-response channel as ordinary autowater, so the `is_bias_water_intervention` flag is what distinguishes it. `False` otherwise. Deliberately **not** gated on `is_rewarded`: these columns record what the anti-bias algorithm did, and the intervention fires at the go cue regardless of how the animal's own choice resolves, so they can be `True` where `auto_waterL` / `auto_waterR` is `0`. | | `anti_bias_lickspout_movement` | Signed horizontal displacement (mm, positive is rightward) the anti-bias algorithm moved the lickspouts on this trial: `trial.lickspout_offset_delta` when `trial.metadata.extra.is_bias_stage_intervention` is `True`, else `0.0`. | | `bait_left` / `bait_right` | Boolean. `bait_right` is `True` if `p_reward_right == 1` and `is_auto_reward_right` is `None` or `False`. `bait_left` is `True` if `p_reward_left == 1` and `is_auto_reward_right` is `None` or `True`. | | `response_duration` | `response_deadline_duration`. | @@ -191,3 +192,7 @@ These were mapped during exploration but are no longer in scope: | 2026-08-06 | Confirmed and documented that the legacy `delay_*` columns describe the acquisition software's **quiescence period**: `delay_start_time` is the `QuiescentPeriod` timestamp (always equal to the new `quiescent_start_time`) and `delay_duration` / `delay_beta` / `delay_min` / `delay_max` summarize `quiescence_period_duration`. `delay_duration` is the *configured* duration — each lick restarts the quiescent period, so the realized `quiescent_stop_time - quiescent_start_time` can be longer. Column descriptions updated accordingly. | | 2026-08-12 | `rewarded_historyL` / `rewarded_historyR` now record **earned** reward only: an auto-reward trial (`is_auto_reward_right` set to either side) is `False` on *both* sides, since `TrialOutcome.is_rewarded` is `True` for autowater too and that water is already reported by `auto_waterL` / `auto_waterR`. This matches the `earned` / `automatic` split used for the NWB reward-delivery annotations. | | 2026-08-12 | `min_reward_each_block` is now `0` rather than `NULL` when the trial generator exposes no `min_block_reward` — no per-block minimum is a floor of zero, not an unknown. The column is non-nullable (`float`, default `0`). | +| 2026-08-17 | Reward-delivery annotations now match each delivery to its trial by the `Response` software-event timestamp rather than the `TrialOutcome` timestamp. `TrialOutcome` fires at the *end* of a trial (after the reward-consumption and ITI periods), so a delivery could land nearer the *previous* trial's outcome and inherit its `is_auto_reward_right`, flipping `earned` and `auto`. The valve opens within milliseconds of the response, so the response anchors the delivery to its own trial. | +| 2026-08-17 | `auto_waterL` / `auto_waterR` are now gated on `is_rewarded`: autowater on a trial that did not pay out is `0`. Free water is triggered immediately at the go cue and the trial then "continues normally", so `is_rewarded` reports the outcome of the animal's own choice — a separate event from the autowater. Gating keeps these columns equal to the `auto` count in the reward-delivery series, which drops the same deliveries, and makes the retained delivery total equal the metadata mapper's `sum(is_rewarded)`. This is narrower than the legacy ungated column of the same name. | +| 2026-08-17 | `anti_bias_left_water` / `anti_bias_right_water` are explicitly **not** gated on `is_rewarded`, unlike `auto_water*`: they record what the anti-bias algorithm did, and the intervention fires regardless of the trial's outcome. They are therefore no longer a subset of `auto_water*` — an intervention on a trial that did not pay out appears in the anti-bias column and not in the autowater column. | +| 2026-08-17 | Added `auto_water_offeredL` / `auto_water_offeredR`: autowater ungated by `is_rewarded`, i.e. the legacy `dynamic-foraging-task` meaning of `auto_waterL` / `auto_waterR` (`B_AutoWaterTrial`). With `auto_water*` now reward-keyed to match the reward-delivery series, these keep every autowater delivery recorded in the trial table, so nothing is lost: `auto_water_offered* - auto_water*` is the autowater delivered on trials that did not pay out. | From 6ee6238a19dd012e4bb8d453830a35049419ec71 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Tue, 18 Aug 2026 11:02:59 -0700 Subject: [PATCH 13/13] refactor: use the trial metadata to classify type --- .../processing/_trial_table.py | 92 ++++++------------- .../processing/models/trial_config.py | 17 +--- .../utils/rewards.py | 46 ++++++---- .../utils/trial_metadata.py | 42 +++++++++ 4 files changed, 101 insertions(+), 96 deletions(-) create mode 100644 src/dynamic_foraging_processing/utils/trial_metadata.py diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 98c8fce..52974d2 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -25,6 +25,7 @@ from contraqctor.contract import Dataset from dynamic_foraging_processing.processing.models import TrialConfig +from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata logger = logging.getLogger(__name__) @@ -442,77 +443,48 @@ def _is_baited(trial: Trial, *, is_right: bool) -> bool: return trial.p_reward_left == 1 and auto in (None, True) @staticmethod - def _auto_water(trial: Trial, outcome: TrialOutcome, *, is_right: bool) -> int: - """Return whether autowater was rewarded on the requested side. - - ``is_auto_reward_right`` triggers an immediate reward to one side - (``True`` right, ``False`` left, ``None`` no autowater). The anti-bias - water intervention is itself autowater delivered through this channel, so - every auto-triggered reward counts here; ``anti_bias_left_water`` / - ``anti_bias_right_water`` mark the subset the anti-bias algorithm drove. - - Autowater on a trial that did not pay out (``is_rewarded`` is ``False``) - is ``0``, so this column matches the reward-keyed reward-delivery series: - the water fires at the go cue whether or not the animal's own choice - later pays out, and the series reports only water that was reward. - - Parameters - ---------- - trial : Trial - The per-trial task-logic model. - outcome : TrialOutcome - The trial's outcome, read for ``is_rewarded``. - is_right : bool - ``True`` for the right port, ``False`` for the left port. - - Returns - ------- - int - ``1`` when rewarded autowater targeted the requested side, else ``0``. - """ - if trial.is_auto_reward_right is None or not outcome.is_rewarded: - return 0 - return int(trial.is_auto_reward_right is is_right) - - @staticmethod - def _auto_water_offered(trial: Trial, *, is_right: bool) -> int: - """Return whether autowater was *offered* to the requested side. - - The ungated counterpart of :meth:`_auto_water`: ``1`` whenever the trial - auto-triggered a reward to this side, whether or not the trial went on to - pay out. This is the legacy ``dynamic-foraging-task`` meaning of - ``auto_waterL``/``auto_waterR`` (its ``B_AutoWaterTrial``, which that GUI - also passes to foraging efficiency as ``autowater_offered``), kept so the - trial table still records every autowater delivery while ``auto_water*`` - stays reward-keyed to match the reward-delivery series. + def _auto_water(trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool) -> int: + """Return whether scheduled autowater was delivered to the requested side. + + ``is_auto_reward_right`` is only the delivery *channel* -- it says free + water was triggered and to which side (``True`` right, ``False`` left, + ``None`` none), not what kind. Scheduled autowater and the anti-bias water + intervention share that channel, so the mechanism comes from + ``is_autowater`` and the side from the channel -- the mirror of + :meth:`_anti_bias_water`. Free water driven by the anti-bias algorithm is + ``0`` here and is reported by + ``anti_bias_left_water``/``anti_bias_right_water`` instead. + + Like the anti-bias columns, this records what the *task* did and so is not + gated on ``is_rewarded``: the water fires at the go cue regardless of how + the animal's own choice later resolves. The reward-delivery series is + reward-keyed and drops free water on trials that did not pay out, so this + column can exceed that series' ``auto`` count. Parameters ---------- trial : Trial The per-trial task-logic model. + bias_metadata : BlockBasedTrialMetadata + The trial's extra metadata (see ``_bias_metadata``). is_right : bool ``True`` for the right port, ``False`` for the left port. Returns ------- int - ``1`` when autowater was offered to the requested side, else ``0``. + ``1`` when scheduled autowater targeted the requested side, else ``0``. """ - if trial.is_auto_reward_right is None: + if not bias_metadata.is_autowater: return 0 return int(trial.is_auto_reward_right is is_right) @staticmethod def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: - """Return the block-based extra metadata carrying the anti-bias flags. + """Return the block-based extra metadata carrying the free-water flags. - The anti-bias flags (``is_bias_water_intervention``, - ``is_bias_stage_intervention``) live on ``trial.metadata.extra``. That - field is schema-typed ``Any``, so it deserializes off the stream as a - plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` - instance is also accepted. When metadata or extra is missing (e.g. a - non-block-based generator), the model's all-``False`` default is returned - so the anti-bias columns are simply inert. + Thin wrapper over :func:`get_bias_metadata`, shared with the reward + annotation so both classify autowater and anti-bias water identically. Parameters ---------- @@ -525,13 +497,7 @@ def _bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: The parsed extra metadata, or an all-``False`` default when absent or unrecognized. """ - metadata = trial.metadata - extra = metadata.extra if metadata is not None else None - if isinstance(extra, BlockBasedTrialMetadata): - return extra - if isinstance(extra, dict): - return BlockBasedTrialMetadata.model_validate(extra) - return BlockBasedTrialMetadata() + return get_bias_metadata(trial) @staticmethod def _anti_bias_water( @@ -953,10 +919,8 @@ def _build_row( reward_consumption_duration=trial.reward_consumption_duration, ITI_duration=trial.inter_trial_interval_duration, delay_duration=trial.quiescence_period_duration, - auto_waterL=self._auto_water(trial, outcome, is_right=False), - auto_waterR=self._auto_water(trial, outcome, is_right=True), - auto_water_offeredL=self._auto_water_offered(trial, is_right=False), - auto_water_offeredR=self._auto_water_offered(trial, is_right=True), + auto_waterL=self._auto_water(trial, bias_metadata, is_right=False), + auto_waterR=self._auto_water(trial, bias_metadata, is_right=True), anti_bias_left_water=self._anti_bias_water(trial, bias_metadata, is_right=False), anti_bias_right_water=self._anti_bias_water(trial, bias_metadata, is_right=True), anti_bias_lickspout_movement=self._anti_bias_lickspout_movement(trial, bias_metadata), diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index 57a4211..94b0f47 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -192,26 +192,13 @@ class TrialConfig(BaseModel): auto_waterL: int = Field( default=0, description=( - "Rewarded autowater given at Left: 1 only when the trial auto-triggered a reward to the left (is_auto_reward_right is False) AND the trial was rewarded (is_rewarded). Autowater on a trial that did not pay out is 0, so this column counts autowater that was reward and matches the left reward-delivery series. The anti-bias water intervention uses the same autowater channel and is included here; anti_bias_left_water reports it ungated." + "Scheduled autowater at Left: 1 when the trial's free water was scheduled autowater (trial.metadata.extra.is_autowater) and was directed left (is_auto_reward_right is False). is_auto_reward_right is only the delivery channel, so the mechanism comes from the metadata: free water driven by the anti-bias algorithm is 0 here and is reported by anti_bias_left_water instead. Records what the task did, so like the anti-bias columns it is NOT conditioned on is_rewarded; the reward-delivery series is reward-keyed and drops free water on trials that did not pay out, so this column can exceed that series' auto count." ), ) auto_waterR: int = Field( default=0, description=( - "Rewarded autowater given at Right: 1 only when the trial auto-triggered a reward to the right (is_auto_reward_right is True) AND the trial was rewarded (is_rewarded). Autowater on a trial that did not pay out is 0, so this column counts autowater that was reward and matches the right reward-delivery series. The anti-bias water intervention uses the same autowater channel and is included here; anti_bias_right_water reports it ungated." - ), - ) - - auto_water_offeredL: int = Field( - default=0, - description=( - "Autowater offered at Left: 1 whenever the trial auto-triggered a reward to the left (is_auto_reward_right is False), whether or not the trial paid out. This is the legacy dynamic-foraging-task meaning of auto_waterL, kept so no delivery is lost from the trial table; auto_waterL is the reward-keyed subset that matches the reward-delivery series." - ), - ) - auto_water_offeredR: int = Field( - default=0, - description=( - "Autowater offered at Right: 1 whenever the trial auto-triggered a reward to the right (is_auto_reward_right is True), whether or not the trial paid out. This is the legacy dynamic-foraging-task meaning of auto_waterR, kept so no delivery is lost from the trial table; auto_waterR is the reward-keyed subset that matches the reward-delivery series." + "Scheduled autowater at Right: 1 when the trial's free water was scheduled autowater (trial.metadata.extra.is_autowater) and was directed right (is_auto_reward_right is True). is_auto_reward_right is only the delivery channel, so the mechanism comes from the metadata: free water driven by the anti-bias algorithm is 0 here and is reported by anti_bias_right_water instead. Records what the task did, so like the anti-bias columns it is NOT conditioned on is_rewarded; the reward-delivery series is reward-keyed and drops free water on trials that did not pay out, so this column can exceed that series' auto count." ), ) diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index ee3b75d..ad6d59d 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -7,6 +7,7 @@ from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome from dynamic_foraging_processing.utils.timestamps import find_closest_timestamps +from dynamic_foraging_processing.utils.trial_metadata import get_bias_metadata def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: @@ -31,13 +32,18 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: def _free_water_label(trial: t.Optional[Trial]) -> str: - """Classify a delivery's trial as ``auto`` or ``earned``. + """Classify a delivery's trial as ``anti_bias``, ``auto``, or ``earned``. - ``is_auto_reward_right`` triggers an immediate reward to one side; the - anti-bias water intervention is itself autowater delivered through that - channel, so any auto-triggered reward is ``auto``. This is the same condition - ``auto_waterL``/``auto_waterR`` use in the trials table, so the two agree - trial for trial. + ``is_auto_reward_right`` is only the delivery *channel* -- it says free water + was triggered and to which side, not what kind -- so the mechanism comes from + the block-based metadata: ``is_bias_water_intervention`` for an anti-bias + correction, ``is_autowater`` for scheduled autowater. These are the same + conditions ``anti_bias_left_water``/``anti_bias_right_water`` and + ``auto_waterL``/``auto_waterR`` use in the trials table, so the labels and the + columns classify each trial identically. + + A trial with no free water, or free water the metadata flags as neither + mechanism, is ``earned``. Parameters ---------- @@ -48,11 +54,16 @@ def _free_water_label(trial: t.Optional[Trial]) -> str: Returns ------- str - ``"auto"`` when the trial auto-triggered a reward, else ``"earned"``. + ``"anti_bias"``, ``"auto"``, or ``"earned"``. """ if trial is None or trial.is_auto_reward_right is None: return "earned" - return "auto" + metadata = get_bias_metadata(trial) + if metadata.is_bias_water_intervention: + return "anti_bias" + if metadata.is_autowater: + return "auto" + return "earned" def get_reward_deliveries( @@ -71,16 +82,17 @@ def get_reward_deliveries( ``GiveManualWater`` software event for this port. The software-event timestamps are correlated to the reward-delivery timestamps with :func:`find_closest_timestamps`. - - ``auto`` -- otherwise, when the trial auto-triggered a reward - (``is_auto_reward_right is not None``). The anti-bias water intervention is - autowater delivered through that same channel, so it is ``auto`` too; the - trials table's ``anti_bias_left_water``/``anti_bias_right_water`` mark - which of these the anti-bias algorithm drove. + - ``anti_bias`` -- otherwise, when the trial's free water came from the + anti-bias algorithm (``is_bias_water_intervention``). + - ``auto`` -- otherwise, when the trial's free water was scheduled autowater + (``is_autowater``). - ``earned`` -- otherwise: water the animal worked for. - The ``auto`` condition is the one ``auto_waterL``/``auto_waterR`` use in the - trials table, and the drop below matches those columns' ``is_rewarded`` gate, - so the ``auto`` count here equals the trials table's autowater count. + ``is_auto_reward_right`` is only the delivery *channel*, so the mechanism + behind free water comes from the block-based metadata (see + :func:`get_bias_metadata`) -- the same fields the trials table's + ``auto_waterL``/``auto_waterR`` and + ``anti_bias_left_water``/``anti_bias_right_water`` read. Deliveries on a trial reporting ``is_rewarded=False`` are dropped rather than annotated, so the series reports only water that counted as reward. In @@ -127,7 +139,7 @@ def get_reward_deliveries( the deliveries on unrewarded trials. numpy.ndarray The matching annotations, one per retained timestamp, each ``"earned"``, - ``"auto"``, or ``"manual"``. + ``"auto"``, ``"anti_bias"``, or ``"manual"``. Raises ------ diff --git a/src/dynamic_foraging_processing/utils/trial_metadata.py b/src/dynamic_foraging_processing/utils/trial_metadata.py new file mode 100644 index 0000000..87ef4d9 --- /dev/null +++ b/src/dynamic_foraging_processing/utils/trial_metadata.py @@ -0,0 +1,42 @@ +"""Helpers for reading a trial's block-based extra metadata.""" + +from aind_behavior_dynamic_foraging.task_logic.trial_generators.block_based_trial_generator import ( + BlockBasedTrialMetadata, +) +from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial + + +def get_bias_metadata(trial: Trial) -> BlockBasedTrialMetadata: + """Return the block-based extra metadata naming a trial's free-water mechanism. + + ``trial.is_auto_reward_right`` is only the delivery *channel*: it says free + water was triggered and on which side, not what kind. Scheduled autowater and + the anti-bias water intervention are told apart here, by ``is_autowater`` and + ``is_bias_water_intervention``. (``is_bias_stage_intervention`` marks the + anti-bias algorithm's other lever, moving the lickspouts.) + + The field is schema-typed ``Any``, so it deserializes off the stream as a + plain ``dict`` rather than a model; a ``BlockBasedTrialMetadata`` instance is + also accepted. When metadata or extra is missing (e.g. a non-block-based + generator), the model's all-``False`` default is returned, so a trial whose + mechanism the data does not record is reported as neither kind rather than + guessed at. + + Parameters + ---------- + trial : Trial + The per-trial task-logic model. + + Returns + ------- + BlockBasedTrialMetadata + The parsed extra metadata, or an all-``False`` default when absent + or unrecognized. + """ + metadata = trial.metadata + extra = metadata.extra if metadata is not None else None + if isinstance(extra, BlockBasedTrialMetadata): + return extra + if isinstance(extra, dict): + return BlockBasedTrialMetadata.model_validate(extra) + return BlockBasedTrialMetadata()