diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index a63fda2..ebd5270 100644 --- a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py +++ b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py @@ -182,7 +182,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 ---------- @@ -223,7 +223,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" ), ) diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index 068a144..02234a8 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, ) -> 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,10 +45,13 @@ 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 ``TrialOutcome`` software-event + timestamp: each delivery takes the annotation of the closest trial. + Parameters ---------- reward_delivery_times : numpy.ndarray @@ -64,7 +67,7 @@ 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"``. """ reward_times = np.asarray(reward_delivery_times) if reward_times.size == 0: @@ -83,9 +86,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 diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index 4b1f91f..f16915e 100644 --- a/tests/test_nwb/test_acquisition/test_acquisition_builder.py +++ b/tests/test_nwb/test_acquisition/test_acquisition_builder.py @@ -175,7 +175,9 @@ 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()), + } ), } ) @@ -265,7 +267,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..2be8269 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -27,7 +27,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]}, @@ -45,14 +45,26 @@ def test_get_annotated_rewards_marks_default_trials_as_earned(): 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]) trial_outcome_df = _trial_outcome_df(np.array([0.1, 0.4]), autos=[True, False]) annotations = get_annotated_rewards(reward_times, trial_outcome_df, 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_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, np.array([])) + + np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) def test_get_annotated_rewards_marks_manual_water_as_manual(): @@ -67,7 +79,7 @@ def test_get_annotated_rewards_marks_manual_water_as_manual(): 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]) @@ -99,7 +111,7 @@ def test_get_annotated_rewards_accepts_json_and_model_payloads(): annotations = get_annotated_rewards(reward_times, trial_outcome_df, 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():