From a710fd7af4ae44c9bfeb289e7f06343ffab017a3 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 <109561860+arjunsridhar12345@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:18:02 -0700 Subject: [PATCH] fix: autowater on ignore trials annotated as earned (#105) * fix: use trial outcome window to match to reward deliveries instead of response event * test: update tests * refactor: remove verbose ai comments * fix: linting --- .../nwb/acquisition/acquisition_builder.py | 28 -- .../utils/rewards.py | 57 ++-- tests/test_utils/test_rewards.py | 245 ++++++++++-------- 3 files changed, 164 insertions(+), 166 deletions(-) diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index be51f43..a890d23 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -76,26 +76,6 @@ 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 _software_event_times(self, stream_name: str) -> np.ndarray: """Get one ``Behavior/SoftwareEvents`` stream's event timestamps. @@ -218,7 +198,6 @@ def _reward_delivery_series( writes: pd.DataFrame, trial_outcomes: pd.DataFrame, manual_water: ManualWaterTimes, - response_times: np.ndarray, *, port_column: str, name: str, @@ -241,9 +220,6 @@ def _reward_delivery_series( manual_water : ManualWaterTimes This port's experimenter-triggered water times, already side-specific (see :meth:`get_manual_water_times`). - 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). @@ -263,7 +239,6 @@ def _reward_delivery_series( delivery_times, trial_outcomes, manual_water, - response_times, ) return AcquisitionSeries( name=name, @@ -303,7 +278,6 @@ def build_acquisition( trial_outcomes = self.get_trial_outcomes() left_manual_water = self.get_manual_water_times(is_right=False) right_manual_water = self.get_manual_water_times(is_right=True) - response_times = self.get_response_times() acquisition_streams = self.loader.get_all_raw_data() acqusition_streams_descriptions = self.loader.raw_data_stream_descriptions @@ -327,7 +301,6 @@ def build_acquisition( rewards, trial_outcomes, left_manual_water, - response_times, port_column="SupplyPort0", name="left_reward_delivery_time", side_label="left", @@ -338,7 +311,6 @@ def build_acquisition( rewards, trial_outcomes, right_manual_water, - response_times, port_column="SupplyPort1", name="right_reward_delivery_time", side_label="right", diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 73287ec..bd6d30f 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -95,7 +95,6 @@ def get_reward_deliveries( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, manual_water: ManualWaterTimes, - response_times: np.ndarray, ) -> np.ndarray: """Classify one lick port's reward deliveries by how the water was given. @@ -138,17 +137,18 @@ def get_reward_deliveries( contribute an ``auto`` delivery on the side the task watered and an ``earned`` delivery on the side the animal chose. - 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``. + Deliveries are matched to trials by *containment*, not proximity: a + ``TrialOutcome`` fires at the end of a trial, after the reward-consumption + and ITI periods, so trial ``i`` spans ``(outcome[i - 1], outcome[i]]`` and a + delivery belongs to the first trial whose outcome it precedes. - ``response_times`` is aligned to ``trial_outcome_df`` positionally: entry - ``i`` is the response of the trial in row ``i``. + Proximity to the ``Response`` event was used previously and is unsound: it + has no notion of a trial boundary, so a delivery near the start of its trial + can be closer to the *previous* trial's response and inherit that trial's + ``is_auto_reward_right``, flipping ``auto`` to ``earned``. Autowater on an + ignore trial is the worst case -- the water lands at the go cue while the + trial's own ``Response`` waits out the full response deadline -- but the + boundary is what makes containment correct, whatever the response latency. Parameters ---------- @@ -161,9 +161,6 @@ def get_reward_deliveries( This port's experimenter-triggered water times, split into ``unaligned`` and ``go_cue_aligned``. Either field may be empty when the session has no water of that kind. - response_times : numpy.ndarray - ``Response`` software-event timestamps, one per trial, positionally - aligned with the rows of ``trial_outcome_df``. Returns ------- @@ -174,23 +171,31 @@ def get_reward_deliveries( Raises ------ ValueError - If ``response_times`` has a different length than ``trial_outcome_df``, - since the two are paired by position. + If ``trial_outcome_df`` is empty, or its index is unsorted or contains + ``NaN``, since the index is used as the trial window boundaries. """ - 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, response_times) + # The trial's own end time bounds it, so the index doubles as the window + # edges. Guard the assumptions searchsorted makes rather than letting a bad + # index silently push every delivery onto one trial: an all-NaN index would + # otherwise assign them all to the last trial. + trial_end_times = np.asarray(trial_outcome_df.index, dtype=float) + if trial_end_times.size == 0: + raise ValueError("trial_outcome_df is empty; deliveries cannot be matched to a trial.") + if np.isnan(trial_end_times).any(): + raise ValueError("trial_outcome_df index contains NaN; trial windows are undefined.") + if np.any(np.diff(trial_end_times) < 0): + raise ValueError("trial_outcome_df index must be sorted to bound trials.") + + # side="left" so a delivery landing exactly on a trial's outcome belongs to + # that trial rather than the next. + trial_indices_in_reward_times = np.minimum( + np.searchsorted(trial_end_times, reward_times, side="left"), + trial_end_times.size - 1, + ) trial_labels = [] for trial_index in trial_indices_in_reward_times: diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index 6ce87bc..6c40d4d 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -44,7 +44,12 @@ def _outcome_payload(auto=None, is_rewarded: bool = True, mechanism: str = "auto def _trial_outcome_df( trial_times: np.ndarray, autos=None, rewarded=None, mechanism: str = "autowater" ) -> pd.DataFrame: - """Build a trial outcome DataFrame with one row per entry of ``trial_times``.""" + """Build a trial outcome DataFrame with one row per entry of ``trial_times``. + + ``trial_times`` are the ``TrialOutcome`` event timestamps, which fire at the + *end* of each trial and so bound it: trial ``i`` spans + ``(trial_times[i - 1], trial_times[i]]``. + """ 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( @@ -55,26 +60,20 @@ def _trial_outcome_df( 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]) - response_times = np.array([0.1, 0.4, 0.9]) - trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4, 1.9])) + reward_times = np.array([0.15, 1.42, 2.95]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0, 3.0])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["earned", "earned", "earned"])) 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]) - response_times = np.array([0.1, 0.4]) - trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4]), autos=[True, False]) + reward_times = np.array([0.15, 1.42]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0]), autos=[True, False]) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) @@ -89,37 +88,73 @@ def test_get_reward_deliveries_marks_all_free_water_as_auto(mechanism): ``anti_bias_left_water``/``anti_bias_right_water`` record the mechanism per trial instead. """ - reward_times = np.array([0.15]) - response_times = np.array([0.1]) trial_outcome_df = pd.DataFrame( {"data": [_outcome_payload(True, mechanism=mechanism)]}, - index=pd.Index([1.1], name="time"), + index=pd.Index([1.0], name="time"), ) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(np.array([0.15]), trial_outcome_df, ManualWaterTimes()) 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. +def test_get_reward_deliveries_matches_the_containing_trial_not_the_nearest(): + """A delivery takes the trial whose window contains it, however near others are. - 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. + The first delivery sits just past trial 0's outcome, so it belongs to trial + 1 even though it is far closer to trial 0's boundary than to trial 1's. + Proximity matching would charge it to trial 0 and read the wrong + ``is_auto_reward_right``. """ - reward_times = np.array([0.95, 1.05]) - 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]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 9.0]), autos=[None, True]) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(np.array([1.05]), trial_outcome_df, ManualWaterTimes()) - np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) + np.testing.assert_array_equal(annotations, np.array(["auto"])) + + +def test_get_reward_deliveries_labels_autowater_on_an_ignore_trial_as_auto(): + """Autowater at the go cue of an ignore trial is ``auto``, not ``earned``. + + Regression for the mislabel seen on 872547_2026-09-11_13-05-03 trial 499. On + an ignore trial the ``Response`` event waits out the full response deadline + while autowater lands at the go cue, so the *previous* trial's response is + the nearer one. Matching by proximity therefore read the previous trial's + ``is_auto_reward_right`` (unset) and annotated task-given free water as water + the animal worked for. Containment is immune: the delivery falls inside its + own trial's window whatever the response latency. + """ + # Trial 0 has no free water and ends at t=10; trial 1 is an autowater trial + # whose water lands immediately after that boundary, at its own go cue. + trial_outcome_df = _trial_outcome_df(np.array([10.0, 20.0]), autos=[None, True]) + autowater_at_go_cue = np.array([10.001]) + + annotations = get_reward_deliveries(autowater_at_go_cue, trial_outcome_df, ManualWaterTimes()) + + np.testing.assert_array_equal(annotations, np.array(["auto"])) + + +def test_get_reward_deliveries_assigns_a_delivery_on_the_boundary_to_that_trial(): + """A delivery landing exactly on a trial's outcome belongs to that trial.""" + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0]), autos=[True, None]) + + annotations = get_reward_deliveries(np.array([1.0]), trial_outcome_df, ManualWaterTimes()) + + np.testing.assert_array_equal(annotations, np.array(["auto"])) + + +def test_get_reward_deliveries_charges_water_after_the_last_trial_to_that_trial(): + """Water delivered past the final outcome has no trial of its own. + + End-of-session experimenter water can land after the last ``TrialOutcome``. + It is charged to the last trial rather than indexing off the end; in practice + the manual labels overwrite it. + """ + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0]), autos=[None, True]) + + annotations = get_reward_deliveries(np.array([99.0]), trial_outcome_df, ManualWaterTimes()) + + np.testing.assert_array_equal(annotations, np.array(["auto"])) def test_get_reward_deliveries_keeps_deliveries_on_unrewarded_trials(): @@ -129,17 +164,14 @@ def test_get_reward_deliveries_keeps_deliveries_on_unrewarded_trials(): ``is_rewarded`` describes the animal's own choice rather than the water. The series records every valve opening, so nothing is filtered out. """ - reward_times = np.array([0.15, 0.42, 0.95]) - response_times = np.array([0.1, 0.4, 0.9]) + reward_times = np.array([0.15, 1.42, 2.95]) trial_outcome_df = _trial_outcome_df( - np.array([1.1, 1.4, 1.9]), + np.array([1.0, 2.0, 3.0]), autos=[None, True, True], rewarded=[True, False, True], ) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["earned", "auto", "auto"])) @@ -149,37 +181,29 @@ def test_get_reward_deliveries_labels_both_sides_of_a_split_trial(): When free water goes to one port and the animal earns reward at the other, the trial contributes an ``auto`` delivery and an ``earned`` delivery. The - label follows the matched trial, so both deliveries on that trial read + label follows the containing trial, so both deliveries on that trial read ``auto`` from this port's perspective; the sides are separate series. """ - trial_outcome_df = _trial_outcome_df(np.array([1.1]), autos=[False], rewarded=[True]) - response_times = np.array([0.1]) + trial_outcome_df = _trial_outcome_df(np.array([1.0]), autos=[False], rewarded=[True]) # This port saw one opening on that trial; the trial gave free water. - annotations = get_reward_deliveries( - np.array([0.15]), trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(np.array([0.15]), trial_outcome_df, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["auto"])) # A trial with no free water at all yields ``earned`` on whichever port opened. - earned_only = _trial_outcome_df(np.array([1.1]), autos=[None], rewarded=[True]) - annotations = get_reward_deliveries( - np.array([0.15]), earned_only, ManualWaterTimes(), response_times - ) + earned_only = _trial_outcome_df(np.array([1.0]), autos=[None], rewarded=[True]) + annotations = get_reward_deliveries(np.array([0.15]), earned_only, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["earned"])) def test_get_reward_deliveries_marks_manual_water_as_manual(): """Deliveries closest to an unaligned manual-water event are ``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(np.array([1.1, 1.4, 1.9])) - # Software event near the second delivery (0.42). - manual_water = ManualWaterTimes(unaligned=np.array([0.43])) - - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + reward_times = np.array([0.15, 1.42, 2.95]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0, 3.0])) + # Software event near the second delivery (1.42). + manual_water = ManualWaterTimes(unaligned=np.array([1.43])) + + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal(annotations, np.array(["earned", "manual", "earned"])) @@ -191,14 +215,11 @@ def test_get_reward_deliveries_marks_manual_auto_reward_as_go_cue_aligned(): set ``is_auto_reward_right`` and the trial-derived label would read ``earned`` -- water the animal never worked for. """ - 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])) - manual_water = ManualWaterTimes(go_cue_aligned=np.array([0.43])) + reward_times = np.array([0.15, 1.42, 2.95]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0, 3.0])) + manual_water = ManualWaterTimes(go_cue_aligned=np.array([1.43])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal( annotations, np.array(["earned", "manual_go_cue_aligned", "earned"]) @@ -207,14 +228,11 @@ def test_get_reward_deliveries_marks_manual_auto_reward_as_go_cue_aligned(): 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]) - 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 = ManualWaterTimes(unaligned=np.array([0.42])) + reward_times = np.array([0.15, 1.42]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0]), autos=[None, True]) + manual_water = ManualWaterTimes(unaligned=np.array([1.42])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal(annotations, np.array(["earned", "manual"])) @@ -225,14 +243,11 @@ def test_get_reward_deliveries_go_cue_aligned_takes_precedence_over_auto(): Free water can fire on the same trial the experimenter watered; the label reports who caused *this* delivery, and the experimenter is the narrower fact. """ - 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]) - manual_water = ManualWaterTimes(go_cue_aligned=np.array([0.42])) + reward_times = np.array([0.15, 1.42]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0]), autos=[None, True]) + manual_water = ManualWaterTimes(go_cue_aligned=np.array([1.42])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal(annotations, np.array(["earned", "manual_go_cue_aligned"])) @@ -243,28 +258,22 @@ def test_get_reward_deliveries_unaligned_manual_outranks_go_cue_aligned(): Unaligned manual water is the stronger claim: it says the delivery is tied to no go cue at all, so it is written last and overwrites the aligned label. """ - 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])) - manual_water = ManualWaterTimes(unaligned=np.array([0.42]), go_cue_aligned=np.array([0.42])) + reward_times = np.array([0.15, 1.42]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0])) + manual_water = ManualWaterTimes(unaligned=np.array([1.42]), go_cue_aligned=np.array([1.42])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal(annotations, np.array(["earned", "manual"])) def test_get_reward_deliveries_labels_both_manual_kinds_in_one_session(): """A session can carry both kinds of experimenter water on the same port.""" - 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])) - manual_water = ManualWaterTimes(unaligned=np.array([0.15]), go_cue_aligned=np.array([0.95])) + reward_times = np.array([0.15, 1.42, 2.95]) + trial_outcome_df = _trial_outcome_df(np.array([1.0, 2.0, 3.0])) + manual_water = ManualWaterTimes(unaligned=np.array([0.15]), go_cue_aligned=np.array([2.95])) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water, response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, manual_water) np.testing.assert_array_equal( annotations, np.array(["manual", "earned", "manual_go_cue_aligned"]) @@ -275,9 +284,7 @@ def test_get_reward_deliveries_empty_deliveries_returns_empty(): """No reward deliveries yields an empty annotation array.""" trial_outcome_df = _trial_outcome_df(np.array([0.0])) - annotations = get_reward_deliveries( - np.array([]), trial_outcome_df, ManualWaterTimes(), np.array([0.0]) - ) + annotations = get_reward_deliveries(np.array([]), trial_outcome_df, ManualWaterTimes()) assert isinstance(annotations, np.ndarray) assert annotations.size == 0 @@ -285,39 +292,53 @@ def test_get_reward_deliveries_empty_deliveries_returns_empty(): 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]) + reward_times = np.array([0.15, 1.42]) payload = _outcome_payload(True) trial_outcome_df = pd.DataFrame( {"data": [json.dumps(payload), TrialOutcome.model_validate(payload)]}, - index=pd.Index([1.1, 1.4], name="time"), + index=pd.Index([1.0, 2.0], name="time"), ) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), response_times - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, ManualWaterTimes()) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) -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])) +def test_get_reward_deliveries_rejects_an_empty_trial_table(): + """Deliveries cannot be matched when there are no trials to match them to.""" + empty = _trial_outcome_df(np.array([])) + + with pytest.raises(ValueError, match="trial_outcome_df is empty"): + get_reward_deliveries(np.array([0.15]), empty, ManualWaterTimes()) + + +def test_get_reward_deliveries_rejects_a_nan_trial_index(): + """A ``NaN`` trial boundary is rejected rather than silently swallowing deliveries. + + ``searchsorted`` against a ``NaN``-bearing index would quietly charge every + delivery to the last trial, which is how an upstream stream failure turns + into wrong annotations instead of an error. + """ + trial_outcome_df = _trial_outcome_df(np.array([1.0, np.nan])) + + with pytest.raises(ValueError, match="contains NaN"): + get_reward_deliveries(np.array([0.15]), trial_outcome_df, ManualWaterTimes()) - with pytest.raises(ValueError, match="paired by position"): - get_reward_deliveries( - np.array([0.15]), trial_outcome_df, ManualWaterTimes(), np.array([0.1]) - ) + +def test_get_reward_deliveries_rejects_an_unsorted_trial_index(): + """Trial boundaries must be ordered for containment to be meaningful.""" + trial_outcome_df = _trial_outcome_df(np.array([2.0, 1.0])) + + with pytest.raises(ValueError, match="must be sorted"): + get_reward_deliveries(np.array([0.15]), trial_outcome_df, ManualWaterTimes()) def test_get_reward_deliveries_returns_one_annotation_per_delivery(): """The result is a :class:`numpy.ndarray` aligned with the input deliveries.""" trial_outcome_df = _trial_outcome_df(np.array([1.0, 1.5])) - reward_times = np.array([0.1, 0.2, 0.6]) + reward_times = np.array([0.1, 0.2, 1.2]) - annotations = get_reward_deliveries( - reward_times, trial_outcome_df, ManualWaterTimes(), np.array([0.0, 0.5]) - ) + annotations = get_reward_deliveries(reward_times, trial_outcome_df, ManualWaterTimes()) assert isinstance(annotations, np.ndarray) assert annotations.shape == reward_times.shape