diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index c69f3db..f1725e7 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -30,6 +30,42 @@ Earlier mapping used `Response.json` (`SoftwareEvents`) for lick times (where Lick times now come from the `Behavior/Lickometer` stream, and reward delivery times use the Harp valve open times. +### Reward-delivery annotations + +Each reward-delivery timestamp carries a label in the series' `data` field: + +| Label | Meaning | +| --- | --- | +| `earned` | Water the animal worked for: the matched trial has no free water (`is_auto_reward_right` is `None`). | +| `auto` | Free water: the matched trial has `is_auto_reward_right` set. Scheduled autowater and the anti-bias intervention share that channel and are **not** split here — `auto_waterL` / `auto_waterR` and `anti_bias_left_water` / `anti_bias_right_water` record the mechanism per trial. | +| `manual` | The delivery is the closest valve opening to a `GiveManualWater` software event for this port. Takes precedence over the other labels, since manual water is not aligned to a go cue. | + +Two properties of this series are worth stating explicitly, because both differ +from "every time the valve opened": + +**Deliveries are matched to trials by the `Response` timestamp**, not the +`TrialOutcome` timestamp. `TrialOutcome` fires at the *end* of a trial, after the +reward-consumption and ITI periods, so a delivery can sit nearer the *previous* +trial's outcome and inherit its `is_auto_reward_right`. The valve opens within +milliseconds of the response, so the response anchors each delivery to its own +trial. Verified on `864253_2026-08-11_12-52-08`, `864253_2026-08-04_12-54-55`, +and `864253_2026-08-10_13-09-30`: nearest-`Response` agrees with the trial whose +`[quiescent_start_time, ITI_start_time)` window contains the delivery on all 859 +valve openings, and every `earned` delivery follows a lick on that same port +(median 2.8 ms, max 32 ms). + +**The series is reward-keyed**: a delivery whose matched trial reports +`is_rewarded=False` is dropped rather than annotated, so the retained count +equals the metadata mapper's `sum(is_rewarded)` (248 / 298 / 288 on those three +sessions, against 261 / 298 / 300 valve openings). Manual water is +experimenter-driven, unrelated to the trial outcome, and is never dropped. 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 free water. The consequence is that free water delivered +on a trial that did not pay out is real water the animal received and is **not** +in this series; the trials table's ungated `auto_water*` /`anti_bias_*` columns +still record it. + ## Trials Table Columns are grouped by the raw source they map from. @@ -61,8 +97,8 @@ 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` | **Scheduled autowater only**: `1` when `trial.metadata.extra.is_autowater` is `True` **and** `is_auto_reward_right` points to that side. `0` otherwise, including when the trial's free water came from the anti-bias algorithm — that is reported by `anti_bias_left_water` / `anti_bias_right_water`. `is_auto_reward_right` is only the delivery *channel* (free water fired, and to which side); the mechanism comes from the metadata, so the two columns are mutually exclusive. Not gated on `is_rewarded`: the column records what the task did, and free water fires at the go cue regardless of how the animal's own choice resolves. Note this is narrower than the legacy `dynamic-foraging-task` column of the same name, which was the ungated channel ("Autowater given at Left", straight from `B_AutoWaterTrial`) and predates anti-bias water. | +| `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 scheduled autowater, so the `is_bias_water_intervention` flag is what distinguishes it and the two columns are mutually exclusive. `False` otherwise. Like `auto_water*`, **not** gated on `is_rewarded`: these columns record what the algorithm did, and the intervention fires at the go cue regardless of how the animal's own choice resolves. The reward-delivery series *is* reward-keyed, so this column can exceed the series' `auto` count. | | `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`. | @@ -70,7 +106,7 @@ Columns are grouped by the raw source they map from. | `reward_probabilityL` / `reward_probabilityR` | The **block** probability from `Trial -> metadata -> p_reward_left` / `p_reward_right`. The top-level `trial.p_reward_left` / `p_reward_right` is the per-trial probability, not the block probability, so it is not used here. `None` when the trial or its metadata is missing. | | `reward_size_left` | `Trial -> reward_size.left` — the reward volume (uL) at the left port. Defaults to `2.0` when not set on the trial. `None` when the trial is missing. | | `reward_size_right` | `Trial -> reward_size.right` — the reward volume (uL) at the right port. Defaults to `2.0` when not set on the trial. `None` when the trial is missing. | -| `rewarded_historyL` / `rewarded_historyR` | **Earned** reward only: filter `is_rewarded == True`, then on `is_right_choice`. `False` on both sides when `is_auto_reward_right` is set (either side) — that trial's water is autowater and is reported by `auto_waterL` / `auto_waterR`. | +| `rewarded_historyL` / `rewarded_historyR` | **Earned** reward only: filter `is_rewarded == True`, then on `is_right_choice`. `False` on both sides when `is_auto_reward_right` is set (either side) — that trial's water is free water, reported by `auto_waterL` / `auto_waterR` (scheduled autowater) or `anti_bias_left_water` / `anti_bias_right_water` (anti-bias intervention). This matches the `earned` / `auto` split in the reward-delivery series: per side, `rewarded_history*` equals that series' `earned` count exactly. | ### From `TrialGeneratorSpec.json` (`SoftwareEvents` stream) @@ -191,3 +227,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. Verified 0-mismatch against trial-window containment on 859 valve openings across three sessions. | +| 2026-08-17 | The reward-delivery series is now reward-keyed: a delivery whose matched trial reports `is_rewarded=False` is dropped rather than annotated (manual water exempt), so the retained count equals the metadata mapper's `sum(is_rewarded)`. Free water fires at the go cue and the trial then continues normally, so `is_rewarded` describes the animal's own choice, not the free water — meaning free water on an unrewarded trial is real water the animal received and is excluded from this series. The trials table's `auto_water*` / `anti_bias_*` columns remain ungated and still record it. | +| 2026-08-17 | `auto_waterL` / `auto_waterR` now read `trial.metadata.extra.is_autowater` rather than the `is_auto_reward_right` channel, making them **scheduled autowater only** and mutually exclusive with `anti_bias_left_water` / `anti_bias_right_water`. `is_auto_reward_right` says free water fired and on which side but not what kind; the mechanism is in the metadata. Neither column is gated on `is_rewarded`, since both record what the task did. This is narrower than the legacy `dynamic-foraging-task` column of the same name, which was the ungated channel and predates anti-bias water. | +| 2026-08-17 | The reward-delivery labels stay `earned` / `auto` / `manual`: free water is `auto` whatever mechanism produced it, so the series does not split scheduled autowater from anti-bias water. That split lives in the trials table. Consequence: the series' `auto` count tracks the channel while `auto_waterL` / `auto_waterR` track `is_autowater`, so the two are not expected to be equal. | diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index ebd5270..c46acef 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 ------- @@ -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, @@ -181,8 +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_annotated_rewards`. + 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 ---------- @@ -193,6 +216,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). @@ -211,10 +237,11 @@ 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_reward_deliveries( delivery_times, trial_outcomes, manual_water_times, + response_times, ) return AcquisitionSeries( name=name, @@ -248,9 +275,10 @@ 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() acquisition_streams = self.loader.get_all_raw_data() acqusition_streams_descriptions = self.loader.raw_data_stream_descriptions @@ -274,6 +302,7 @@ def build_acquisition( rewards, trial_outcomes, manual_water, + response_times, port_column="SupplyPort0", is_right=False, name="left_reward_delivery_time", @@ -285,6 +314,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/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 724828b..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__) @@ -353,7 +354,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``). @@ -442,28 +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, *, 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 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. - 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: + 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 classify autowater and anti-bias water identically. Parameters ---------- @@ -476,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( @@ -497,6 +512,13 @@ def _anti_bias_water( was a bias-water intervention *and* the auto-response was to the requested side. + 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 @@ -897,8 +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, is_right=False), - auto_waterR=self._auto_water(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 22e0e96..94b0f47 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -189,20 +189,30 @@ 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=( + "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=( + "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." + ), + ) # --- 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/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 02234a8..af7048a 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -4,7 +4,7 @@ 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 @@ -30,12 +30,39 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: return TrialOutcome.model_validate(payload) -def get_annotated_rewards( +def _free_water_label(trial: t.Optional[Trial]) -> str: + """Classify a delivery's trial as ``auto`` (free water) or ``earned``. + + ``is_auto_reward_right`` triggers an immediate reward to one side, so any + trial with it set gave free water rather than water the animal worked for. + Scheduled autowater and the anti-bias water intervention share that channel + and are both ``auto`` here; which mechanism gave the water is recorded per + trial by ``auto_waterL``/``auto_waterR`` and + ``anti_bias_left_water``/``anti_bias_right_water`` in the trials table. + + Parameters + ---------- + trial : Trial or None + The per-trial task-logic model, or ``None`` when the outcome payload was + missing. + + Returns + ------- + str + ``"auto"`` when the trial delivered free water, else ``"earned"``. + """ + if trial is None or trial.is_auto_reward_right is None: + return "earned" + return "auto" + + +def get_reward_deliveries( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, manual_water_times: np.ndarray, -) -> np.ndarray: - """Annotate each reward delivery as ``earned``, ``auto``, or ``manual``. + response_times: np.ndarray, +) -> 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 @@ -45,12 +72,37 @@ 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`. - - ``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. + - ``auto`` -- otherwise, when the trial delivered free water + (``is_auto_reward_right is not None``). Scheduled autowater and the + anti-bias water intervention are both delivered through that channel, so + both are ``auto`` here; which mechanism gave the water is recorded per + trial by the trials table's ``auto_waterL``/``auto_waterR`` and + ``anti_bias_left_water``/``anti_bias_right_water``. + - ``earned`` -- otherwise: water the animal worked for. + + 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 + 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,31 +114,46 @@ 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"``. + 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"``, 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) + 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, trial_outcome_df.index.to_numpy() - ) + 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 - if trial is None or trial.is_auto_reward_right is None: - annotated_rewards.append("earned") - else: - annotated_rewards.append("auto") + 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 "manual" and would truncate it. @@ -97,8 +164,14 @@ 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" + + # 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] 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() diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index f16915e..be0f544 100644 --- a/tests/test_nwb/test_acquisition/test_acquisition_builder.py +++ b/tests/test_nwb/test_acquisition/test_acquisition_builder.py @@ -58,7 +58,11 @@ def _make_output_set_frame() -> pd.DataFrame: def _outcome_payload(auto) -> dict: - """Return a serialized ``TrialOutcome`` payload with the given auto-response.""" + """Return a serialized ``TrialOutcome`` payload with the given auto-response. + + Free water is flagged as scheduled autowater in ``metadata.extra`` so the + annotation can attribute it; the channel alone does not name the mechanism. + """ return { "trial": { "p_reward_left": 1.0, @@ -68,6 +72,7 @@ def _outcome_payload(auto) -> dict: "quiescence_period_duration": 0.5, "inter_trial_interval_duration": 4.0, "is_auto_reward_right": auto, + "metadata": {"extra": {"is_autowater": auto is not None}}, }, "is_right_choice": True, "is_rewarded": True, @@ -82,6 +87,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 +133,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 +161,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 +191,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 +239,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_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index b213ba1..f66c0da 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -96,6 +96,7 @@ def _outcome( reward_size_left=None, reward_size_right=None, lickspout_offset_delta=None, + is_autowater=None, is_bias_water_intervention=None, is_bias_stage_intervention=None, ): @@ -106,9 +107,9 @@ def _outcome( probabilities stored under ``trial.metadata`` (the source of the ``reward_probability`` columns). ``reward_size_left`` / ``reward_size_right`` override the default per-trial reward volumes (uL). ``lickspout_offset_delta`` - sets the per-trial horizontal spout displacement (mm), and the - ``is_bias_*_intervention`` flags populate the anti-bias ``metadata.extra`` - (``BlockBasedTrialMetadata``) block. + sets the per-trial horizontal spout displacement (mm), and ``is_autowater`` + plus the ``is_bias_*_intervention`` flags populate the ``metadata.extra`` + (``BlockBasedTrialMetadata``) block naming the free-water mechanism. """ trial = { "p_reward_left": p_left, @@ -128,9 +129,14 @@ def _outcome( } if block_p_left is not None or block_p_right is not None: trial["metadata"] = {"p_reward_left": block_p_left, "p_reward_right": block_p_right} - if is_bias_water_intervention is not None or is_bias_stage_intervention is not None: + if ( + is_autowater is not None + or is_bias_water_intervention is not None + or is_bias_stage_intervention is not None + ): metadata = trial.setdefault("metadata", {}) metadata["extra"] = { + "is_autowater": bool(is_autowater), "is_bias_water_intervention": bool(is_bias_water_intervention), "is_bias_stage_intervention": bool(is_bias_stage_intervention), } @@ -689,17 +695,70 @@ 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.""" + """Scheduled autowater encodes ``1`` on its side and ``0`` on the other.""" trial = TrialOutcome.model_validate( - _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=True) + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=True, is_autowater=True) ).trial - assert TrialTableBuilder._auto_water(trial, is_right=True) == 1 - assert TrialTableBuilder._auto_water(trial, is_right=False) == 0 - # No auto-response counts as no autowater (0). + meta = TrialTableBuilder._bias_metadata(trial) + assert TrialTableBuilder._auto_water(trial, meta, is_right=True) == 1 + assert TrialTableBuilder._auto_water(trial, meta, is_right=False) == 0 + # No free water at all counts as no autowater (0). no_auto = TrialOutcome.model_validate( - _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=None) + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=None, is_autowater=True) ).trial - assert TrialTableBuilder._auto_water(no_auto, is_right=True) == 0 + assert ( + TrialTableBuilder._auto_water( + no_auto, TrialTableBuilder._bias_metadata(no_auto), is_right=True + ) + == 0 + ) + + +def test_auto_water_includes_unrewarded_trials(): + """Autowater on a trial that did not pay out still counts. + + The column records what the task did, and free water fires at the go cue + regardless of how the animal's own choice resolves. The reward-delivery + series is reward-keyed and drops those deliveries, so this column can exceed + that series' ``auto`` count. + """ + unrewarded = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=False, is_rewarded=False, auto=True, is_autowater=True) + ).trial + meta = TrialTableBuilder._bias_metadata(unrewarded) + assert TrialTableBuilder._auto_water(unrewarded, meta, is_right=True) == 1 + + +def test_auto_water_excludes_anti_bias_water(): + """Anti-bias water is reported by ``anti_bias_*_water``, not ``auto_water*``. + + ``is_auto_reward_right`` is only the delivery channel, shared by scheduled + autowater and the anti-bias intervention, so the mechanism comes from the + metadata flags. The two columns are mutually exclusive. + """ + bias_water = TrialOutcome.model_validate( + _outcome( + 1.0, + 1.0, + is_right_choice=True, + is_rewarded=True, + auto=True, + is_bias_water_intervention=True, + ) + ).trial + meta = TrialTableBuilder._bias_metadata(bias_water) + assert TrialTableBuilder._auto_water(bias_water, meta, is_right=True) == 0 + assert TrialTableBuilder._anti_bias_water(bias_water, meta, is_right=True) is True + + +def test_auto_water_excludes_free_water_with_no_mechanism_flag(): + """Free water the metadata flags as neither mechanism is ``0`` on both columns.""" + unflagged = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, auto=True) + ).trial + meta = TrialTableBuilder._bias_metadata(unflagged) + assert TrialTableBuilder._auto_water(unflagged, meta, is_right=True) == 0 + assert TrialTableBuilder._anti_bias_water(unflagged, meta, is_right=True) is False def test_bias_metadata_parses_dict_model_and_default(): @@ -750,18 +809,40 @@ 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, 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. autowater = TrialOutcome.model_validate( _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, auto_meta, is_right=False) is False + + +def test_anti_bias_water_includes_unrewarded_trials(): + """An intervention on a trial that did not pay out still counts. + + 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. The + reward-delivery series is reward-keyed and drops that delivery, so this + column can exceed the series' ``anti_bias`` count. + """ + unrewarded = TrialOutcome.model_validate( + _outcome( + 1.0, + 1.0, + is_right_choice=False, + is_rewarded=False, + auto=True, + is_bias_water_intervention=True, + ) ).trial - auto_meta = TrialTableBuilder._bias_metadata(autowater) - assert TrialTableBuilder._anti_bias_water(autowater, auto_meta, is_right=False) is False + meta = TrialTableBuilder._bias_metadata(unrewarded) + assert TrialTableBuilder._anti_bias_water(unrewarded, meta, is_right=True) is True 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 2be8269..e6172c9 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -4,121 +4,264 @@ 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 - - -def _outcome_payload(auto=None) -> dict: - """Return a serialized ``TrialOutcome`` payload with the given auto-response.""" +from dynamic_foraging_processing.utils.rewards import get_reward_deliveries + + +def _outcome_payload(auto=None, is_rewarded: bool = True, mechanism: str = "autowater") -> dict: + """Return a serialized ``TrialOutcome`` payload with the given auto-response. + + ``mechanism`` names which free-water flag ``metadata.extra`` carries when + ``auto`` is set: ``"autowater"`` for scheduled autowater, ``"anti_bias"`` for + an anti-bias intervention, or ``None`` for neither (the mechanism is not + recorded, so the delivery cannot be attributed). + """ + trial = { + "p_reward_left": 1.0, + "p_reward_right": 1.0, + "response_deadline_duration": 3.0, + "reward_consumption_duration": 1.0, + "quiescence_period_duration": 0.5, + "inter_trial_interval_duration": 4.0, + "is_auto_reward_right": auto, + } + if mechanism is not None: + trial["metadata"] = { + "extra": { + "is_autowater": mechanism == "autowater", + "is_bias_water_intervention": mechanism == "anti_bias", + } + } return { - "trial": { - "p_reward_left": 1.0, - "p_reward_right": 1.0, - "response_deadline_duration": 3.0, - "reward_consumption_duration": 1.0, - "quiescence_period_duration": 0.5, - "inter_trial_interval_duration": 4.0, - "is_auto_reward_right": auto, - }, + "trial": trial, "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, mechanism: str = "autowater" +) -> 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, mechanism) 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. +@pytest.mark.parametrize("mechanism", ["autowater", "anti_bias", None]) +def test_get_reward_deliveries_marks_all_free_water_as_auto(mechanism): + """Every free-water delivery is ``auto``, whatever mechanism gave it. + + ``is_auto_reward_right`` is the delivery channel, shared by scheduled + autowater and the anti-bias intervention, and the series does not split them: + ``auto_waterL``/``auto_waterR`` and + ``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"), + ) + + 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. + + 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_auto_water_on_unrewarded_trials(): + """Autowater on a trial reporting ``is_rewarded=False`` is dropped, not annotated. + + 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]) + 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_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 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]) + 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 + ) + + assert times.size == 0 + assert annotations.size == 0 + + +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])) + + with pytest.raises(ValueError, match="paired by position"): + get_reward_deliveries(np.array([0.15]), trial_outcome_df, np.array([]), np.array([0.1])) + - result = get_annotated_rewards(reward_times, trial_outcome_df, np.array([])) +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)