From f4b27012339f8d1efec0dfe99612db3612914c5a Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 <109561860+arjunsridhar12345@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:19:50 -0700 Subject: [PATCH] feat: ingest new manual water software events (#102) * feat: use new software events for manual water * feat: update qc plot for manual and manual aligned go cue * build: update acquisition dependency * test: update tests * fix: lint * docs: update trial table docs * docs: refactor wording to be more smooth --- docs/trials_table_mapping.md | 8 +- pyproject.toml | 2 +- .../nwb/acquisition/acquisition_builder.py | 92 +++++++++------ .../pipeline/_pipeline.py | 65 ++++++---- .../qc/processed/plots.py | 69 +++++++---- .../qc/processed/results.py | 16 +-- .../qc/processed/stage.py | 14 ++- .../utils/rewards.py | 106 +++++++++++++---- .../test_acquisition_builder.py | 105 +++++++++++------ tests/test_pipeline/test_pipeline.py | 40 ++++--- tests/test_qc/test_plots.py | 67 +++++++++-- tests/test_utils/test_rewards.py | 111 +++++++++++++++--- uv.lock | 8 +- 13 files changed, 497 insertions(+), 206 deletions(-) diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index e505f60..6763d97 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -38,7 +38,12 @@ Each reward-delivery timestamp carries a label in the series' `data` field: | --- | --- | | `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. | +| `manual_go_cue_aligned` | The delivery is the closest valve opening to a `LeftManualAutoReward` / `RightManualAutoReward` software event for this port: water the *experimenter* triggered to land on the go cue. It fires at the go cue like autowater, but the task did not schedule it, so it is neither `auto` nor `earned`. Takes precedence over both trial-derived labels. | +| `manual` | The delivery is the closest valve opening to a `LeftManualWater` / `RightManualWater` software event for this port: experimenter water given at an arbitrary moment, not tied to a go cue. Highest precedence of all four. | + +The side of an experimenter-water event comes from the **stream name**, not from +an event payload. Each of the four streams exists only when the experimenter gave +water of that kind, so a session with none of them is normal. Two properties of this series are worth stating explicitly, because both differ from "every time the valve opened": @@ -230,3 +235,4 @@ These were mapped during exploration but are no longer in scope: | 2026-08-20 | `block_max` is now one below `block_length`'s configured maximum, which accounts for the floor applied upstream: a block is a whole number of trials, so the configured bound is never itself reachable. `block_min`, `block_beta`, and the `ITI_*` / `delay_*` bounds are unchanged — those durations are continuous and take no such adjustment. | | 2026-08-20 | `ITI_min` now reports `inter_trial_interval_duration`'s scaling `offset` instead of its truncation minimum: the sampled ITI is shifted by the offset, so the offset is the shortest ITI the generator can produce. Falls back to the truncation minimum when no scaling parameters are configured. | | 2026-08-20 | `bait_left` / `bait_right` now read `trial.metadata.extra.is_left_baited` / `is_right_baited` from the acquisition software instead of being re-derived from `p_reward_left` / `p_reward_right` and the `is_auto_reward_right` channel. The software is the authority on bait state, so the two can disagree — notably a port with `p_reward == 1` is no longer assumed baited. `False` when the trial carries no extra metadata. | +| 2026-09-15 | **Breaking:** experimenter water is now read from the four side-specific software-event streams the acquisition software emits (`LeftManualWater` / `RightManualWater`, not aligned to a go cue, and `LeftManualAutoReward` / `RightManualAutoReward`, aligned to it) instead of the single `GiveManualWaterRight` stream whose `data` payload selected the side. The `GiveManualWaterRight` path is removed, not deprecated. A fourth reward-delivery label, `manual_go_cue_aligned`, joins `earned` / `auto` / `manual`. **This fixes a mislabel:** manual auto-rewards fire at the go cue but leave `is_auto_reward_right` unset, so they previously fell through to `earned` — water the animal never worked for, counted as earned. The QC `side_bias.png` behavior raster gains an `L` / `R Manual Water (go cue)` row per side (now 11 rows, y-limits `[-0.8, 1.8]`), drawn dotted where unaligned manual water is dashed. | diff --git a/pyproject.toml b/pyproject.toml index 52d16a7..7054454 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ readme = "README.md" version = "0.3.8" dependencies = [ - "aind-behavior-dynamic-foraging[data] @ git+https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git@933a7aee2627de1979a982847f3553cf7435cbd2", + "aind-behavior-dynamic-foraging[data] @ git+https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git@4e4faa4941c8d9cc8024e760fdaec30339b626dc", "ipykernel", ] diff --git a/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py b/src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py index 6fd4f1f..be51f43 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_reward_deliveries +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes, get_reward_deliveries class LickSource(t.NamedTuple): @@ -96,28 +96,54 @@ def get_response_times(self) -> np.ndarray: ) return responses.index.to_numpy() - def get_manual_water_times(self) -> pd.DataFrame: - """Get the manual-water software-event stream. + def _software_event_times(self, stream_name: str) -> np.ndarray: + """Get one ``Behavior/SoftwareEvents`` stream's event timestamps. + + Parameters + ---------- + stream_name : str + The software-event stream to read (e.g. ``"LeftManualWater"``). Returns ------- - pandas.DataFrame - The ``GiveManualWaterRight`` stream under ``Behavior/SoftwareEvents``, - indexed by event timestamp with a ``data`` column that is ``True`` - for right-port manual water and ``False`` for left-port manual water. - An empty frame (with a ``data`` column) is returned when the stream - is absent. + numpy.ndarray + The stream's event timestamps, or an empty array when the stream is + absent. Only the timestamps are used; these events carry no payload + this pipeline reads. """ try: - return ( - self.loader.dataset.at("Behavior") - .at("SoftwareEvents") - .at("GiveManualWaterRight") - .load() - .data + data = ( + self.loader.dataset.at("Behavior").at("SoftwareEvents").at(stream_name).load().data ) except (KeyError, FileNotFoundError): - return pd.DataFrame({"data": []}) + return np.array([]) + return data.index.to_numpy() + + def get_manual_water_times(self, *, is_right: bool) -> ManualWaterTimes: + """Get one lick port's experimenter-triggered water times. + + The acquisition software emits four side-specific streams: + ``{Left,Right}ManualWater`` for water given at an arbitrary moment and + ``{Left,Right}ManualAutoReward`` for water triggered to land on the go + cue. The side comes from the stream name, so no payload inspection is + needed. Each stream is optional -- a session where the experimenter gave + no water of that kind has no file -- and reads as an empty array. + + Parameters + ---------- + is_right : bool + ``True`` for the right lick port, ``False`` for the left. + + Returns + ------- + ManualWaterTimes + This port's ``unaligned`` and ``go_cue_aligned`` event timestamps. + """ + side = "Right" if is_right else "Left" + return ManualWaterTimes( + unaligned=self._software_event_times(f"{side}ManualWater"), + go_cue_aligned=self._software_event_times(f"{side}ManualAutoReward"), + ) def get_lick_times(self, device: str, stream_name: str, port: str) -> np.ndarray: """Get the lick times for one lick port from a Harp digital-input stream. @@ -191,20 +217,20 @@ def _reward_delivery_series( self, writes: pd.DataFrame, trial_outcomes: pd.DataFrame, - manual_water: pd.DataFrame, + manual_water: ManualWaterTimes, response_times: np.ndarray, *, port_column: str, - is_right: bool, name: str, side_label: str, ) -> AcquisitionSeries: """Build one lick port's reward-delivery series with reward annotations. Only valve-open events (``port_column`` is truthy) are reward - deliveries; the ``data`` field annotates each as earned, manual, or auto - via :func:`get_reward_deliveries`. Every valve opening is reported, so - the series is a complete record of the water delivered at this port. + deliveries; the ``data`` field annotates each as earned, auto, manual, or + manual-go-cue-aligned via :func:`get_reward_deliveries`. Every valve + opening is reported, so the series is a complete record of the water + delivered at this port. Parameters ---------- @@ -212,17 +238,15 @@ def _reward_delivery_series( ``OutputSet`` ``WRITE`` messages indexed by timestamp. trial_outcomes : pandas.DataFrame The ``TrialOutcome`` stream, indexed by trial timestamp. - manual_water : pandas.DataFrame - The ``GiveManualWaterRight`` stream; the ``data`` column selects the - side (``True`` right, ``False`` left). + 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). - is_right : bool - ``True`` for the right lick port, ``False`` for the left. name : str Acquisition series name. side_label : str @@ -235,11 +259,10 @@ def _reward_delivery_series( """ open_writes = writes[writes[port_column].fillna(False).astype(bool)] delivery_times = open_writes.index.to_numpy() - manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy() annotations = get_reward_deliveries( delivery_times, trial_outcomes, - manual_water_times, + manual_water, response_times, ) return AcquisitionSeries( @@ -249,7 +272,9 @@ def _reward_delivery_series( unit="second", description=( f"The reward delivery time of the {side_label} lick port. The data field " - "annotates whether the reward was earned, manual, or auto" + "annotates whether the reward was earned, auto (task-triggered free water), " + "manual (experimenter water not aligned to a go cue), or " + "manual_go_cue_aligned (experimenter water delivered at the go cue)" ), ) @@ -276,7 +301,8 @@ def build_acquisition( """ rewards = self.get_valve_writes() trial_outcomes = self.get_trial_outcomes() - manual_water = self.get_manual_water_times() + 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() @@ -300,10 +326,9 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, - manual_water, + left_manual_water, response_times, port_column="SupplyPort0", - is_right=False, name="left_reward_delivery_time", side_label="left", ) @@ -312,10 +337,9 @@ def build_acquisition( self._reward_delivery_series( rewards, trial_outcomes, - manual_water, + right_manual_water, response_times, port_column="SupplyPort1", - is_right=True, name="right_reward_delivery_time", side_label="right", ) diff --git a/src/dynamic_foraging_processing/pipeline/_pipeline.py b/src/dynamic_foraging_processing/pipeline/_pipeline.py index 6c6f013..843d3b8 100644 --- a/src/dynamic_foraging_processing/pipeline/_pipeline.py +++ b/src/dynamic_foraging_processing/pipeline/_pipeline.py @@ -48,6 +48,11 @@ from dynamic_foraging_processing.processing import TrialConfig, TrialTableBuilder from dynamic_foraging_processing.qc import ProcessedQC, RawQC, build_quality_control from dynamic_foraging_processing.raw_data_loader import RawDataLoader +from dynamic_foraging_processing.utils.rewards import ( + MANUAL, + MANUAL_GO_CUE_ALIGNED, + ManualWaterTimes, +) #: Default lick-port sources on the standard behavior board. _DEFAULT_LEFT_LICK = LickSource("HarpBehavior", "DigitalInputState", "DIPort0") @@ -84,9 +89,6 @@ _LEFT_REWARD_SERIES = "left_reward_delivery_time" _RIGHT_REWARD_SERIES = "right_reward_delivery_time" -#: Reward-delivery annotation marking a manual-water event. -_MANUAL_ANNOTATION = "manual" - class Pipeline: """Package a raw dynamic foraging acquisition to NWB and run QC. @@ -345,7 +347,7 @@ def _write_processing( # ------------------------------------------------------------------ # def _read_processed_inputs( self, nwb_file: pynwb.NWBFile - ) -> t.Tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + ) -> t.Tuple[pd.DataFrame, np.ndarray, np.ndarray, ManualWaterTimes, ManualWaterTimes]: """Read the processed-QC inputs from an NWB file. Parameters @@ -356,25 +358,38 @@ def _read_processed_inputs( Returns ------- tuple - ``(trials, left_lick_times, right_lick_times, manual_left_times, - manual_right_times)``. + ``(trials, left_lick_times, right_lick_times, manual_left, + manual_right)``, where the last two are :class:`ManualWaterTimes`. """ trials = nwb_file.trials.to_dataframe() left_lick_times = np.asarray(nwb_file.acquisition[_LEFT_LICK_SERIES].timestamps) right_lick_times = np.asarray(nwb_file.acquisition[_RIGHT_LICK_SERIES].timestamps) - manual_left_times, manual_right_times = self._manual_water_times(nwb_file) - return trials, left_lick_times, right_lick_times, manual_left_times, manual_right_times + manual_left, manual_right = self._manual_water_times(nwb_file) + return trials, left_lick_times, right_lick_times, manual_left, manual_right @staticmethod - def _manual_water_times(nwb_file: pynwb.NWBFile) -> t.Tuple[np.ndarray, np.ndarray]: - """Return the ``(left, right)`` manual-water delivery times from the NWB. + def _manual_water_times( + nwb_file: pynwb.NWBFile, + ) -> t.Tuple[ManualWaterTimes, ManualWaterTimes]: + """Return the ``(left, right)`` experimenter-water times from the NWB. - Manual-water deliveries are the reward-delivery events annotated - ``"manual"`` on each side's ``*_reward_delivery_time`` acquisition series. + Each side's times are the reward-delivery events annotated ``"manual"`` + (not aligned to a go cue) and ``"manual_go_cue_aligned"`` on that side's + ``*_reward_delivery_time`` acquisition series. The two are read back + separately so the QC figure can keep them on their own rows. """ - left = Pipeline._annotated_times(nwb_file, _LEFT_REWARD_SERIES, _MANUAL_ANNOTATION) - right = Pipeline._annotated_times(nwb_file, _RIGHT_REWARD_SERIES, _MANUAL_ANNOTATION) - return left, right + return ( + Pipeline._side_manual_water_times(nwb_file, _LEFT_REWARD_SERIES), + Pipeline._side_manual_water_times(nwb_file, _RIGHT_REWARD_SERIES), + ) + + @staticmethod + def _side_manual_water_times(nwb_file: pynwb.NWBFile, series_name: str) -> ManualWaterTimes: + """Split one reward-delivery series' experimenter-water times by alignment.""" + return ManualWaterTimes( + unaligned=Pipeline._annotated_times(nwb_file, series_name, MANUAL), + go_cue_aligned=Pipeline._annotated_times(nwb_file, series_name, MANUAL_GO_CUE_ALIGNED), + ) @staticmethod def _annotated_times(nwb_file: pynwb.NWBFile, series_name: str, annotation: str) -> np.ndarray: @@ -389,8 +404,8 @@ def _assemble_quality_control( trials: pd.DataFrame, left_lick_times: np.ndarray, right_lick_times: np.ndarray, - manual_left_times: np.ndarray, - manual_right_times: np.ndarray, + manual_left: ManualWaterTimes, + manual_right: ManualWaterTimes, results_folder: t.Optional[str] = None, ) -> QualityControl: """Run the raw and processed QC stages and assemble one ``QualityControl``. @@ -401,9 +416,9 @@ def _assemble_quality_control( The trials table, consumed by the processed (behavior) QC stage. left_lick_times, right_lick_times : numpy.ndarray Left/right-port lick times for the processed QC stage. - manual_left_times, manual_right_times : numpy.ndarray - Left/right manual-water delivery times passed through to the side-bias - figure. + manual_left, manual_right : ManualWaterTimes + Left/right experimenter-water delivery times, split into unaligned + and go-cue-aligned, passed through to the side-bias figure. results_folder : str, optional Directory to write figure assets into so the metric references resolve. If ``None``, assets are skipped. @@ -420,8 +435,8 @@ def _assemble_quality_control( left_lick_times, right_lick_times, results_folder, - manual_left_times=manual_left_times, - manual_right_times=manual_right_times, + manual_left=manual_left, + manual_right=manual_right, ) return build_quality_control([*raw_metrics, *processed_metrics]) @@ -486,7 +501,7 @@ def run_qc( ``output_path`` is given, the QC JSON and figure assets are written to disk as a side effect. """ - trials, left_lick_times, right_lick_times, manual_left_times, manual_right_times = ( + trials, left_lick_times, right_lick_times, manual_left, manual_right = ( self._read_processed_inputs(nwb_file) ) @@ -500,8 +515,8 @@ def run_qc( trials, left_lick_times, right_lick_times, - manual_left_times, - manual_right_times, + manual_left, + manual_right, results_folder, ) if output_path is not None: diff --git a/src/dynamic_foraging_processing/qc/processed/plots.py b/src/dynamic_foraging_processing/qc/processed/plots.py index 65bd0d8..fba2438 100644 --- a/src/dynamic_foraging_processing/qc/processed/plots.py +++ b/src/dynamic_foraging_processing/qc/processed/plots.py @@ -21,6 +21,13 @@ SIDE_BIAS_PLOT, lick_latency_by_side, ) +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes + +#: Line styles separating the two kinds of experimenter water. Both rows are +#: broken lines so neither is confused with the solid earned/auto ticks, and the +#: two differ from each other so the rows stay distinguishable in grayscale. +_MANUAL_STYLE = "dashed" +_MANUAL_GO_CUE_STYLE = "dotted" #: Vertical offset (in side-bias units) of the anti-bias lickspout-move markers #: from the zero-bias line: rightward moves sit this far above it, leftward moves @@ -334,8 +341,8 @@ def _add_behavior_plot( rewarded_right: t.Optional[np.ndarray], autowater_left: t.Optional[np.ndarray], autowater_right: t.Optional[np.ndarray], - manual_left_times: t.Optional[np.ndarray], - manual_right_times: t.Optional[np.ndarray], + manual_left: ManualWaterTimes, + manual_right: ManualWaterTimes, go_cue_times: t.Optional[np.ndarray], ) -> None: """Draw the per-trial behavior raster (choices, rewards, water). @@ -343,6 +350,11 @@ def _add_behavior_plot( Each water type gets its own row per side: manual water is not autowater, so sharing a band with it (as an earlier version did) made the two indistinguishable and left manual deliveries reading as mislabeled autowater. + + For the same reason the two kinds of experimenter water get separate rows. + Go-cue-aligned manual water fires at the go cue exactly as task-scheduled + autowater does, so plotting it on the manual row (or on the autowater row) + would hide which of the two gave the water. """ choices = np.asarray(animal_response) ax.vlines(np.where(choices == 1)[0], 0.8, 1, linewidth=1, color="gray", label="Choice") @@ -374,17 +386,24 @@ def _add_behavior_plot( color="cyan", ) - # Manual water sits outside the autowater rows and is dashed, so it reads as - # distinct from the solid earned and auto ticks even where colour alone is - # hard to judge. It is labelled only when the session actually has - # deliveries, so the legend never claims manual water for a session that had - # none. - manual_label: t.Optional[str] = "Manual Water" - for times, bottom, top in ( - (manual_right_times, 1.4, 1.6), - (manual_left_times, -0.6, -0.4), + # Experimenter water sits outside the autowater rows and is drawn with broken + # lines, so it reads as distinct from the solid earned and auto ticks even + # where colour alone is hard to judge. The two kinds are told apart by dash + # pattern as well as by row. Each kind is labelled only when the session + # actually has deliveries of it, so the legend never claims water the session + # never got; labels are popped so a kind appears at most once across the two + # sides. + remaining_labels = { + _MANUAL_STYLE: "Manual Water", + _MANUAL_GO_CUE_STYLE: "Manual Water (go cue aligned)", + } + for times, bottom, top, linestyle in ( + (manual_right.go_cue_aligned, 1.6, 1.8, _MANUAL_GO_CUE_STYLE), + (manual_right.unaligned, 1.4, 1.6, _MANUAL_STYLE), + (manual_left.unaligned, -0.6, -0.4, _MANUAL_STYLE), + (manual_left.go_cue_aligned, -0.8, -0.6, _MANUAL_GO_CUE_STYLE), ): - if times is None or go_cue_times is None: + if go_cue_times is None: continue trial_indices = _time_to_trial_index(go_cue_times, times) if not trial_indices: @@ -395,17 +414,17 @@ def _add_behavior_plot( top, linewidth=1, color="blue", - linestyles="dashed", - label=manual_label, + linestyles=linestyle, + label=remaining_labels.pop(linestyle, None), ) - manual_label = None - ax.set_ylim([-0.6, 1.6]) + ax.set_ylim([-0.8, 1.8]) ax.set_xlim([_TRIAL_AXIS_LEFT, len(choices)]) ax.set_xlabel("Trial #") ax.set_yticks( - [-0.5, -0.3, -0.1, 0.1, 0.5, 0.9, 1.1, 1.3, 1.5], + [-0.7, -0.5, -0.3, -0.1, 0.1, 0.5, 0.9, 1.1, 1.3, 1.5, 1.7], labels=[ + "L Manual Water (go cue)", "L Manual Water", "L Auto Water", "L Reward", @@ -415,6 +434,7 @@ def _add_behavior_plot( "R Reward", "R Auto Water", "R Manual Water", + "R Manual Water (go cue)", ], ) _legend_outside(ax) @@ -451,8 +471,8 @@ def plot_side_bias( go_cue_times: t.Optional[np.ndarray] = None, autowater_left: t.Optional[np.ndarray] = None, autowater_right: t.Optional[np.ndarray] = None, - manual_left_times: t.Optional[np.ndarray] = None, - manual_right_times: t.Optional[np.ndarray] = None, + manual_left: ManualWaterTimes = ManualWaterTimes(), + manual_right: ManualWaterTimes = ManualWaterTimes(), anti_bias_left_water: t.Optional[np.ndarray] = None, anti_bias_right_water: t.Optional[np.ndarray] = None, anti_bias_lickspout_movement: t.Optional[np.ndarray] = None, @@ -477,11 +497,12 @@ def plot_side_bias( reward_probability_left, reward_probability_right : numpy.ndarray, optional Per-trial reward probabilities. go_cue_times : numpy.ndarray, optional - Go-cue timestamps (s), used to map manual-water times to trials. + Go-cue timestamps (s), used to map experimenter-water times to trials. autowater_left, autowater_right : numpy.ndarray, optional Per-trial autowater indicator arrays. - manual_left_times, manual_right_times : numpy.ndarray, optional - Manual-water delivery timestamps (s). + manual_left, manual_right : ManualWaterTimes, optional + Left/right experimenter-water delivery timestamps (s), split into + ``unaligned`` and ``go_cue_aligned``; each gets its own raster row. anti_bias_left_water, anti_bias_right_water : numpy.ndarray, optional Boolean per-trial arrays flagging anti-bias water interventions on each side; overlaid on the side-bias trace. @@ -524,8 +545,8 @@ def plot_side_bias( rewarded_right, autowater_left, autowater_right, - manual_left_times, - manual_right_times, + manual_left, + manual_right, go_cue_times, ) _add_reward_probabilities(ax[3], reward_probability_left, reward_probability_right) diff --git a/src/dynamic_foraging_processing/qc/processed/results.py b/src/dynamic_foraging_processing/qc/processed/results.py index 1c708b3..692f432 100644 --- a/src/dynamic_foraging_processing/qc/processed/results.py +++ b/src/dynamic_foraging_processing/qc/processed/results.py @@ -25,6 +25,7 @@ plot_lick_latency, plot_side_bias, ) +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes # Logical input -> trials-table column name. Centralized so the mapping is easy # to correct against the trial-table builder. The lickspout columns are the @@ -67,8 +68,8 @@ def behavior_qc_results( right_lick_times: np.ndarray, results_folder: t.Optional[str] = None, *, - manual_left_times: t.Optional[np.ndarray] = None, - manual_right_times: t.Optional[np.ndarray] = None, + manual_left: ManualWaterTimes = ManualWaterTimes(), + manual_right: ManualWaterTimes = ManualWaterTimes(), ) -> t.List[QCResult]: """Build the behavior QC results (side bias + lick intervals). @@ -90,9 +91,10 @@ def behavior_qc_results( per-trial, so they are passed explicitly rather than read from ``trials``. results_folder : str, optional Directory to write the plots into. If ``None``, plots are skipped. - manual_left_times, manual_right_times : numpy.ndarray, optional - Manual-water delivery timestamps (s); event-time arrays passed through - to the side-bias figure. + manual_left, manual_right : ManualWaterTimes, optional + Left/right experimenter-water delivery timestamps (s), split into + unaligned and go-cue-aligned; event-time arrays passed through to the + side-bias figure. Returns ------- @@ -124,8 +126,8 @@ def behavior_qc_results( go_cue_times=_column(trials, "go_cue_times"), autowater_left=_column(trials, "autowater_left"), autowater_right=_column(trials, "autowater_right"), - manual_left_times=manual_left_times, - manual_right_times=manual_right_times, + manual_left=manual_left, + manual_right=manual_right, anti_bias_left_water=_column(trials, "anti_bias_left_water"), anti_bias_right_water=_column(trials, "anti_bias_right_water"), anti_bias_lickspout_movement=_column(trials, "anti_bias_lickspout_movement"), diff --git a/src/dynamic_foraging_processing/qc/processed/stage.py b/src/dynamic_foraging_processing/qc/processed/stage.py index 18d5f10..3078d79 100644 --- a/src/dynamic_foraging_processing/qc/processed/stage.py +++ b/src/dynamic_foraging_processing/qc/processed/stage.py @@ -14,6 +14,7 @@ from dynamic_foraging_processing.qc._core.base import BaseQC from dynamic_foraging_processing.qc._core.result import to_metrics from dynamic_foraging_processing.qc.processed.results import behavior_qc_results +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes class ProcessedQC(BaseQC): @@ -26,8 +27,8 @@ def run( right_lick_times: np.ndarray, results_folder: t.Optional[str] = None, *, - manual_left_times: t.Optional[np.ndarray] = None, - manual_right_times: t.Optional[np.ndarray] = None, + manual_left: ManualWaterTimes = ManualWaterTimes(), + manual_right: ManualWaterTimes = ManualWaterTimes(), ) -> t.List[QCMetric]: """Compute the behavior QC checks and return them as metrics. @@ -45,8 +46,9 @@ def run( Timestamps (s) of left/right-port licks. results_folder : str, optional Directory to write the plots into. If ``None``, plots are skipped. - manual_left_times, manual_right_times : numpy.ndarray, optional - Manual-water delivery timestamps passed through to the side-bias figure. + manual_left, manual_right : ManualWaterTimes, optional + Left/right experimenter-water delivery timestamps, split into + unaligned and go-cue-aligned, passed through to the side-bias figure. Returns ------- @@ -59,7 +61,7 @@ def run( left_lick_times, right_lick_times, results_folder, - manual_left_times=manual_left_times, - manual_right_times=manual_right_times, + manual_left=manual_left, + manual_right=manual_right, ) return to_metrics(results) diff --git a/src/dynamic_foraging_processing/utils/rewards.py b/src/dynamic_foraging_processing/utils/rewards.py index ac487e1..73287ec 100644 --- a/src/dynamic_foraging_processing/utils/rewards.py +++ b/src/dynamic_foraging_processing/utils/rewards.py @@ -8,6 +8,41 @@ from dynamic_foraging_processing.utils.timestamps import find_closest_timestamps +#: Reward-delivery annotation for water the animal worked for. +EARNED = "earned" +#: Reward-delivery annotation for task-triggered free water (autowater, anti-bias). +AUTO = "auto" +#: Reward-delivery annotation for experimenter water not aligned to a go cue. +MANUAL = "manual" +#: Reward-delivery annotation for experimenter water delivered at the go cue. +MANUAL_GO_CUE_ALIGNED = "manual_go_cue_aligned" + +#: Shared empty default for the ``ManualWaterTimes`` fields. Never mutated. +_NO_TIMES = np.array([]) + + +class ManualWaterTimes(t.NamedTuple): + """One lick port's experimenter-triggered water times, split by alignment. + + The acquisition software emits the two kinds on separate software-event + streams, so they are kept apart here rather than conflated into a single + "manual" bucket: only ``unaligned`` water is independent of the trial + structure, while ``go_cue_aligned`` water fires at the go cue like autowater + but was triggered by the experimenter, not the task. + + Attributes + ---------- + unaligned : numpy.ndarray + Timestamps of the ``LeftManualWater`` / ``RightManualWater`` events: + water given at an arbitrary moment, not tied to a go cue. + go_cue_aligned : numpy.ndarray + Timestamps of the ``LeftManualAutoReward`` / ``RightManualAutoReward`` + events: water the experimenter triggered to land on the go cue. + """ + + unaligned: np.ndarray = _NO_TIMES + go_cue_aligned: np.ndarray = _NO_TIMES + def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]: """Parse a ``TrialOutcome`` software-event payload into its domain model. @@ -52,26 +87,31 @@ def _free_water_label(trial: t.Optional[Trial]) -> 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" + return EARNED + return AUTO def get_reward_deliveries( reward_delivery_times: np.ndarray, trial_outcome_df: pd.DataFrame, - manual_water_times: np.ndarray, + manual_water: ManualWaterTimes, response_times: np.ndarray, ) -> np.ndarray: """Classify one lick port's reward deliveries by how the water was given. Annotates the deliveries of a single lick port. Each delivery is classified - as follows, with ``manual`` taking precedence because manual water is not - aligned to a go cue: + as follows, with the experimenter-triggered labels taking precedence over the + trial-derived ones because the experimenter acts outside the task logic: - ``manual`` -- the delivery is the closest hardware (harp) timestamp to a - ``GiveManualWater`` software event for this port. The software-event - timestamps are correlated to the reward-delivery timestamps with - :func:`find_closest_timestamps`. + ``LeftManualWater``/``RightManualWater`` software event for this port: + water given at an arbitrary moment, not aligned to a go cue. Highest + precedence. + - ``manual_go_cue_aligned`` -- the delivery is the closest timestamp to a + ``LeftManualAutoReward``/``RightManualAutoReward`` event: water the + experimenter triggered to land on the go cue. It fires at the go cue like + autowater but the task did not schedule it, so it is neither ``auto`` nor + ``earned``. - ``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 @@ -80,6 +120,15 @@ def get_reward_deliveries( ``anti_bias_left_water``/``anti_bias_right_water``. - ``earned`` -- otherwise: water the animal worked for. + The two manual labels are kept apart rather than collapsed into ``manual`` + because they differ in kind: only unaligned manual water is independent of + the trial structure. Collapsing them would also make the go-cue-aligned + deliveries indistinguishable from task-scheduled ``auto`` water in the QC + figure, which is what the separate rows there exist to show. + + The software-event timestamps are correlated to the reward-delivery + timestamps with :func:`find_closest_timestamps`. + Every valve opening is annotated and none is filtered out, so the series is a complete record of the water this port delivered. In particular a trial reporting ``is_rewarded=False`` keeps its delivery: free water is triggered @@ -108,9 +157,10 @@ def get_reward_deliveries( trial_outcome_df : pandas.DataFrame 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``). + manual_water : ManualWaterTimes + 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``. @@ -119,7 +169,7 @@ def get_reward_deliveries( ------- numpy.ndarray Array of the same shape as ``reward_delivery_times`` whose entries are - ``"earned"``, ``"auto"``, or ``"manual"``. + ``"earned"``, ``"auto"``, ``"manual_go_cue_aligned"``, or ``"manual"``. Raises ------ @@ -142,22 +192,28 @@ def get_reward_deliveries( # get one trial position per reward delivery. trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times) - annotated_rewards = [] + trial_labels = [] for trial_index in trial_indices_in_reward_times: outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"]) - annotated_rewards.append(_free_water_label(outcome.trial if outcome is not None else None)) + trial_labels.append(_free_water_label(outcome.trial if outcome is not None else None)) # Object dtype, not the inferred fixed-width string dtype: a run of only "auto" - # and "earned" entries would be too narrow to hold "manual" and would truncate it. - 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 - # 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) - 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" + # and "earned" entries would be too narrow to hold the manual labels and would + # truncate them. + annotated_rewards = np.array(trial_labels, dtype=object) + + # Experimenter water is independent of trials (multiple deliveries can occur + # within one trial), so annotate those deliveries directly rather than through + # the trial they fall in. Correlate each software event to its closest reward + # delivery; the returned positions index into reward_times, i.e. the deliveries + # the experimenter caused. Unaligned manual water is written last so it wins + # where both kinds land on the same delivery. + for times, label in ( + (manual_water.go_cue_aligned, MANUAL_GO_CUE_ALIGNED), + (manual_water.unaligned, MANUAL), + ): + times = np.asarray(times) + if times.size: + annotated_rewards[find_closest_timestamps(times, reward_times)] = label return annotated_rewards diff --git a/tests/test_nwb/test_acquisition/test_acquisition_builder.py b/tests/test_nwb/test_acquisition/test_acquisition_builder.py index be0f544..10ef615 100644 --- a/tests/test_nwb/test_acquisition/test_acquisition_builder.py +++ b/tests/test_nwb/test_acquisition/test_acquisition_builder.py @@ -95,9 +95,12 @@ def _make_response_frame() -> pd.DataFrame: ) -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")) +def _manual_water_frame(*times: float) -> pd.DataFrame: + """Build a manual-water software-event stream firing at ``times``. + + Only the index is read, so the payload is a placeholder. + """ + return pd.DataFrame({"data": [True] * len(times)}, index=pd.Index(list(times), name="time")) def _make_digital_input_frame() -> pd.DataFrame: @@ -116,10 +119,14 @@ def _make_digital_input_frame() -> pd.DataFrame: ) -def _make_dataset(manual_water=None): - """Build a path-aware fake dataset rooted at ``Behavior``.""" - if manual_water is None: - manual_water = _empty_manual_water_frame() +def _make_dataset(**manual_water_streams): + """Build a path-aware fake dataset rooted at ``Behavior``. + + Manual-water streams are optional per session, so only the ones named in + ``manual_water_streams`` (e.g. ``RightManualWater=_manual_water_frame(0.49)``) + are present; the rest are absent, as they are for a session where the + experimenter gave no water of that kind. + """ return _FakeNode( { "Behavior": _FakeNode( @@ -134,7 +141,10 @@ def _make_dataset(manual_water=None): { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), "Response": _FakeStream(_make_response_frame()), - "GiveManualWaterRight": _FakeStream(manual_water), + **{ + name: _FakeStream(frame) + for name, frame in manual_water_streams.items() + }, } ), } @@ -171,39 +181,37 @@ def test_get_valve_writes_filters_to_write_messages(): assert list(result.index) == [0.1, 0.3, 0.5] -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")) - builder = AcquisitionBuilder(loader=_make_loader(_make_dataset(manual))) +def test_get_manual_water_times_selects_streams_by_side(): + """The side comes from the stream name, not from an event payload.""" + dataset = _make_dataset( + LeftManualWater=_manual_water_frame(0.11), + LeftManualAutoReward=_manual_water_frame(0.12), + RightManualWater=_manual_water_frame(0.49), + RightManualAutoReward=_manual_water_frame(0.31, 0.32), + ) + builder = AcquisitionBuilder(loader=_make_loader(dataset)) - result = builder.get_manual_water_times() + left = builder.get_manual_water_times(is_right=False) + right = builder.get_manual_water_times(is_right=True) - pd.testing.assert_frame_equal(result, manual) + np.testing.assert_array_equal(left.unaligned, np.array([0.11])) + np.testing.assert_array_equal(left.go_cue_aligned, np.array([0.12])) + np.testing.assert_array_equal(right.unaligned, np.array([0.49])) + np.testing.assert_array_equal(right.go_cue_aligned, np.array([0.31, 0.32])) def test_get_manual_water_times_returns_empty_when_absent(): - """A missing manual-water stream yields an empty frame with a ``data`` column.""" - dataset = _FakeNode( - { - "Behavior": _FakeNode( - { - "HarpBehavior": _FakeNode({"OutputSet": _FakeStream(_make_output_set_frame())}), - "SoftwareEvents": _FakeNode( - { - "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), - "Response": _FakeStream(_make_response_frame()), - } - ), - } - ) - } - ) - builder = AcquisitionBuilder(loader=_make_loader(dataset)) + """Missing manual-water streams yield empty arrays, not an error. - result = builder.get_manual_water_times() + Each stream only exists when the experimenter gave water of that kind, so a + session with no manual water at all has none of the four files. + """ + builder = AcquisitionBuilder(loader=_make_loader(_make_dataset())) - assert list(result.columns) == ["data"] - assert result.empty + for is_right in (False, True): + times = builder.get_manual_water_times(is_right=is_right) + assert times.unaligned.size == 0 + assert times.go_cue_aligned.size == 0 def test_get_lick_times_selects_di_port_by_side(): @@ -240,7 +248,6 @@ def test_get_lick_times_returns_empty_when_absent(): { "TrialOutcome": _FakeStream(_make_trial_outcome_frame()), "Response": _FakeStream(_make_response_frame()), - "GiveManualWaterRight": _FakeStream(_empty_manual_water_frame()), } ), } @@ -257,9 +264,10 @@ def test_get_lick_times_returns_empty_when_absent(): def test_build_acquisition_returns_populated_list(): """``build_acquisition`` returns the table plus reward and lick port series.""" - # A right-side manual-water event (data=True) near the second right delivery. - manual = pd.DataFrame({"data": [True]}, index=pd.Index([0.49], name="time")) - builder = AcquisitionBuilder(loader=_make_loader(_make_dataset(manual))) + # A right-side manual-water event near the second right delivery. + builder = AcquisitionBuilder( + loader=_make_loader(_make_dataset(RightManualWater=_manual_water_frame(0.49))) + ) acquisition = builder.build_acquisition() @@ -299,6 +307,27 @@ def test_build_acquisition_returns_populated_list(): assert "DIPort1" in right_lick.description +def test_build_acquisition_annotates_manual_auto_reward_separately(): + """Manual go-cue-aligned water is its own label, on its own side. + + The left stream must not reach the right series: the side is carried by the + stream name now, so a left event near a right delivery is not that delivery's. + """ + builder = AcquisitionBuilder( + loader=_make_loader( + _make_dataset( + RightManualAutoReward=_manual_water_frame(0.49), + LeftManualWater=_manual_water_frame(0.1), + ) + ) + ) + + _, left_reward, right_reward, _, _ = builder.build_acquisition() + + np.testing.assert_array_equal(left_reward.data, np.array(["manual"])) + np.testing.assert_array_equal(right_reward.data, np.array(["auto", "manual_go_cue_aligned"])) + + def test_build_acquisition_defaults_none_description_to_empty_string(): """A None stream description falls back to "" so the table validates.""" loader = _make_loader() diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py index 9e777d4..c510d65 100644 --- a/tests/test_pipeline/test_pipeline.py +++ b/tests/test_pipeline/test_pipeline.py @@ -16,6 +16,7 @@ from dynamic_foraging_processing.pipeline import _pipeline from dynamic_foraging_processing.pipeline._pipeline import Pipeline from dynamic_foraging_processing.processing.models import TrialConfig +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes def _make_loader() -> MagicMock: @@ -196,16 +197,14 @@ def run(self, dataset, results_folder): class _FakeProcessedQC: """Processed QC stub returning one sentinel metric.""" - def run( - self, trials, left, right, results_folder, *, manual_left_times, manual_right_times - ): + def run(self, trials, left, right, results_folder, *, manual_left, manual_right): """Record args and return processed metrics.""" captured["processed"] = { "left": left, "right": right, "results_folder": results_folder, - "manual_left": manual_left_times, - "manual_right": manual_right_times, + "manual_left": manual_left, + "manual_right": manual_right, } return ["proc1"] @@ -220,8 +219,8 @@ def run( trials, np.array([1.0]), np.array([2.0]), - np.array([0.1]), - np.array([0.2]), + ManualWaterTimes(unaligned=np.array([0.1])), + ManualWaterTimes(go_cue_aligned=np.array([0.2])), "out", ) @@ -229,8 +228,10 @@ def run( assert captured["raw"] == (pipeline.loader.dataset, "out") np.testing.assert_array_equal(captured["processed"]["left"], np.array([1.0])) np.testing.assert_array_equal(captured["processed"]["right"], np.array([2.0])) - np.testing.assert_array_equal(captured["processed"]["manual_left"], np.array([0.1])) - np.testing.assert_array_equal(captured["processed"]["manual_right"], np.array([0.2])) + np.testing.assert_array_equal(captured["processed"]["manual_left"].unaligned, np.array([0.1])) + np.testing.assert_array_equal( + captured["processed"]["manual_right"].go_cue_aligned, np.array([0.2]) + ) def test_add_acquisition_series_builds_time_series(): @@ -342,20 +343,23 @@ def test_add_trials_skips_when_time_columns_missing(): nwb_file.add_trial.assert_not_called() -def test_manual_water_times_reads_manual_annotations(): - """Manual-water times are the reward events annotated ``"manual"`` per side.""" +def test_manual_water_times_splits_annotations_by_alignment(): + """Each side's experimenter water is read back split by go-cue alignment.""" nwb_file = MagicMock() nwb_file.acquisition = { "left_reward_delivery_time": _FakeSeries( - np.array(["manual", "earned", "manual"]), np.array([0.1, 0.2, 0.3]) + np.array(["manual", "earned", "manual", "manual_go_cue_aligned"]), + np.array([0.1, 0.2, 0.3, 0.4]), ), "right_reward_delivery_time": _FakeSeries(np.array(["auto"]), np.array([0.5])), } left, right = Pipeline._manual_water_times(nwb_file) - np.testing.assert_array_equal(left, np.array([0.1, 0.3])) - assert right.size == 0 + np.testing.assert_array_equal(left.unaligned, np.array([0.1, 0.3])) + np.testing.assert_array_equal(left.go_cue_aligned, np.array([0.4])) + assert right.unaligned.size == 0 + assert right.go_cue_aligned.size == 0 def test_read_processed_inputs_reads_from_nwb(): @@ -369,7 +373,9 @@ def test_read_processed_inputs_reads_from_nwb(): "left_reward_delivery_time": _FakeSeries( np.array(["manual", "earned"]), np.array([0.1, 0.2]) ), - "right_reward_delivery_time": _FakeSeries(np.array(["manual"]), np.array([0.3])), + "right_reward_delivery_time": _FakeSeries( + np.array(["manual_go_cue_aligned"]), np.array([0.3]) + ), } pipeline = _make_pipeline() @@ -379,8 +385,8 @@ def test_read_processed_inputs_reads_from_nwb(): assert got_trials is trials np.testing.assert_array_equal(left, np.array([1.0])) np.testing.assert_array_equal(right, np.array([2.0, 2.5])) - np.testing.assert_array_equal(manual_left, np.array([0.1])) - np.testing.assert_array_equal(manual_right, np.array([0.3])) + np.testing.assert_array_equal(manual_left.unaligned, np.array([0.1])) + np.testing.assert_array_equal(manual_right.go_cue_aligned, np.array([0.3])) def test_run_nwb_writes_nwb_and_processing(tmp_path, monkeypatch): diff --git a/tests/test_qc/test_plots.py b/tests/test_qc/test_plots.py index 2ddbb02..dbb782f 100644 --- a/tests/test_qc/test_plots.py +++ b/tests/test_qc/test_plots.py @@ -6,6 +6,7 @@ import numpy as np from dynamic_foraging_processing.qc.processed import plots as _plots +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes def test_plot_lick_intervals_writes_file(tmp_path): @@ -66,8 +67,11 @@ def test_plot_side_bias_full_inputs(tmp_path): go_cue_times=np.array([0.5, 1.5, 2.5, 3.5, 4.5, 5.5]), autowater_left=np.array([1, 0, 0, 0, 0, 0]), autowater_right=np.array([0, 0, 0, 1, 0, 0]), - manual_left_times=np.array([0.1, 3.6]), # 0.1 -> -1, 3.6 -> trial index - manual_right_times=np.array([5.6]), + manual_left=ManualWaterTimes( + unaligned=np.array([0.1, 3.6]), # 0.1 -> -1, 3.6 -> trial index + go_cue_aligned=np.array([2.6]), + ), + manual_right=ManualWaterTimes(unaligned=np.array([5.6]), go_cue_aligned=np.array([1.6])), anti_bias_left_water=np.array([False, False, True, False, False, False]), anti_bias_right_water=np.array([False, False, False, False, False, True]), anti_bias_lickspout_movement=np.array([0.0, 0.5, 0.0, 0.0, -0.3, 0.0]), @@ -142,8 +146,8 @@ def test_add_behavior_plot_gives_manual_water_its_own_rows(): rewarded_right=None, autowater_left=np.array([1, 0, 0, 0]), autowater_right=np.array([0, 0, 0, 1]), - manual_left_times=np.array([1.6]), - manual_right_times=np.array([2.6]), + manual_left=ManualWaterTimes(unaligned=np.array([1.6])), + manual_right=ManualWaterTimes(unaligned=np.array([2.6])), go_cue_times=np.array([0.5, 1.5, 2.5, 3.5]), ) labels = [text.get_text() for text in ax.get_legend().get_texts()] @@ -153,7 +157,55 @@ def test_add_behavior_plot_gives_manual_water_its_own_rows(): tick_labels = [text.get_text() for text in ax.get_yticklabels()] assert "L Manual Water" in tick_labels and "R Manual Water" in tick_labels # The manual rows are outside the autowater bands, so the two never overlap. - assert ax.get_ylim() == (-0.6, 1.6) + assert ax.get_ylim() == (-0.8, 1.8) + plt.close(fig) + + +def test_add_behavior_plot_separates_go_cue_aligned_manual_water(): + """Go-cue-aligned manual water gets its own rows and its own legend entry. + + It fires at the go cue exactly as autowater does, so sharing either the + autowater row or the manual row would hide which gave the water. + """ + fig, ax = plt.subplots() + _plots._add_behavior_plot( + ax, + np.array([0, 1, 0, 1]), + rewarded_left=None, + rewarded_right=None, + autowater_left=None, + autowater_right=None, + manual_left=ManualWaterTimes(unaligned=np.array([1.6]), go_cue_aligned=np.array([2.6])), + manual_right=ManualWaterTimes(go_cue_aligned=np.array([3.6])), + go_cue_times=np.array([0.5, 1.5, 2.5, 3.5]), + ) + labels = [text.get_text() for text in ax.get_legend().get_texts()] + # One entry per kind, however many sides carry it. + assert labels.count("Manual Water") == 1 + assert labels.count("Manual Water (go cue aligned)") == 1 + tick_labels = [text.get_text() for text in ax.get_yticklabels()] + assert "L Manual Water (go cue)" in tick_labels + assert "R Manual Water (go cue)" in tick_labels + plt.close(fig) + + +def test_add_behavior_plot_labels_only_the_kinds_present(): + """A session with only go-cue-aligned water gets only that legend entry.""" + fig, ax = plt.subplots() + _plots._add_behavior_plot( + ax, + np.array([0, 1]), + rewarded_left=None, + rewarded_right=None, + autowater_left=None, + autowater_right=None, + manual_left=ManualWaterTimes(go_cue_aligned=np.array([1.6])), + manual_right=ManualWaterTimes(), + go_cue_times=np.array([0.5, 1.5]), + ) + labels = [text.get_text() for text in ax.get_legend().get_texts()] + assert "Manual Water (go cue aligned)" in labels + assert "Manual Water" not in labels plt.close(fig) @@ -167,12 +219,13 @@ def test_add_behavior_plot_omits_manual_water_legend_when_absent(): rewarded_right=None, autowater_left=None, autowater_right=None, - manual_left_times=np.array([]), - manual_right_times=np.array([]), + manual_left=ManualWaterTimes(), + manual_right=ManualWaterTimes(), go_cue_times=np.array([0.5, 1.5]), ) labels = [text.get_text() for text in ax.get_legend().get_texts()] assert "Manual Water" not in labels + assert "Manual Water (go cue aligned)" not in labels plt.close(fig) diff --git a/tests/test_utils/test_rewards.py b/tests/test_utils/test_rewards.py index 695c5d3..6ce87bc 100644 --- a/tests/test_utils/test_rewards.py +++ b/tests/test_utils/test_rewards.py @@ -7,7 +7,7 @@ import pytest from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome -from dynamic_foraging_processing.utils.rewards import get_reward_deliveries +from dynamic_foraging_processing.utils.rewards import ManualWaterTimes, get_reward_deliveries def _outcome_payload(auto=None, is_rewarded: bool = True, mechanism: str = "autowater") -> dict: @@ -60,7 +60,7 @@ def test_get_reward_deliveries_marks_default_trials_as_earned(): trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4, 1.9])) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["earned", "earned", "earned"])) @@ -73,7 +73,7 @@ def test_get_reward_deliveries_marks_auto_response_trials_as_auto(): trial_outcome_df = _trial_outcome_df(np.array([1.1, 1.4]), autos=[True, False]) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) @@ -97,7 +97,7 @@ def test_get_reward_deliveries_marks_all_free_water_as_auto(mechanism): ) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["auto"])) @@ -116,7 +116,7 @@ def test_get_reward_deliveries_matches_closest_response_time(): trial_outcome_df = _trial_outcome_df(np.array([0.9, 5.0]), autos=[None, True]) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) @@ -138,7 +138,7 @@ def test_get_reward_deliveries_keeps_deliveries_on_unrewarded_trials(): ) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["earned", "auto", "auto"])) @@ -157,51 +157,126 @@ def test_get_reward_deliveries_labels_both_sides_of_a_split_trial(): # This port saw one opening on that trial; the trial gave free water. annotations = get_reward_deliveries( - np.array([0.15]), trial_outcome_df, np.array([]), response_times + np.array([0.15]), trial_outcome_df, ManualWaterTimes(), response_times ) 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, np.array([]), response_times) + annotations = get_reward_deliveries( + np.array([0.15]), earned_only, ManualWaterTimes(), response_times + ) np.testing.assert_array_equal(annotations, np.array(["earned"])) def test_get_reward_deliveries_marks_manual_water_as_manual(): - """Deliveries closest to a manual-water event are annotated as ``manual``.""" + """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_times = np.array([0.43]) + manual_water = ManualWaterTimes(unaligned=np.array([0.43])) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water_times, response_times + reward_times, trial_outcome_df, manual_water, response_times ) np.testing.assert_array_equal(annotations, np.array(["earned", "manual", "earned"])) +def test_get_reward_deliveries_marks_manual_auto_reward_as_go_cue_aligned(): + """Deliveries closest to a manual auto-reward event get their own label. + + The experimenter triggered this water at the go cue, so the task logic never + 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])) + + annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water, response_times + ) + + np.testing.assert_array_equal( + annotations, np.array(["earned", "manual_go_cue_aligned", "earned"]) + ) + + 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_times = np.array([0.42]) + manual_water = ManualWaterTimes(unaligned=np.array([0.42])) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, manual_water_times, response_times + reward_times, trial_outcome_df, manual_water, response_times ) np.testing.assert_array_equal(annotations, np.array(["earned", "manual"])) +def test_get_reward_deliveries_go_cue_aligned_takes_precedence_over_auto(): + """Manual go-cue-aligned water outranks the trial's own free-water label. + + 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])) + + annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water, response_times + ) + + np.testing.assert_array_equal(annotations, np.array(["earned", "manual_go_cue_aligned"])) + + +def test_get_reward_deliveries_unaligned_manual_outranks_go_cue_aligned(): + """When both kinds match one delivery, the unaligned label wins. + + 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])) + + annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water, response_times + ) + + 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])) + + annotations = get_reward_deliveries( + reward_times, trial_outcome_df, manual_water, response_times + ) + + np.testing.assert_array_equal( + annotations, np.array(["manual", "earned", "manual_go_cue_aligned"]) + ) + + 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, np.array([]), np.array([0.0]) + np.array([]), trial_outcome_df, ManualWaterTimes(), np.array([0.0]) ) assert isinstance(annotations, np.ndarray) @@ -219,7 +294,7 @@ def test_get_reward_deliveries_accepts_json_and_model_payloads(): ) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), response_times + reward_times, trial_outcome_df, ManualWaterTimes(), response_times ) np.testing.assert_array_equal(annotations, np.array(["auto", "auto"])) @@ -230,7 +305,9 @@ def test_get_reward_deliveries_rejects_misaligned_response_times(): 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])) + get_reward_deliveries( + np.array([0.15]), trial_outcome_df, ManualWaterTimes(), np.array([0.1]) + ) def test_get_reward_deliveries_returns_one_annotation_per_delivery(): @@ -239,7 +316,7 @@ def test_get_reward_deliveries_returns_one_annotation_per_delivery(): reward_times = np.array([0.1, 0.2, 0.6]) annotations = get_reward_deliveries( - reward_times, trial_outcome_df, np.array([]), np.array([0.0, 0.5]) + reward_times, trial_outcome_df, ManualWaterTimes(), np.array([0.0, 0.5]) ) assert isinstance(annotations, np.ndarray) diff --git a/uv.lock b/uv.lock index 50f7fa1..5579f6b 100644 --- a/uv.lock +++ b/uv.lock @@ -39,7 +39,7 @@ wheels = [ [[package]] name = "aind-behavior-dynamic-foraging" version = "0.0.2rc37" -source = { git = "https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git?rev=933a7aee2627de1979a982847f3553cf7435cbd2#933a7aee2627de1979a982847f3553cf7435cbd2" } +source = { git = "https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git?rev=4e4faa4941c8d9cc8024e760fdaec30339b626dc#4e4faa4941c8d9cc8024e760fdaec30339b626dc" } dependencies = [ { name = "aind-behavior-services" }, { name = "contraqctor" }, @@ -364,7 +364,7 @@ name = "cffi" version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9e/ef/008a1939e372c06329a3fce4279c02f328488f3526744906eeec3da7ad5f/cffi-2.1.1.tar.gz", hash = "sha256:dd31f52ea1086513bb9df30f8fcee9b8918323ae067a3d5b78bc826a000712be", size = 530807, upload-time = "2026-08-03T21:21:18.939Z" } wheels = [ @@ -884,7 +884,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "aind-behavior-dynamic-foraging", extras = ["data"], git = "https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git?rev=933a7aee2627de1979a982847f3553cf7435cbd2" }, + { name = "aind-behavior-dynamic-foraging", extras = ["data"], git = "https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git?rev=4e4faa4941c8d9cc8024e760fdaec30339b626dc" }, { name = "aind-data-schema", marker = "extra == 'qc'", specifier = ">=2.4.1" }, { name = "aind-nwb-utils", marker = "extra == 'nwb'" }, { name = "dynamic-foraging-processing", extras = ["qc", "nwb"], marker = "extra == 'full'" }, @@ -1933,7 +1933,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [