From a163968a81147a93ab6f3aa8dde84d86f3e705ae Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Fri, 7 Aug 2026 11:52:25 -0700 Subject: [PATCH 1/4] fix: use repsonse event for correlating with reward delivery times --- .../nwb/acquisition/acquisition_builder.py | 23 +++++++++- .../utils/rewards.py | 42 ++++++++++++++----- 2 files changed, 53 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 a63fda2..add7bcb 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -76,6 +76,17 @@ def get_trial_outcomes(self) -> pd.DataFrame: self.loader.dataset.at("Behavior").at("SoftwareEvents").at("TrialOutcome").load().data ) + def get_responses(self) -> pd.DataFrame: + """Get the ``Response`` software-event stream. + + Returns + ------- + pandas.DataFrame + The ``Response`` stream under ``Behavior/SoftwareEvents``, indexed + by the timestamp at which the animal responded, one row per trial. + """ + return self.loader.dataset.at("Behavior").at("SoftwareEvents").at("Response").load().data + def get_manual_water_times(self) -> pd.DataFrame: """Get the manual-water software-event stream. @@ -171,6 +182,7 @@ def _reward_delivery_series( self, writes: pd.DataFrame, trial_outcomes: pd.DataFrame, + responses: pd.DataFrame, manual_water: pd.DataFrame, *, port_column: str, @@ -182,7 +194,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 - automatic via :func:`get_annotated_rewards`. + auto via :func:`get_annotated_rewards`. Parameters ---------- @@ -190,6 +202,9 @@ def _reward_delivery_series( ``OutputSet`` ``WRITE`` messages indexed by timestamp. trial_outcomes : pandas.DataFrame The ``TrialOutcome`` stream, indexed by trial timestamp. + responses : pandas.DataFrame + The ``Response`` stream, indexed by the time the animal responded; + positionally aligned with ``trial_outcomes``. manual_water : pandas.DataFrame The ``GiveManualWaterRight`` stream; the ``data`` column selects the side (``True`` right, ``False`` left). @@ -214,6 +229,7 @@ def _reward_delivery_series( annotations = get_annotated_rewards( delivery_times, trial_outcomes, + responses.index.to_numpy(), manual_water_times, ) return AcquisitionSeries( @@ -223,7 +239,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 automatic" + "annotates whether the reward was earned, manual, or auto" ), ) @@ -250,6 +266,7 @@ def build_acquisition( """ rewards = self.get_reward_delivery() trial_outcomes = self.get_trial_outcomes() + responses = self.get_responses() manual_water = self.get_manual_water_times() acquisition_streams = self.loader.get_all_raw_data() @@ -273,6 +290,7 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, + responses, manual_water, port_column="SupplyPort0", is_right=False, @@ -284,6 +302,7 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, + responses, manual_water, port_column="SupplyPort1", is_right=True, diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 068a144..43a5ea4 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -33,9 +33,10 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: def get_annotated_rewards( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, + response_times: np.ndarray, manual_water_times: np.ndarray, ) -> np.ndarray: - """Annotate each reward delivery as ``earned``, ``automatic``, or ``manual``. + """Annotate each reward delivery as ``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 @@ -45,17 +46,26 @@ def get_annotated_rewards( ``GiveManualWater`` software event for this port. The software-event timestamps are correlated to the reward-delivery timestamps with :func:`find_closest_timestamps`. - - ``automatic`` -- otherwise, when the matching trial auto-responded + - ``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 are matched to trials by the ``Response`` software event, which + is emitted when the animal responds and so sits next to the delivery it + caused. The ``Response`` and ``TrialOutcome`` streams are emitted once per + trial and are positionally aligned, so the matched ``Response`` position + indexes the corresponding ``TrialOutcome`` row. + Parameters ---------- reward_delivery_times : numpy.ndarray Hardware (harp) timestamps of this port's reward deliveries. trial_outcome_df : pandas.DataFrame - Trial outcome table indexed by trial timestamp; each row's ``data`` - field is a :class:`TrialOutcome` payload. + Trial outcome table with one row per trial; each row's ``data`` field + is a :class:`TrialOutcome` payload. + response_times : numpy.ndarray + Software-event timestamps of the per-trial ``Response`` events, in + trial order (one per row of ``trial_outcome_df``). manual_water_times : numpy.ndarray Software-event timestamps of this port's manual water deliveries (``GiveManualWaterLeft`` / ``GiveManualWaterRight``). @@ -64,17 +74,27 @@ def get_annotated_rewards( ------- numpy.ndarray Array of the same shape as ``reward_delivery_times`` whose entries are - ``"earned"``, ``"automatic"``, or ``"manual"``. + ``"earned"``, ``"auto"``, or ``"manual"``. + + Raises + ------ + ValueError + If ``response_times`` and ``trial_outcome_df`` have different lengths, + i.e. the per-trial streams are misaligned. """ reward_times = np.asarray(reward_delivery_times) + response_times = np.asarray(response_times) + if response_times.size != len(trial_outcome_df): + raise ValueError( + "Response and TrialOutcome streams are misaligned: " + f"{response_times.size} responses vs {len(trial_outcome_df)} trial outcomes." + ) 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: @@ -83,9 +103,11 @@ def get_annotated_rewards( if trial is None or trial.is_auto_reward_right is None: annotated_rewards.append("earned") else: - annotated_rewards.append("automatic") + annotated_rewards.append("auto") - annotated_rewards = np.array(annotated_rewards) + # 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. + annotated_rewards = np.array(annotated_rewards, dtype=object) # Manual water is independent of trials (multiple can occur within a trial) and # takes precedence, so annotate the manual deliveries directly. Correlate each From c2e9f9cbcc0ff95b4bed5f1d4e54e6a6683eaf2f Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Fri, 7 Aug 2026 11:52:38 -0700 Subject: [PATCH 2/4] test: update tests --- .../test_acquisition_builder.py | 26 +++++- tests/test_pipeline/test_pipeline.py | 2 +- tests/test_utils/test_rewards.py | 82 ++++++++++++++----- 3 files changed, 87 insertions(+), 23 deletions(-) diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index 4b1f91f..db9b2b5 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: + """The per-trial ``Response`` events, aligned with the trial outcome frame.""" + return pd.DataFrame( + {"data": [{"Item1": 0.1, "Item2": False}, {"Item1": 0.4, "Item2": True}]}, + index=pd.Index([0.1, 0.4], 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), } ), @@ -157,6 +166,15 @@ def test_get_reward_delivery_filters_to_write_messages(): assert list(result.index) == [0.1, 0.3, 0.5] +def test_get_responses_returns_stream(): + """``get_responses`` returns the per-trial ``Response`` stream.""" + builder = AcquisitionBuilder(loader=_make_loader()) + + result = builder.get_responses() + + pd.testing.assert_frame_equal(result, _make_response_frame()) + + def test_get_manual_water_times_returns_stream(): """``get_manual_water_times`` returns the GiveManualWaterRight stream.""" manual = pd.DataFrame({"data": [True]}, index=pd.Index([0.49], name="time")) @@ -175,7 +193,10 @@ def test_get_manual_water_times_returns_empty_when_absent(): { "HarpBehavior": _FakeNode({"OutputSet": _FakeStream(_make_output_set_frame())}), "SoftwareEvents": _FakeNode( - {"TrialOutcome": _FakeStream(_make_trial_outcome_frame())} + { + "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), + "Response": _FakeStream(_make_response_frame()), + } ), } ) @@ -222,6 +243,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()), } ), @@ -265,7 +287,7 @@ def test_build_acquisition_returns_populated_list(): # the second is overridden to manual by the right-side manual-water event. assert isinstance(right_reward, AcquisitionSeries) np.testing.assert_array_equal(right_reward.timestamps, np.array([0.3, 0.5])) - np.testing.assert_array_equal(right_reward.data, np.array(["automatic", "manual"])) + np.testing.assert_array_equal(right_reward.data, np.array(["auto", "manual"])) assert right_reward.name == "right_reward_delivery_time" assert "right lick port" in right_reward.description diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py index 1124adc..c897ea7 100644 --- a/tests/test_pipeline/test_pipeline.py +++ b/tests/test_pipeline/test_pipeline.py @@ -348,7 +348,7 @@ def test_manual_water_times_reads_manual_annotations(): "left_reward_delivery_time": _FakeSeries( np.array(["manual", "earned", "manual"]), np.array([0.1, 0.2, 0.3]) ), - "right_reward_delivery_time": _FakeSeries(np.array(["automatic"]), np.array([0.5])), + "right_reward_delivery_time": _FakeSeries(np.array(["auto"]), np.array([0.5])), } left, right = Pipeline._manual_water_times(nwb_file) diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index 3fcc4e0..833158b 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -4,6 +4,7 @@ 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 @@ -27,7 +28,7 @@ def _outcome_payload(auto=None) -> dict: def _trial_outcome_df(trial_times: np.ndarray, autos=None) -> pd.DataFrame: - """Build a trial outcome DataFrame indexed by ``trial_times``.""" + """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) return pd.DataFrame( {"data": [_outcome_payload(auto) for auto in autos]}, @@ -38,75 +39,116 @@ def _trial_outcome_df(trial_times: np.ndarray, autos=None) -> pd.DataFrame: def test_get_annotated_rewards_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(response_times) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, np.array([]) + ) np.testing.assert_array_equal(annotations, np.array(["earned", "earned", "earned"])) -def test_get_annotated_rewards_marks_auto_response_trials_as_automatic(): - """Trials with ``is_auto_reward_right`` set (either side) are ``automatic``.""" +def test_get_annotated_rewards_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]) + response_times = np.array([0.1, 0.4]) + trial_outcome_df = _trial_outcome_df(response_times, autos=[True, False]) + + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, np.array([]) + ) + + np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) + + +def test_get_annotated_rewards_matches_response_times_not_outcome_times(): + """Deliveries are matched to trials by ``Response`` time, not outcome time.""" + # The TrialOutcome timestamps here trail the deliveries; only the Response + # times sit next to them, so matching on the outcome index would pick the + # wrong trial and annotate the second delivery as earned. 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([0.4, 5.0]), autos=[None, True]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, np.array([]) + ) - np.testing.assert_array_equal(annotations, np.array(["automatic", "automatic"])) + np.testing.assert_array_equal(annotations, np.array(["earned", "auto"])) def test_get_annotated_rewards_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(response_times) # 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) + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, manual_water_times + ) np.testing.assert_array_equal(annotations, np.array(["earned", "manual", "earned"])) -def test_get_annotated_rewards_manual_takes_precedence_over_automatic(): +def test_get_annotated_rewards_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(response_times, autos=[None, True]) manual_water_times = np.array([0.42]) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, manual_water_times) + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, manual_water_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.""" - trial_outcome_df = _trial_outcome_df(np.array([0.0])) + response_times = np.array([0.0]) + trial_outcome_df = _trial_outcome_df(response_times) - result = get_annotated_rewards(np.array([]), trial_outcome_df, np.array([])) + result = get_annotated_rewards(np.array([]), trial_outcome_df, response_times, np.array([])) assert isinstance(result, np.ndarray) assert result.size == 0 +def test_get_annotated_rewards_misaligned_responses_raises(): + """A ``Response`` count that disagrees with the trial count is an error.""" + trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4])) + + with pytest.raises(ValueError, match="misaligned"): + get_annotated_rewards(np.array([0.15]), trial_outcome_df, np.array([0.1]), np.array([])) + + def test_get_annotated_rewards_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(response_times, name="time"), ) - annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + annotations = get_annotated_rewards( + reward_times, trial_outcome_df, response_times, np.array([]) + ) - np.testing.assert_array_equal(annotations, np.array(["automatic", "automatic"])) + 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])) + response_times = np.array([0.0]) + trial_outcome_df = _trial_outcome_df(response_times) - result = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) + result = get_annotated_rewards(reward_times, trial_outcome_df, response_times, np.array([])) assert isinstance(result, np.ndarray) From 683a755ba7f9fca7bb1e537c09fe53941f007799 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Wed, 12 Aug 2026 14:32:51 -0700 Subject: [PATCH 3/4] refactor: revert to using trial outcome and set automatic label to auto --- .../nwb/acquisition/acquisition_builder.py | 19 ------------ .../utils/rewards.py | 31 +++++-------------- 2 files changed, 7 insertions(+), 43 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index add7bcb..ebd5270 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -76,17 +76,6 @@ def get_trial_outcomes(self) -> pd.DataFrame: self.loader.dataset.at("Behavior").at("SoftwareEvents").at("TrialOutcome").load().data ) - def get_responses(self) -> pd.DataFrame: - """Get the ``Response`` software-event stream. - - Returns - ------- - pandas.DataFrame - The ``Response`` stream under ``Behavior/SoftwareEvents``, indexed - by the timestamp at which the animal responded, one row per trial. - """ - return self.loader.dataset.at("Behavior").at("SoftwareEvents").at("Response").load().data - def get_manual_water_times(self) -> pd.DataFrame: """Get the manual-water software-event stream. @@ -182,7 +171,6 @@ def _reward_delivery_series( self, writes: pd.DataFrame, trial_outcomes: pd.DataFrame, - responses: pd.DataFrame, manual_water: pd.DataFrame, *, port_column: str, @@ -202,9 +190,6 @@ def _reward_delivery_series( ``OutputSet`` ``WRITE`` messages indexed by timestamp. trial_outcomes : pandas.DataFrame The ``TrialOutcome`` stream, indexed by trial timestamp. - responses : pandas.DataFrame - The ``Response`` stream, indexed by the time the animal responded; - positionally aligned with ``trial_outcomes``. manual_water : pandas.DataFrame The ``GiveManualWaterRight`` stream; the ``data`` column selects the side (``True`` right, ``False`` left). @@ -229,7 +214,6 @@ def _reward_delivery_series( annotations = get_annotated_rewards( delivery_times, trial_outcomes, - responses.index.to_numpy(), manual_water_times, ) return AcquisitionSeries( @@ -266,7 +250,6 @@ def build_acquisition( """ rewards = self.get_reward_delivery() trial_outcomes = self.get_trial_outcomes() - responses = self.get_responses() manual_water = self.get_manual_water_times() acquisition_streams = self.loader.get_all_raw_data() @@ -290,7 +273,6 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, - responses, manual_water, port_column="SupplyPort0", is_right=False, @@ -302,7 +284,6 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, - responses, manual_water, port_column="SupplyPort1", is_right=True, diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 43a5ea4..02234a8 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -33,7 +33,6 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: def get_annotated_rewards( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, - response_times: np.ndarray, manual_water_times: np.ndarray, ) -> np.ndarray: """Annotate each reward delivery as ``earned``, ``auto``, or ``manual``. @@ -50,22 +49,16 @@ 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 ``Response`` software event, which - is emitted when the animal responds and so sits next to the delivery it - caused. The ``Response`` and ``TrialOutcome`` streams are emitted once per - trial and are positionally aligned, so the matched ``Response`` position - indexes the corresponding ``TrialOutcome`` row. + Deliveries are matched to trials by the ``TrialOutcome`` software-event + timestamp: each delivery takes the annotation of the closest trial. Parameters ---------- reward_delivery_times : numpy.ndarray Hardware (harp) timestamps of this port's reward deliveries. trial_outcome_df : pandas.DataFrame - Trial outcome table with one row per trial; each row's ``data`` field - is a :class:`TrialOutcome` payload. - response_times : numpy.ndarray - Software-event timestamps of the per-trial ``Response`` events, in - trial order (one per row of ``trial_outcome_df``). + Trial outcome table indexed by trial timestamp; each row's ``data`` + field is a :class:`TrialOutcome` payload. manual_water_times : numpy.ndarray Software-event timestamps of this port's manual water deliveries (``GiveManualWaterLeft`` / ``GiveManualWaterRight``). @@ -75,26 +68,16 @@ def get_annotated_rewards( numpy.ndarray Array of the same shape as ``reward_delivery_times`` whose entries are ``"earned"``, ``"auto"``, or ``"manual"``. - - Raises - ------ - ValueError - If ``response_times`` and ``trial_outcome_df`` have different lengths, - i.e. the per-trial streams are misaligned. """ reward_times = np.asarray(reward_delivery_times) - response_times = np.asarray(response_times) - if response_times.size != len(trial_outcome_df): - raise ValueError( - "Response and TrialOutcome streams are misaligned: " - f"{response_times.size} responses vs {len(trial_outcome_df)} trial outcomes." - ) 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, response_times) + trial_indices_in_reward_times = find_closest_timestamps( + reward_times, trial_outcome_df.index.to_numpy() + ) annotated_rewards = [] for trial_index in trial_indices_in_reward_times: From 45160bed6b6dcac6cdde202d260cc6955860689f Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Wed, 12 Aug 2026 14:33:04 -0700 Subject: [PATCH 4/4] test: update tests --- .../test_acquisition_builder.py | 20 ----- tests/test_utils/test_rewards.py | 74 ++++++------------- 2 files changed, 22 insertions(+), 72 deletions(-) diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index db9b2b5..f16915e 100644 --- a/tests/test_nwb/test_acquisition/test_acquisition_builder.py +++ b/tests/test_nwb/test_acquisition/test_acquisition_builder.py @@ -82,14 +82,6 @@ def _make_trial_outcome_frame() -> pd.DataFrame: ) -def _make_response_frame() -> pd.DataFrame: - """The per-trial ``Response`` events, aligned with the trial outcome frame.""" - return pd.DataFrame( - {"data": [{"Item1": 0.1, "Item2": False}, {"Item1": 0.4, "Item2": True}]}, - index=pd.Index([0.1, 0.4], 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")) @@ -128,7 +120,6 @@ def _make_dataset(manual_water=None): "SoftwareEvents": _FakeNode( { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), - "Response": _FakeStream(_make_response_frame()), "GiveManualWaterRight": _FakeStream(manual_water), } ), @@ -166,15 +157,6 @@ def test_get_reward_delivery_filters_to_write_messages(): assert list(result.index) == [0.1, 0.3, 0.5] -def test_get_responses_returns_stream(): - """``get_responses`` returns the per-trial ``Response`` stream.""" - builder = AcquisitionBuilder(loader=_make_loader()) - - result = builder.get_responses() - - pd.testing.assert_frame_equal(result, _make_response_frame()) - - def test_get_manual_water_times_returns_stream(): """``get_manual_water_times`` returns the GiveManualWaterRight stream.""" manual = pd.DataFrame({"data": [True]}, index=pd.Index([0.49], name="time")) @@ -195,7 +177,6 @@ def test_get_manual_water_times_returns_empty_when_absent(): "SoftwareEvents": _FakeNode( { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), - "Response": _FakeStream(_make_response_frame()), } ), } @@ -243,7 +224,6 @@ 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 833158b..2be8269 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -4,7 +4,6 @@ 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 @@ -39,12 +38,9 @@ def _trial_outcome_df(trial_times: np.ndarray, autos=None) -> pd.DataFrame: def test_get_annotated_rewards_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]) - response_times = np.array([0.1, 0.4, 0.9]) - trial_outcome_df = _trial_outcome_df(response_times) + trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4, 0.9])) - annotations = get_annotated_rewards( - reward_times, trial_outcome_df, response_times, np.array([]) - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) np.testing.assert_array_equal(annotations, np.array(["earned", "earned", "earned"])) @@ -52,43 +48,33 @@ def test_get_annotated_rewards_marks_default_trials_as_earned(): def test_get_annotated_rewards_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]) - response_times = np.array([0.1, 0.4]) - trial_outcome_df = _trial_outcome_df(response_times, autos=[True, False]) + trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4]), autos=[True, False]) - annotations = get_annotated_rewards( - reward_times, trial_outcome_df, response_times, np.array([]) - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -def test_get_annotated_rewards_matches_response_times_not_outcome_times(): - """Deliveries are matched to trials by ``Response`` time, not outcome time.""" - # The TrialOutcome timestamps here trail the deliveries; only the Response - # times sit next to them, so matching on the outcome index would pick the - # wrong trial and annotate the second delivery as earned. - reward_times = np.array([0.15, 0.42]) - response_times = np.array([0.1, 0.4]) - trial_outcome_df = _trial_outcome_df(np.array([0.4, 5.0]), autos=[None, True]) +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. + reward_times = np.array([0.95, 1.05]) + trial_outcome_df = _trial_outcome_df(np.array([0.1, 1.0]), autos=[None, True]) - annotations = get_annotated_rewards( - reward_times, trial_outcome_df, response_times, np.array([]) - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) - np.testing.assert_array_equal(annotations, np.array(["earned", "auto"])) + np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) def test_get_annotated_rewards_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]) - response_times = np.array([0.1, 0.4, 0.9]) - trial_outcome_df = _trial_outcome_df(response_times) + trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4, 0.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, response_times, manual_water_times - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, manual_water_times) np.testing.assert_array_equal(annotations, np.array(["earned", "manual", "earned"])) @@ -96,49 +82,34 @@ def test_get_annotated_rewards_marks_manual_water_as_manual(): def test_get_annotated_rewards_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]) - response_times = np.array([0.1, 0.4]) - trial_outcome_df = _trial_outcome_df(response_times, autos=[None, True]) + trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4]), autos=[None, True]) manual_water_times = np.array([0.42]) - annotations = get_annotated_rewards( - reward_times, trial_outcome_df, response_times, manual_water_times - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, manual_water_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.""" - response_times = np.array([0.0]) - trial_outcome_df = _trial_outcome_df(response_times) + trial_outcome_df = _trial_outcome_df(np.array([0.0])) - result = get_annotated_rewards(np.array([]), trial_outcome_df, response_times, np.array([])) + result = get_annotated_rewards(np.array([]), trial_outcome_df, np.array([])) assert isinstance(result, np.ndarray) assert result.size == 0 -def test_get_annotated_rewards_misaligned_responses_raises(): - """A ``Response`` count that disagrees with the trial count is an error.""" - trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4])) - - with pytest.raises(ValueError, match="misaligned"): - get_annotated_rewards(np.array([0.15]), trial_outcome_df, np.array([0.1]), np.array([])) - - def test_get_annotated_rewards_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(response_times, name="time"), + index=pd.Index([0.1, 0.4], name="time"), ) - annotations = get_annotated_rewards( - reward_times, trial_outcome_df, response_times, np.array([]) - ) + annotations = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) @@ -146,9 +117,8 @@ def test_get_annotated_rewards_accepts_json_and_model_payloads(): def test_get_annotated_rewards_returns_ndarray(): """The return value is a ``numpy.ndarray``.""" reward_times = np.array([0.1]) - response_times = np.array([0.0]) - trial_outcome_df = _trial_outcome_df(response_times) + trial_outcome_df = _trial_outcome_df(np.array([0.0])) - result = get_annotated_rewards(reward_times, trial_outcome_df, response_times, np.array([])) + result = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) assert isinstance(result, np.ndarray)