From 10230fa99a55e6d0b10fcfe3715481ae9d7e3990 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 <109561860+arjunsridhar12345@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:16:21 -0700 Subject: [PATCH 1/2] feat: update the columns in the trials table from the acquisition software (#63) * feat: add additional columns from acquisition * test: update tests * docs: update trial table docs with changelog and updates --- docs/trials_table_mapping.md | 60 +++++-- .../pipeline/_pipeline.py | 62 ++++++-- .../processing/_trial_table.py | 148 ++++++++++++++---- .../processing/models/trial_config.py | 53 ++++++- tests/test_pipeline/test_pipeline.py | 44 +++++- tests/test_processing/test_trial_table.py | 19 ++- 6 files changed, 323 insertions(+), 63 deletions(-) diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index 7e1d457..12223b3 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -77,18 +77,58 @@ Columns are grouped by the raw source they map from. | `base_reward_probability_sum` | If `type == "CoupledTrialGenerator"`, look at `reward_probability_parameters`. | | `min_reward_each_block` | Present when `type == "CoupledWarmupTrialGenerator"` (has `min_block_reward`); otherwise `None`. | -### From `QuiescentPeriod.json` (`SoftwareEvents` stream) +### Trial period timing (the four period `SoftwareEvents` streams) + +Each of `QuiescentPeriod.json`, `ResponsePeriod.json`, +`RewardConsumptionPeriod.json`, and `ItiPeriod.json` emits one event per trial at +the **start** of its period, and the periods run back-to-back (the AIND DF v2 +trial structure): + +``` + go cue response ITI start next trial + | registered | | + quiescent | response | reward consumption | ITI | quiescent +|--------------->|------------------->|------------------->|-------------->|----------> +q[i] r[i] c[i] iti[i] q[i+1] +``` + +So each period's stop time is the next period's start time. All four streams are +aligned with `TrialOutcome` by position (a length mismatch is reported by +`_check_aligned`; a short stream pads with `NaN`). + +Verified on +`864253_2026-07-29_11-50-18` (753 trials): all five streams have equal length, +`q[i] < r[i] < c[i] < iti[i] < q[i+1]` holds for every trial, the `SoundCard` go +cue falls within 1.2 ms of `r[i]` on every trial, and the realized period +durations track the configured ones (reward consumption ≈ +`reward_consumption_duration`, `iti[i] → q[i+1]` ≈ `ITI_duration`). | Trials column | Mapping | | --- | --- | -| `delay_start_time` | `timestamp`. | -| `start_time` | `timestamp` column. | - -### From `ITI_period.json` (`SoftwareEvents` stream) - -| Trials column | Mapping | -| --- | --- | -| `stop_time` | `timestamp` column. Possible QC check: length should match `QuiescentPeriod.json`. | +| `quiescent_start_time` | `QuiescentPeriod` `timestamp`. | +| `quiescent_stop_time` | `ResponsePeriod` `timestamp` (the quiescent period ends where the response period begins). | +| `response_start_time` | `ResponsePeriod` `timestamp`. | +| `response_stop_time` | `RewardConsumptionPeriod` `timestamp`. | +| `reward_consumption_start_time` | `RewardConsumptionPeriod` `timestamp`. | +| `reward_consumption_stop_time` | `ItiPeriod` `timestamp`. | +| `ITI_start_time` | `ItiPeriod` `timestamp`. | +| `ITI_stop_time` | The **next** trial's `QuiescentPeriod` `timestamp`; `NaN` on the last trial of the session. | +| `delay_start_time` | `QuiescentPeriod` `timestamp` — the legacy name for `quiescent_start_time` (see the note below). | + +There are no `start_time` / `stop_time` trial columns. NWB's `TimeIntervals` +requires a native `start_time` / `stop_time` per trial, so the pipeline derives +the trial extent when writing: `start_time` is `quiescent_start_time` and +`stop_time` is `ITI_stop_time`, falling back to `ITI_start_time` on the last +trial. + +> **`delay` means `quiescent`.** The legacy `delay_*` columns describe the +> acquisition software's *quiescence period* — the lick-free interval preceding +> the go cue. `delay_start_time` is therefore the `QuiescentPeriod` timestamp and +> always equals `quiescent_start_time`, and `delay_duration` / +> `delay_beta` / `delay_min` / `delay_max` summarize +> `quiescence_period_duration`. Note `delay_duration` is the *configured* +> duration: each lick restarts the quiescent period, so the realized duration +> (`quiescent_stop_time - quiescent_start_time`) can be longer. ### From `HarpBehavior` (`PulseSupplyPort{0,1}`) @@ -144,3 +184,5 @@ These were mapped during exploration but are no longer in scope: | 2026-06-20 | `reward_probabilityL` / `reward_probabilityR` now read the block probability from `trial.metadata.p_reward_left` / `p_reward_right` instead of the top-level per-trial `trial.p_reward_left` / `p_reward_right`. | | 2026-07-22 | `lickspout_position_x` / `y1` / `y2` / `z` now derive from the `HarpManipulator` `AccumulatedSteps` stream (microsteps → mm via the `InputSchemas.Rig` manipulator calibration, `full_step_to_mm / microstep_resolution`), sampled per trial via the closest sample in the `[start_time, stop_time)` window and re-referenced to the session-start position (displacement relative to session start, mm), replacing the static `InitialManipulatorPosition` software event. `Motor{i}` maps to `Axis(i + 1)` (X, Y1, Y2, Z). The rig and `AccumulatedSteps` streams are required when there are trials (`build` raises if either is missing). Column descriptions corrected from `um` to `mm`. | | 2026-07-24 | `reward_size_left` / `reward_size_right` moved from session-level `task_parameters.reward_size` to per-trial `Trial.reward_size` (fields `.left` / `.right`). The columns are now nullable — `None` when the trial is missing. A missing `TaskLogic` stream no longer raises; session distribution columns are simply null. `min_reward_each_block` moved from `CoupledTrialGenerator` to `CoupledWarmupTrialGenerator`. | +| 2026-08-06 | **Breaking:** the trial `start_time` / `stop_time` columns are removed and replaced by one start/stop pair per task period: `quiescent_start_time` / `quiescent_stop_time`, `response_start_time` / `response_stop_time`, `reward_consumption_start_time` / `reward_consumption_stop_time`, and `ITI_start_time` / `ITI_stop_time`, read from the `ResponsePeriod` and `RewardConsumptionPeriod` streams in addition to `QuiescentPeriod` and `ItiPeriod`. Each period event marks its period's start, so each stop is the next period's start; `ITI_stop_time` is the next trial's `QuiescentPeriod` timestamp (`NaN` on the last trial). The two new streams are also checked for positional alignment with `TrialOutcome`. NWB's required native `start_time` / `stop_time` are now derived when writing (`quiescent_start_time` → `ITI_stop_time`, falling back to `ITI_start_time`), so the NWB trials table changes in two ways: the old `start_time` / `stop_time` columns are gone, and the native trial extent now ends at the *end* of the ITI rather than at its start. | +| 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. | diff --git a/src/dynamic_foraging_processing/pipeline/_pipeline.py b/src/dynamic_foraging_processing/pipeline/_pipeline.py index 022694a..2a08af9 100644 --- a/src/dynamic_foraging_processing/pipeline/_pipeline.py +++ b/src/dynamic_foraging_processing/pipeline/_pipeline.py @@ -53,9 +53,16 @@ _DEFAULT_LEFT_LICK = LickSource("HarpBehavior", "DigitalInputState", "DIPort0") _DEFAULT_RIGHT_LICK = LickSource("HarpBehavior", "DigitalInputState", "DIPort1") -#: Trials-table columns NWB models natively; every other column is added as an -#: extra trial column. -_TRIAL_TIME_COLUMNS = ("start_time", "stop_time") +#: Trials-table column NWB's required native ``start_time`` is taken from. The +#: trials table itself no longer carries trial ``start_time`` / ``stop_time`` +#: columns — it carries one start/stop pair per task period instead — but +#: ``TimeIntervals`` requires both, so they are derived here. +_NWB_START_COLUMN = "quiescent_start_time" + +#: Columns NWB's required native ``stop_time`` is taken from, in order of +#: preference: the end of the ITI, falling back to its start on the last trial +#: of the session (where the ITI end is unknown). +_NWB_STOP_COLUMNS = ("ITI_stop_time", "ITI_start_time") #: Source repository recorded in the ``processing.json`` data process. _CODE_URL = "https://github.com/AllenNeuralDynamics/dynamic-foraging-processing" @@ -224,23 +231,54 @@ def _add_acquisition_table(nwb_file: pynwb.NWBFile, table: AcquisitionTable) -> ) @staticmethod - def _add_trials(nwb_file: pynwb.NWBFile, trials: pd.DataFrame) -> None: + def _trial_extent(row: pd.Series) -> t.Tuple[float, float]: + """Return NWB's required native ``(start_time, stop_time)`` for one trial. + + The trials table has no trial start/stop columns of its own, so the + trial's extent is taken from its period bounds: it starts with the + quiescent period and ends with the ITI, falling back to the ITI start on + the last trial of the session (whose ITI end is unknown). + + Parameters + ---------- + row : pandas.Series + One row of the trials table. + + Returns + ------- + tuple of (float, float) + The trial start and stop time (seconds). + """ + stops = [row[column] for column in _NWB_STOP_COLUMNS if pd.notnull(row[column])] + stop = stops[0] if stops else np.nan + return float(row[_NWB_START_COLUMN]), float(stop) + + @classmethod + def _add_trials(cls, nwb_file: pynwb.NWBFile, trials: pd.DataFrame) -> None: """Add the trials table to the NWB file's native ``trials`` table. - ``start_time`` / ``stop_time`` are modeled natively by NWB; every other - column is registered as an extra trial column (described by the matching - :class:`TrialConfig` field) and populated per row. The DataFrame index + Every trials-table column is registered as an extra trial column + (described by the matching :class:`TrialConfig` field) and populated per + row. NWB additionally requires a native ``start_time`` / ``stop_time`` per + trial, which are derived from the period columns (see + :meth:`_trial_extent`) rather than stored as columns. The DataFrame index (named ``id``) is replicated as each trial's NWB ``id``. An empty table - (or one missing the required time columns) is skipped. + (or one missing the period columns the extent is derived from) is skipped. """ - if trials.empty or any(col not in trials.columns for col in _TRIAL_TIME_COLUMNS): + required = (_NWB_START_COLUMN, *_NWB_STOP_COLUMNS) + if trials.empty or any(col not in trials.columns for col in required): return descriptions = TrialConfig.column_descriptions() - extra_columns = [col for col in trials.columns if col not in _TRIAL_TIME_COLUMNS] - for column in extra_columns: + for column in trials.columns: nwb_file.add_trial_column(name=column, description=descriptions.get(column, column)) for row_id, row in trials.iterrows(): - nwb_file.add_trial(id=int(row_id), **{column: row[column] for column in trials.columns}) + start_time, stop_time = cls._trial_extent(row) + nwb_file.add_trial( + id=int(row_id), + start_time=start_time, + stop_time=stop_time, + **{column: row[column] for column in trials.columns}, + ) # ------------------------------------------------------------------ # # Writers diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 73963ac..a0dcd9e 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -127,6 +127,30 @@ def _event_times(df: t.Optional[pd.DataFrame]) -> np.ndarray: return np.empty(0) return df.sort_index().index.to_numpy(dtype=float) + @staticmethod + def _time_at(times: np.ndarray, index: int) -> float: + """Return ``times[index]``, or ``NaN`` when the stream is that much shorter. + + The per-trial streams are paired by position, so a stream with fewer + events than there are trials (already reported by ``_check_aligned``) is + padded rather than raising. + + Parameters + ---------- + times : numpy.ndarray + Sorted per-trial event timestamps. + index : int + The trial index to read. + + Returns + ------- + float + The timestamp, or ``numpy.nan`` when ``index`` is out of range. + """ + if index < 0 or index >= times.size: + return np.nan + return float(times[index]) + @staticmethod def _closest_time_in_window(times: np.ndarray, start: float, stop: float) -> t.Optional[float]: """Return the timestamp in ``times`` closest to ``start`` within ``[start, stop)``. @@ -661,12 +685,57 @@ def _sample_lickspout( # ------------------------------------------------------------------ # # Per-trial assembly # ------------------------------------------------------------------ # + @classmethod + def _trial_periods( + cls, + index: int, + *, + quiescent_times: np.ndarray, + response_period_times: np.ndarray, + consumption_times: np.ndarray, + iti_times: np.ndarray, + ) -> t.Dict[str, float]: + """Return the start and stop time of every task period for one trial. + + Each software event marks the *start* of its period and the periods run + back-to-back — quiescent, response, reward consumption, ITI, then the + next trial's quiescent — so every period's stop time is the following + period's start time. The last trial's ITI has no following quiescent + event, so its stop time is ``NaN``. + + Parameters + ---------- + index : int + The trial index; every per-trial stream is paired by position. + quiescent_times, response_period_times, consumption_times, iti_times : numpy.ndarray + Per-trial timestamps of the ``QuiescentPeriod``, ``ResponsePeriod``, + ``RewardConsumptionPeriod``, and ``ItiPeriod`` streams. + + Returns + ------- + dict of str to float + The eight ``*_start_time`` / ``*_stop_time`` period columns; an entry + is ``NaN`` where the corresponding event is missing. + """ + response_start = cls._time_at(response_period_times, index) + consumption_start = cls._time_at(consumption_times, index) + iti_start = cls._time_at(iti_times, index) + return { + "quiescent_start_time": cls._time_at(quiescent_times, index), + "quiescent_stop_time": response_start, + "response_start_time": response_start, + "response_stop_time": consumption_start, + "reward_consumption_start_time": consumption_start, + "reward_consumption_stop_time": iti_start, + "ITI_start_time": iti_start, + "ITI_stop_time": cls._time_at(quiescent_times, index + 1), + } + def _build_row( self, *, outcome: TrialOutcome, - start: float, - stop: float, + periods: t.Dict[str, float], response: t.Any, side_bias: t.Optional[float], left_valve_open_time: t.Optional[float], @@ -675,14 +744,20 @@ def _build_row( session: t.Dict[str, t.Any], lickspout: t.Dict[str, t.Optional[float]], ) -> TrialConfig: - """Assemble a single ``TrialConfig`` from aligned per-trial inputs.""" + """Assemble a single ``TrialConfig`` from aligned per-trial inputs. + + ``periods`` holds the trial's period bounds (see :meth:`_trial_periods`); + the quiescent-period start through the ITI start is also the window used + to pick this trial's go cue out of the unaligned hardware stream. + """ trial = outcome.trial is_right_choice = outcome.is_right_choice is_rewarded = bool(outcome.is_rewarded) + start = periods["quiescent_start_time"] + stop = periods["ITI_start_time"] return TrialConfig( - start_time=start, - stop_time=stop, + **periods, delay_start_time=start, animal_response=self._animal_response(response), rewarded_historyL=self._rewarded_history(is_rewarded, is_right_choice, is_right=False), @@ -712,9 +787,10 @@ def _check_aligned(n_trials: int, counts: t.Mapping[str, int]) -> t.List[str]: """Return human-readable warnings for per-trial streams that mismatch ``n_trials``. The builder aligns the per-trial software-event streams *positionally*: - the i-th ``TrialOutcome`` is paired with the i-th ``QuiescentPeriod`` - (start), the i-th ``ItiPeriod`` (stop), and the i-th ``Response``. If any - of those streams has a different length than the ``TrialOutcome`` stream, + the i-th ``TrialOutcome`` is paired with the i-th event of each period + stream (``QuiescentPeriod``, ``ResponsePeriod``, + ``RewardConsumptionPeriod``, ``ItiPeriod``) and the i-th ``Response``. If + any of those streams has a different length than the ``TrialOutcome`` stream, the pairing slips and every subsequent row is silently misaligned (shorter streams are padded with ``NaN``/``None``; longer streams are ignored past ``n_trials``). @@ -745,18 +821,22 @@ def build(self) -> pd.DataFrame: """Build the trials table. The per-trial software-event streams (``TrialOutcome``, - ``QuiescentPeriod``, ``ItiPeriod``, ``Response``) are emitted once per - trial and are aligned here *by position*: row ``i`` draws its outcome, - start time, stop time, and response from index ``i`` of each stream. The - ``TrialOutcome`` stream defines the trial count; the lengths of the other - streams are checked against it before assembly (see ``_check_aligned``) - so a slipped stream surfaces as a warning (or a ``ValueError`` when - ``raise_on_error`` is set) rather than silently misaligned rows. + ``QuiescentPeriod``, ``ResponsePeriod``, ``RewardConsumptionPeriod``, + ``ItiPeriod``, ``Response``) are emitted once per trial and are aligned + here *by position*: row ``i`` draws its outcome, period start times, and + response from index ``i`` of each stream. The ``TrialOutcome`` stream + defines the trial count; the lengths of the other streams are checked + against it before assembly (see ``_check_aligned``) so a slipped stream + surfaces as a warning (or a ``ValueError`` when ``raise_on_error`` is set) + rather than silently misaligned rows. + + Each period event marks the start of its period, so the periods' stop + times come from the next event in sequence (see ``_trial_periods``). Hardware streams are handled per their nature: the go cue is an event - each trial selects within its ``[start, stop)`` window, while the valve - open duration is a constant session-configured supply-port pulse width - applied to every trial. + each trial selects within its ``[quiescent_start_time, ITI_start_time)`` + window, while the valve open duration is a constant session-configured + supply-port pulse width applied to every trial. Returns ------- @@ -772,6 +852,8 @@ def build(self) -> pd.DataFrame: """ outcomes = self._load("Behavior", "SoftwareEvents", "TrialOutcome") quiescent = self._load("Behavior", "SoftwareEvents", "QuiescentPeriod") + response_period = self._load("Behavior", "SoftwareEvents", "ResponsePeriod") + consumption = self._load("Behavior", "SoftwareEvents", "RewardConsumptionPeriod") iti = self._load("Behavior", "SoftwareEvents", "ItiPeriod") responses = self._load("Behavior", "SoftwareEvents", "Response") metrics = self._load("Behavior", "SoftwareEvents", "TrialMetrics") @@ -785,8 +867,10 @@ def build(self) -> pd.DataFrame: # Per-trial streams: one payload/timestamp per trial, aligned by index. outcome_payloads = self._event_payloads(outcomes) - start_times = self._event_times(quiescent) - stop_times = self._event_times(iti) + quiescent_times = self._event_times(quiescent) + response_period_times = self._event_times(response_period) + consumption_times = self._event_times(consumption) + iti_times = self._event_times(iti) response_payloads = self._event_payloads(responses) metric_payloads = self._event_payloads(metrics) @@ -796,8 +880,10 @@ def build(self) -> pd.DataFrame: warnings = self._check_aligned( n_trials, { - "QuiescentPeriod (start_time)": start_times.size, - "ItiPeriod (stop_time)": stop_times.size, + "QuiescentPeriod (quiescent_start_time)": quiescent_times.size, + "ResponsePeriod (response_start_time)": response_period_times.size, + "RewardConsumptionPeriod (reward_consumption_start_time)": consumption_times.size, + "ItiPeriod (ITI_start_time)": iti_times.size, "Response": len(response_payloads), "TrialMetrics (side_bias)": len(metric_payloads), }, @@ -834,17 +920,25 @@ def build(self) -> pd.DataFrame: outcome = self._parse_outcome(outcome_payload) # Pad with NaN/None when a stream is shorter than the trial count; # _check_aligned has already warned about any such mismatch. - start = float(start_times[i]) if i < start_times.size else np.nan - stop = float(stop_times[i]) if i < stop_times.size else np.nan + periods = self._trial_periods( + i, + quiescent_times=quiescent_times, + response_period_times=response_period_times, + consumption_times=consumption_times, + iti_times=iti_times, + ) response = response_payloads[i] if i < len(response_payloads) else None side_bias = self._side_bias(metric_payloads[i] if i < len(metric_payloads) else None) - lickspout = self._sample_lickspout(lickspout_positions, start, stop) + lickspout = self._sample_lickspout( + lickspout_positions, + periods["quiescent_start_time"], + periods["ITI_start_time"], + ) rows.append( self._build_row( outcome=outcome, - start=start, - stop=stop, + periods=periods, response=response, side_bias=side_bias, left_valve_open_time=left_valve_open_time, diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index ed3425b..494d5fd 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -18,9 +18,49 @@ class TrialConfig(BaseModel): NWB trial-column description. """ - # --- Trial timing (NWB built-ins; no entry in the column-info JSON) --- - start_time: float = Field(description="Trial start time (QuiescentPeriod timestamp).") - stop_time: float = Field(description="Trial stop time (ItiPeriod timestamp).") + # --- Trial period timing (one start/stop pair per task period) --- + # The four periods run back-to-back in this order, each software event marking + # the *start* of its period, so each period's stop is the next period's start: + # quiescent -> response -> reward consumption -> ITI -> (next trial's + # quiescent). No entry in the column-info JSON. + quiescent_start_time: float = Field( + description=( + "Start time of the quiescent period (QuiescentPeriod timestamp). The quiescent period is the lick-free delay preceding the go cue; each lick restarts it, so its realized duration can exceed the configured delay_duration." + ), + ) + quiescent_stop_time: float = Field( + description=( + "End time of the quiescent period, i.e. the start of the response period (ResponsePeriod timestamp); the go cue is played at this boundary." + ), + ) + response_start_time: float = Field( + description=( + "Start time of the response period (ResponsePeriod timestamp), when the go cue is played." + ), + ) + response_stop_time: float = Field( + description=( + "End time of the response period, i.e. the start of the reward consumption period (RewardConsumptionPeriod timestamp). This is when the animal responded, or the response deadline for an ignored trial." + ), + ) + reward_consumption_start_time: float = Field( + description=( + "Start time of the reward consumption period (RewardConsumptionPeriod timestamp)." + ), + ) + reward_consumption_stop_time: float = Field( + description=( + "End time of the reward consumption period, i.e. the start of the inter-trial interval (ItiPeriod timestamp)." + ), + ) + ITI_start_time: float = Field( + description="Start time of the inter-trial interval (ItiPeriod timestamp).", + ) + ITI_stop_time: float = Field( + description=( + "End time of the inter-trial interval, i.e. the start of the next trial's quiescent period (the following QuiescentPeriod timestamp); NaN on the last trial of the session." + ), + ) # --- trial_info --- animal_response: int = Field( @@ -37,7 +77,7 @@ class TrialConfig(BaseModel): delay_start_time: Optional[float] = Field( default=None, description=( - "Start time of the delay (quiescent) period preceding the go cue; equals the trial start time (QuiescentPeriod timestamp)." + "Legacy name for the start of the quiescent period (QuiescentPeriod timestamp); the 'delay' of the legacy delay_* columns is the acquisition software's quiescent period, so this always equals quiescent_start_time." ), ) goCue_start_time: Optional[float] = Field(default=None, description="The go cue start time") @@ -109,7 +149,10 @@ class TrialConfig(BaseModel): default=None, description="The maximum duration(s) allowed for each delay" ) delay_duration: Optional[float] = Field( - default=None, description="The duration (s) between delay start and go cue start" + default=None, + description=( + "The configured duration (s) of the delay (quiescent) period between delay start and go cue start. Each lick restarts the quiescent period, so the realized duration (quiescent_stop_time - quiescent_start_time) can be longer." + ), ) # --- ITI_duration --- diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py index 8967e5e..1124adc 100644 --- a/tests/test_pipeline/test_pipeline.py +++ b/tests/test_pipeline/test_pipeline.py @@ -35,11 +35,16 @@ def _make_pipeline() -> Pipeline: def _trials_frame() -> pd.DataFrame: - """Two-trial table with the time columns plus one modeled, one unmodeled column.""" + """Two-trial table with the period columns plus one modeled, one unmodeled column. + + The second trial's ``ITI_stop_time`` is ``NaN``, as it is for the last trial + of a session. + """ return pd.DataFrame( { - "start_time": [0.0, 1.0], - "stop_time": [0.5, 1.5], + "quiescent_start_time": [0.0, 1.0], + "ITI_start_time": [0.4, 1.4], + "ITI_stop_time": [1.0, np.nan], "animal_response": [0, 1], "not_in_model": [7, 8], } @@ -268,7 +273,7 @@ def test_add_acquisition_table_builds_dynamic_table(): def test_add_trials_populates_columns_and_rows(): - """Extra columns are described via ``TrialConfig``; rows are added.""" + """Every trials column is registered and described via ``TrialConfig``.""" nwb_file = MagicMock() trials = _trials_frame() @@ -278,8 +283,8 @@ def test_add_trials_populates_columns_and_rows(): call.kwargs["name"]: call.kwargs["description"] for call in nwb_file.add_trial_column.call_args_list } - # start_time / stop_time are native and not registered as extra columns. - assert set(added_columns) == {"animal_response", "not_in_model"} + # The table has no trial start/stop columns, so every column is an extra one. + assert set(added_columns) == set(trials.columns) assert added_columns["animal_response"] == TrialConfig.column_descriptions()["animal_response"] # A column absent from TrialConfig falls back to its own name as description. assert added_columns["not_in_model"] == "not_in_model" @@ -288,6 +293,33 @@ def test_add_trials_populates_columns_and_rows(): assert [call.kwargs["id"] for call in nwb_file.add_trial.call_args_list] == [0, 1] +def test_add_trials_derives_native_start_and_stop_from_periods(): + """NWB's native trial extent spans the quiescent start to the ITI end. + + The last trial has no ITI end (no following quiescent period), so its stop + time falls back to the ITI start. + """ + nwb_file = MagicMock() + + Pipeline._add_trials(nwb_file, _trials_frame()) + + extents = [ + (call.kwargs["start_time"], call.kwargs["stop_time"]) + for call in nwb_file.add_trial.call_args_list + ] + assert extents == [(0.0, 1.0), (1.0, 1.4)] + + +def test_add_trials_skips_frame_without_period_columns(): + """A table missing the columns the trial extent is derived from adds nothing.""" + nwb_file = MagicMock() + + Pipeline._add_trials(nwb_file, pd.DataFrame({"animal_response": [0, 1]})) + + nwb_file.add_trial_column.assert_not_called() + nwb_file.add_trial.assert_not_called() + + def test_add_trials_skips_empty_frame(): """An empty trials table adds nothing.""" nwb_file = MagicMock() diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index 65c0eb1..e615bc8 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -215,8 +215,12 @@ def _full_dataset(): ], ) ), + # The four period streams, each emitted at its period's start: + # quiescent -> response -> reward consumption -> ITI, per trial. "QuiescentPeriod": _Stream(_events([10.0, 20.0], [None, None])), - "ItiPeriod": _Stream(_events([20.0, 30.0], [None, None])), + "ResponsePeriod": _Stream(_events([11.0, 21.0], [None, None])), + "RewardConsumptionPeriod": _Stream(_events([12.0, 22.0], [None, None])), + "ItiPeriod": _Stream(_events([15.0, 25.0], [None, None])), "Response": _Stream( _events( [10.5, 20.5], [{"Item1": 10.5, "Item2": True}, {"Item1": 20.5, "Item2": None}] @@ -284,9 +288,16 @@ def test_build_full_dataset(): first, second = table.iloc[0], table.iloc[1] - # Trial windows. - assert first["start_time"] == 10.0 and first["stop_time"] == 20.0 - assert first["delay_start_time"] == 10.0 + # Period bounds: each period ends where the next one starts, and the ITI + # ends at the next trial's quiescent period (NaN on the last trial). + assert first["quiescent_start_time"] == 10.0 and first["quiescent_stop_time"] == 11.0 + assert first["response_start_time"] == 11.0 and first["response_stop_time"] == 12.0 + assert first["reward_consumption_start_time"] == 12.0 + assert first["reward_consumption_stop_time"] == 15.0 + assert first["ITI_start_time"] == 15.0 and first["ITI_stop_time"] == 20.0 + assert second["ITI_start_time"] == 25.0 and np.isnan(second["ITI_stop_time"]) + # delay_start_time is the legacy name for the quiescent period start. + assert first["delay_start_time"] == first["quiescent_start_time"] == 10.0 # Response encoding: True -> right (1), None -> no response (2). assert first["animal_response"] == 1 From 1408e67f84b765ebee28ffa2e8e35d75e185d449 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 <109561860+arjunsridhar12345@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:25:37 -0700 Subject: [PATCH 2/2] feat: add columns for anti bias intervention to trials table and qc plot (#53) * feat: add anti-bias intervention to trial table and qc plot * test: update tests * docs: update trial table doc * chore: fix linting * fix: restore changes lost in merge conflict * refactor: use convention red for right and blue for left * feat: add ticks for anti bias in lickspout plot * feat: add ticks to side bias plots and adjust color --- docs/trials_table_mapping.md | 3 + .../processing/_trial_table.py | 98 +++++++++++ .../processing/models/trial_config.py | 20 +++ .../qc/processed/plots.py | 163 +++++++++++++++++- .../qc/processed/results.py | 6 + tests/test_processing/test_trial_table.py | 149 +++++++++++++++- tests/test_qc/test_plots.py | 60 +++++++ 7 files changed, 492 insertions(+), 7 deletions(-) diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index 12223b3..acce921 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -62,6 +62,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. | +| `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`. | | `reward_consumption_duration` | `Trial -> reward_consumption_duration`. | @@ -182,6 +184,7 @@ These were mapped during exploration but are no longer in scope: | 2026-06-17 | `auto_waterL` / `auto_waterR` now encode no auto-response (`is_auto_reward_right` is `None`) and missing trials as `0` instead of `NULL`. The columns are non-nullable (`int`, default `0`). | | 2026-06-20 | Added `reward_size_left` / `reward_size_right` (reward volume in uL) from `task_parameters.reward_size`, and `side_bias` from the per-trial `TrialMetrics` event (`bias` field). | | 2026-06-20 | `reward_probabilityL` / `reward_probabilityR` now read the block probability from `trial.metadata.p_reward_left` / `p_reward_right` instead of the top-level per-trial `trial.p_reward_left` / `p_reward_right`. | +| 2026-07-27 | Added `anti_bias_left_water` / `anti_bias_right_water` (boolean anti-bias water interventions per side) and `anti_bias_lickspout_movement` (mm the anti-bias algorithm shifted the lickspouts) from `TrialOutcome`'s `trial.metadata.extra` (`is_bias_water_intervention` / `is_bias_stage_intervention`), `is_auto_reward_right`, and `lickspout_offset_delta`. These are also overlaid on the QC `side_bias.png` figure. | | 2026-07-22 | `lickspout_position_x` / `y1` / `y2` / `z` now derive from the `HarpManipulator` `AccumulatedSteps` stream (microsteps → mm via the `InputSchemas.Rig` manipulator calibration, `full_step_to_mm / microstep_resolution`), sampled per trial via the closest sample in the `[start_time, stop_time)` window and re-referenced to the session-start position (displacement relative to session start, mm), replacing the static `InitialManipulatorPosition` software event. `Motor{i}` maps to `Axis(i + 1)` (X, Y1, Y2, Z). The rig and `AccumulatedSteps` streams are required when there are trials (`build` raises if either is missing). Column descriptions corrected from `um` to `mm`. | | 2026-07-24 | `reward_size_left` / `reward_size_right` moved from session-level `task_parameters.reward_size` to per-trial `Trial.reward_size` (fields `.left` / `.right`). The columns are now nullable — `None` when the trial is missing. A missing `TaskLogic` stream no longer raises; session distribution columns are simply null. `min_reward_each_block` moved from `CoupledTrialGenerator` to `CoupledWarmupTrialGenerator`. | | 2026-08-06 | **Breaking:** the trial `start_time` / `stop_time` columns are removed and replaced by one start/stop pair per task period: `quiescent_start_time` / `quiescent_stop_time`, `response_start_time` / `response_stop_time`, `reward_consumption_start_time` / `reward_consumption_stop_time`, and `ITI_start_time` / `ITI_stop_time`, read from the `ResponsePeriod` and `RewardConsumptionPeriod` streams in addition to `QuiescentPeriod` and `ItiPeriod`. Each period event marks its period's start, so each stop is the next period's start; `ITI_stop_time` is the next trial's `QuiescentPeriod` timestamp (`NaN` on the last trial). The two new streams are also checked for positional alignment with `TrialOutcome`. NWB's required native `start_time` / `stop_time` are now derived when writing (`quiescent_start_time` → `ITI_stop_time`, falling back to `ITI_start_time`), so the NWB trials table changes in two ways: the old `start_time` / `stop_time` columns are gone, and the native trial extent now ends at the *end* of the ITI rather than at its start. | diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index a0dcd9e..9daaab9 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -13,6 +13,9 @@ from aind_behavior_dynamic_foraging.rig import AindDynamicForagingRig from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic from aind_behavior_dynamic_foraging.task_logic.trial_generators import TrialGeneratorSpec +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, TrialMetrics, @@ -428,6 +431,97 @@ def _auto_water(trial: Trial, *, is_right: bool) -> int: 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. + + 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. + + 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() + + @staticmethod + def _anti_bias_water( + trial: Trial, bias_metadata: BlockBasedTrialMetadata, *, is_right: bool + ) -> bool: + """Return whether the anti-bias algorithm watered the requested side. + + The anti-bias algorithm delivers its water intervention through the same + auto-response channel as ordinary autowater (``is_auto_reward_right``: + ``True`` right, ``False`` left), so the two are distinguished only by the + ``is_bias_water_intervention`` flag. This is ``True`` only when the trial + was a bias-water intervention *and* the auto-response was to the + requested side. + + 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 + ------- + bool + Whether an anti-bias water intervention targeted the requested side. + """ + if not bias_metadata.is_bias_water_intervention: + return False + return trial.is_auto_reward_right is is_right + + @staticmethod + def _anti_bias_lickspout_movement( + trial: Trial, bias_metadata: BlockBasedTrialMetadata + ) -> float: + """Return the anti-bias lickspout displacement (mm) for this trial. + + The anti-bias algorithm's other intervention shifts the lickspouts + horizontally; the per-trial displacement is ``trial.lickspout_offset_delta`` + (positive is rightward). Reported only when the trial is flagged as a + bias-stage intervention, so a stray offset from another source is not + attributed to the anti-bias algorithm; ``0.0`` otherwise. + + Parameters + ---------- + trial : Trial + The per-trial task-logic model. + bias_metadata : BlockBasedTrialMetadata + The trial's extra metadata (see ``_bias_metadata``). + + Returns + ------- + float + The signed displacement (mm), or ``0.0`` when there was no + lickspout intervention. + """ + if not bias_metadata.is_bias_stage_intervention: + return 0.0 + return trial.lickspout_offset_delta + @staticmethod def _block_reward_probability(trial: Trial, *, is_right: bool) -> t.Optional[float]: """Return the block reward probability for a side from the trial metadata. @@ -753,6 +847,7 @@ def _build_row( trial = outcome.trial is_right_choice = outcome.is_right_choice is_rewarded = bool(outcome.is_rewarded) + bias_metadata = self._bias_metadata(trial) start = periods["quiescent_start_time"] stop = periods["ITI_start_time"] @@ -778,6 +873,9 @@ def _build_row( delay_duration=trial.quiescence_period_duration, auto_waterL=self._auto_water(trial, is_right=False), auto_waterR=self._auto_water(trial, 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), **session, **lickspout, ) diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index 494d5fd..f34a57a 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -183,6 +183,26 @@ class TrialConfig(BaseModel): auto_waterL: int = Field(default=0, description="Autowater given at Left") auto_waterR: int = Field(default=0, description="Autowater given at Right") + # --- 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." + ), + ) + 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." + ), + ) + anti_bias_lickspout_movement: float = Field( + default=0.0, + description=( + "Horizontal distance (mm) the lickspouts were moved by the anti-bias algorithm on this trial (positive is rightward); 0 when no lickspout intervention occurred." + ), + ) + # --- lickspout_position (mapping's `lickspout_positions` -> these four components) --- lickspout_position_x: Optional[float] = Field( default=None, diff --git a/src/dynamic_foraging_processing/qc/processed/plots.py b/src/dynamic_foraging_processing/qc/processed/plots.py index 7750ace..4f40673 100644 --- a/src/dynamic_foraging_processing/qc/processed/plots.py +++ b/src/dynamic_foraging_processing/qc/processed/plots.py @@ -22,6 +22,11 @@ lick_latency_by_side, ) +#: 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 +#: this far below, so the pair reads as arrows straddling y=0. +_MOVE_MARKER_OFFSET = 0.06 + def plot_lick_intervals( left_lick_times: np.ndarray, right_lick_times: np.ndarray, results_folder: str @@ -138,8 +143,31 @@ def plot_lick_latency( return LICK_LATENCY_PLOT -def _add_bias_plot(ax: plt.Axes, side_bias: np.ndarray) -> None: - """Draw the per-trial side-bias trace from the trial-table column.""" +def _legend_outside(ax: plt.Axes) -> None: + """Place the axes legend just outside the right edge of the panel. + + The trial traces span the full width of these panels, so an in-axes legend + sits on top of the data. Anchoring it outside keeps the trace readable; the + figure is saved with ``bbox_inches="tight"``, so the legend is not clipped. + """ + ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), borderaxespad=0.0, fontsize="x-small") + + +def _add_bias_plot( + ax: plt.Axes, + side_bias: np.ndarray, + 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, +) -> None: + """Draw the per-trial side-bias trace with anti-bias interventions overlaid. + + The anti-bias algorithm pushes against a developing side bias, so its two + interventions are drawn on top of the bias trace they respond to: water + interventions as short ticks at the top (right port) and bottom (left + port), and lickspout movements as triangles straddling the zero-bias line, + pointing (and coloured) in the direction the spout was moved. + """ ax.set_xlabel("Trial #") ax.set_ylabel("Side Bias") ax.axhline(+0.7, color="r", linestyle="--") @@ -153,6 +181,67 @@ def _add_bias_plot(ax: plt.Axes, side_bias: np.ndarray) -> None: if len(bias): ax.set_xlim([0, len(bias)]) + plotted = False + if anti_bias_right_water is not None: + right = np.where(np.asarray(anti_bias_right_water, dtype=bool))[0] + ax.vlines(right, 0.9, 1.0, color="darkred", linewidth=1, label="Anti-bias water (R)") + plotted = True + if anti_bias_left_water is not None: + left = np.where(np.asarray(anti_bias_left_water, dtype=bool))[0] + ax.vlines(left, -1.0, -0.9, color="darkblue", linewidth=1, label="Anti-bias water (L)") + plotted = True + if anti_bias_lickspout_movement is not None: + move = np.asarray(anti_bias_lickspout_movement, dtype=float) + # Direction-coded markers straddling the zero-bias line: a rightward move + # is a red up-triangle above it, a leftward move a blue down-triangle + # below it, matching the red/blue sides used elsewhere in the figure. + for indices, offset, marker, color, side in ( + (np.where(move > 0)[0], _MOVE_MARKER_OFFSET, "^", "red", "R"), + (np.where(move < 0)[0], -_MOVE_MARKER_OFFSET, "v", "blue", "L"), + ): + ax.plot( + indices, + np.full(len(indices), offset), + marker=marker, + color=color, + linestyle="none", + markersize=6, + label=f"Anti-bias lickspout move ({side})", + ) + plotted = True + if plotted: + _legend_outside(ax) + + +def _moved_trials(positions: t.Sequence[t.Optional[np.ndarray]]) -> np.ndarray: + """Return the trial indices where any lickspout axis changed position. + + A move is detected as a change from the previous trial on any of the + supplied axes, so a trial is flagged once no matter how many axes shifted. + NaN samples (missing position readings) are ignored rather than counted as + a change. + + Parameters + ---------- + positions : sequence of numpy.ndarray or None + Per-trial position arrays (one per axis); ``None`` and empty arrays are + skipped. + + Returns + ------- + numpy.ndarray + Sorted, unique trial indices at which the lickspouts moved. + """ + moved: t.Set[int] = set() + for position in positions: + if position is None or len(position) < 2: + continue + values = np.asarray(position, dtype=float) + delta = np.diff(values) + changed = np.where(~np.isnan(delta) & (delta != 0))[0] + 1 + moved.update(int(index) for index in changed) + return np.array(sorted(moved), dtype=int) + def _add_lickspout_position_plot( ax: plt.Axes, @@ -160,8 +249,15 @@ def _add_lickspout_position_plot( lickspout_y1: t.Optional[np.ndarray], lickspout_y2: t.Optional[np.ndarray], lickspout_z: t.Optional[np.ndarray], + anti_bias_lickspout_movement: t.Optional[np.ndarray] = None, ) -> None: - """Draw lickspout x/y/z positions relative to session start (mm).""" + """Draw lickspout x/y/z positions relative to session start (mm). + + Trials where the spouts moved are marked with ticks along the bottom of the + panel, split by who moved them: the anti-bias algorithm (same flag as the + markers on the side-bias panel) versus the experimenter, which is any other + change in position. + """ ax.set_xlabel("Trial #") ax.set_ylabel("Lickspout Position \n relative to session start (mm)") positions = [ @@ -179,8 +275,39 @@ def _add_lickspout_position_plot( # trial-table builder), so plot them directly. ax.plot(values, color, label=label) plotted = True + + # Ticks sit in axes-fraction coordinates on y so they stay pinned to the + # bottom of the panel whatever the position data's range is. + transform = ax.get_xaxis_transform() + automatic = np.array([], dtype=int) + if anti_bias_lickspout_movement is not None: + move = np.asarray(anti_bias_lickspout_movement, dtype=float) + automatic = np.where(move != 0)[0] + ax.vlines( + automatic, + 0.0, + 0.08, + transform=transform, + color="g", + linewidth=1, + label="Automatic move", + ) + plotted = True + moved = _moved_trials([lickspout_x, lickspout_y1, lickspout_y2, lickspout_z]) + manual = np.setdiff1d(moved, automatic) + if len(manual): + ax.vlines( + manual, + 0.0, + 0.08, + transform=transform, + color="k", + linewidth=1, + label="Manual move", + ) + plotted = True if plotted: - ax.legend() + _legend_outside(ax) def _time_to_trial_index(go_cue_times: np.ndarray, times: np.ndarray) -> t.List[int]: @@ -305,6 +432,9 @@ def plot_side_bias( autowater_right: t.Optional[np.ndarray] = None, manual_left_times: t.Optional[np.ndarray] = None, manual_right_times: t.Optional[np.ndarray] = None, + 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, ) -> str: """Save the four-panel side-bias figure. @@ -331,6 +461,14 @@ def plot_side_bias( Per-trial autowater indicator arrays. manual_left_times, manual_right_times : numpy.ndarray, optional Manual-water delivery timestamps (s). + 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. + anti_bias_lickspout_movement : numpy.ndarray, optional + Per-trial signed lickspout displacement (mm) applied by the anti-bias + algorithm; nonzero trials are marked on the side-bias trace and ticked + as automatic moves on the lickspout-position panel, where any other + change in position is ticked as a manual move. Returns ------- @@ -343,8 +481,21 @@ def plot_side_bias( axis.spines["top"].set_visible(False) axis.spines["right"].set_visible(False) - _add_bias_plot(ax[0], side_bias) - _add_lickspout_position_plot(ax[1], lickspout_x, lickspout_y1, lickspout_y2, lickspout_z) + _add_bias_plot( + ax[0], + side_bias, + anti_bias_left_water=anti_bias_left_water, + anti_bias_right_water=anti_bias_right_water, + anti_bias_lickspout_movement=anti_bias_lickspout_movement, + ) + _add_lickspout_position_plot( + ax[1], + lickspout_x, + lickspout_y1, + lickspout_y2, + lickspout_z, + anti_bias_lickspout_movement=anti_bias_lickspout_movement, + ) _add_behavior_plot( ax[2], animal_response, diff --git a/src/dynamic_foraging_processing/qc/processed/results.py b/src/dynamic_foraging_processing/qc/processed/results.py index 0657e65..1c708b3 100644 --- a/src/dynamic_foraging_processing/qc/processed/results.py +++ b/src/dynamic_foraging_processing/qc/processed/results.py @@ -42,6 +42,9 @@ "reward_probability_right": "reward_probabilityR", "autowater_left": "auto_waterL", "autowater_right": "auto_waterR", + "anti_bias_left_water": "anti_bias_left_water", + "anti_bias_right_water": "anti_bias_right_water", + "anti_bias_lickspout_movement": "anti_bias_lickspout_movement", "go_cue_times": "goCue_start_time", } @@ -123,6 +126,9 @@ def behavior_qc_results( autowater_right=_column(trials, "autowater_right"), manual_left_times=manual_left_times, manual_right_times=manual_right_times, + 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"), ) plot_lick_intervals(left_lick_times, right_lick_times, results_folder) plot_lick_latency( diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index e615bc8..e5ebe3f 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -93,6 +93,9 @@ def _outcome( block_p_right=None, reward_size_left=None, reward_size_right=None, + lickspout_offset_delta=None, + is_bias_water_intervention=None, + is_bias_stage_intervention=None, ): """Build a serialized ``TrialOutcome`` payload (dict, as delivered by the reader). @@ -100,7 +103,10 @@ def _outcome( fields); ``block_p_left`` / ``block_p_right``, when given, are the block 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). + 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. """ trial = { "p_reward_left": p_left, @@ -111,6 +117,8 @@ def _outcome( "inter_trial_interval_duration": 4.0, "is_auto_reward_right": auto, } + if lickspout_offset_delta is not None: + trial["lickspout_offset_delta"] = lickspout_offset_delta if reward_size_left is not None or reward_size_right is not None: trial["reward_size"] = { "left": reward_size_left if reward_size_left is not None else 2.0, @@ -118,6 +126,12 @@ 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: + metadata = trial.setdefault("metadata", {}) + metadata["extra"] = { + "is_bias_water_intervention": bool(is_bias_water_intervention), + "is_bias_stage_intervention": bool(is_bias_stage_intervention), + } return { "trial": trial, "is_right_choice": is_right_choice, @@ -356,6 +370,13 @@ def test_build_full_dataset(): assert first["side_bias"] == pytest.approx(0.3) assert pd.isna(second["side_bias"]) + # No anti-bias interventions in the base fixture -> inert defaults. + assert bool(first["anti_bias_left_water"]) is False + assert bool(first["anti_bias_right_water"]) is False + assert first["anti_bias_lickspout_movement"] == 0.0 + assert bool(second["anti_bias_left_water"]) is False + assert second["anti_bias_lickspout_movement"] == 0.0 + def test_build_missing_task_logic_leaves_session_columns_null(): """A missing TaskLogic stream still builds; session distribution columns are null.""" @@ -622,6 +643,132 @@ def test_auto_water_encodes_side_from_auto_response(): assert TrialTableBuilder._auto_water(no_auto, is_right=True) == 0 +def test_bias_metadata_parses_dict_model_and_default(): + """``_bias_metadata`` handles a dict extra, a model extra, and missing metadata.""" + from aind_behavior_dynamic_foraging.task_logic.trial_generators.block_based_trial_generator import ( + BlockBasedTrialMetadata, + ) + + # Dict extra (as delivered off the stream) is validated into the model. + from_dict = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, is_bias_water_intervention=True) + ).trial + assert TrialTableBuilder._bias_metadata(from_dict).is_bias_water_intervention is True + + # A ``BlockBasedTrialMetadata`` instance is returned as-is. + trial = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, block_p_left=0.5) + ).trial + trial.metadata.extra = BlockBasedTrialMetadata(is_bias_stage_intervention=True) + assert TrialTableBuilder._bias_metadata(trial).is_bias_stage_intervention is True + + # No metadata -> all-False default. + no_meta = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True) + ).trial + assert no_meta.metadata is None + default = TrialTableBuilder._bias_metadata(no_meta) + assert default.is_bias_water_intervention is False + assert default.is_bias_stage_intervention is False + + # Metadata present but a non-dict / non-model extra -> default. + other = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, block_p_left=0.5) + ).trial + other.metadata.extra = "unexpected" + assert TrialTableBuilder._bias_metadata(other).is_bias_water_intervention is False + + +def test_anti_bias_water_gated_on_intervention_flag_and_side(): + """Anti-bias water is True only for a bias-water intervention on the matching side.""" + # Right-side bias-water intervention. + right = 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(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 + + # 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) + ).trial + auto_meta = TrialTableBuilder._bias_metadata(autowater) + assert TrialTableBuilder._anti_bias_water(autowater, auto_meta, is_right=False) is False + + +def test_anti_bias_lickspout_movement_gated_on_stage_flag(): + """Movement is the offset delta only when flagged a bias-stage intervention.""" + moved = TrialOutcome.model_validate( + _outcome( + 1.0, + 1.0, + is_right_choice=True, + is_rewarded=True, + lickspout_offset_delta=1.5, + is_bias_stage_intervention=True, + ) + ).trial + assert TrialTableBuilder._anti_bias_lickspout_movement( + moved, TrialTableBuilder._bias_metadata(moved) + ) == pytest.approx(1.5) + + # Offset present but not flagged as a stage intervention -> 0.0. + unflagged = TrialOutcome.model_validate( + _outcome(1.0, 1.0, is_right_choice=True, is_rewarded=True, lickspout_offset_delta=1.5) + ).trial + assert ( + TrialTableBuilder._anti_bias_lickspout_movement( + unflagged, TrialTableBuilder._bias_metadata(unflagged) + ) + == 0.0 + ) + + +def test_build_populates_anti_bias_columns(): + """A dataset with anti-bias interventions populates the three anti-bias columns.""" + dataset = _full_dataset() + software_events = dataset.children["Behavior"].children["SoftwareEvents"] + software_events.children["TrialOutcome"] = _Stream( + _events( + [10.1, 20.1], + [ + _outcome( + 1.0, + 1.0, + is_right_choice=True, + is_rewarded=True, + auto=True, + is_bias_water_intervention=True, + ), + _outcome( + 1.0, + 1.0, + is_right_choice=True, + is_rewarded=True, + lickspout_offset_delta=-0.8, + is_bias_stage_intervention=True, + ), + ], + ) + ) + table = TrialTableBuilder(dataset).build() + first, second = table.iloc[0], table.iloc[1] + assert bool(first["anti_bias_right_water"]) is True + assert bool(first["anti_bias_left_water"]) is False + assert first["anti_bias_lickspout_movement"] == 0.0 + assert bool(second["anti_bias_right_water"]) is False + assert second["anti_bias_lickspout_movement"] == pytest.approx(-0.8) + + def test_block_reward_probability_reads_metadata_not_trial(): """The block probability comes from ``trial.metadata``, not the top-level p_reward.""" trial = TrialOutcome.model_validate( diff --git a/tests/test_qc/test_plots.py b/tests/test_qc/test_plots.py index 1366b01..a157f27 100644 --- a/tests/test_qc/test_plots.py +++ b/tests/test_qc/test_plots.py @@ -2,6 +2,7 @@ import os +import matplotlib.pyplot as plt import numpy as np from dynamic_foraging_processing.qc.processed import plots as _plots @@ -67,11 +68,70 @@ def test_plot_side_bias_full_inputs(tmp_path): 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]), + 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]), ) assert name == _plots.SIDE_BIAS_PLOT assert os.path.exists(tmp_path / name) +def test_add_bias_plot_lickspout_markers_are_direction_coded(): + """Rightward moves are red up-triangles above y=0, leftward blue down below.""" + fig, ax = plt.subplots() + _plots._add_bias_plot( + ax, + np.array([0.1, -0.2, 0.3]), + # Trial 0 moves right, trial 2 moves left, trial 1 does not move. + anti_bias_lickspout_movement=np.array([0.5, 0.0, -0.4]), + ) + markers = { + line.get_label(): line + for line in ax.get_lines() + if str(line.get_label()).startswith("Anti-bias lickspout move") + } + right = markers["Anti-bias lickspout move (R)"] + left = markers["Anti-bias lickspout move (L)"] + assert right.get_marker() == "^" and right.get_color() == "red" + assert left.get_marker() == "v" and left.get_color() == "blue" + # Only the moving trials are marked, each just off the zero-bias line. + assert list(right.get_xdata()) == [0] + assert list(right.get_ydata()) == [_plots._MOVE_MARKER_OFFSET] + assert list(left.get_xdata()) == [2] + assert list(left.get_ydata()) == [-_plots._MOVE_MARKER_OFFSET] + plt.close(fig) + + +def test_moved_trials_ignores_missing_and_unchanged_axes(): + """Only real per-trial changes count as moves; NaN and short arrays don't.""" + moved = _plots._moved_trials( + [ + None, + np.array([1.0]), # too short to diff + np.array([2.0, 2.0, 2.5, 2.5]), # change at trial 2 + np.array([3.0, np.nan, 3.0, 4.0]), # NaN ignored; change at trial 3 + ] + ) + assert moved.tolist() == [2, 3] + + +def test_add_lickspout_position_plot_splits_automatic_and_manual_moves(): + """Anti-bias trials tick as automatic; every other move ticks as manual.""" + fig, ax = plt.subplots() + _plots._add_lickspout_position_plot( + ax, + np.array([0.0, 1.0, 1.0, 2.0]), # moves at trials 1 and 3 + None, + None, + None, + anti_bias_lickspout_movement=np.array([0.0, 0.5, 0.0, 0.0]), + ) + labels = [text.get_text() for text in ax.get_legend().get_texts()] + assert "Automatic move" in labels + assert "Manual move" in labels + plt.close(fig) + + def test_plot_side_bias_minimal_inputs(tmp_path): """With only choices supplied, the optional panels are skipped cleanly.""" name = _plots.plot_side_bias(np.array([]), np.array([]), str(tmp_path))