From 700199df9da7f548e71980df87310de7510f9a6e Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 11:45:01 -0700 Subject: [PATCH 1/5] fix: propogate nan value of ITI_stop_time for last trial to NWB stop_time column --- .../pipeline/_pipeline.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/dynamic_foraging_processing/pipeline/_pipeline.py b/src/dynamic_foraging_processing/pipeline/_pipeline.py index 2a08af9..00390f5 100644 --- a/src/dynamic_foraging_processing/pipeline/_pipeline.py +++ b/src/dynamic_foraging_processing/pipeline/_pipeline.py @@ -59,10 +59,11 @@ #: ``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") +#: Trials-table column NWB's required native ``stop_time`` is taken from: the +#: end of the ITI. This is ``NaN`` on the last trial of the session (whose ITI +#: end is unknown) and the ``NaN`` is propagated rather than substituted, so an +#: unknown trial end reads as unknown instead of as a shortened trial. +_NWB_STOP_COLUMN = "ITI_stop_time" #: Source repository recorded in the ``processing.json`` data process. _CODE_URL = "https://github.com/AllenNeuralDynamics/dynamic-foraging-processing" @@ -236,8 +237,9 @@ def _trial_extent(row: pd.Series) -> t.Tuple[float, float]: 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). + quiescent period and ends with the ITI. The last trial of the session has + no ITI end, so its stop time is ``NaN`` — an unknown end is reported as + unknown rather than substituted with an earlier landmark. Parameters ---------- @@ -247,11 +249,10 @@ def _trial_extent(row: pd.Series) -> t.Tuple[float, float]: Returns ------- tuple of (float, float) - The trial start and stop time (seconds). + The trial start and stop time (seconds); the stop time is ``NaN`` + where the ITI end is unknown. """ - 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) + return float(row[_NWB_START_COLUMN]), float(row[_NWB_STOP_COLUMN]) @classmethod def _add_trials(cls, nwb_file: pynwb.NWBFile, trials: pd.DataFrame) -> None: @@ -265,7 +266,7 @@ def _add_trials(cls, nwb_file: pynwb.NWBFile, trials: pd.DataFrame) -> None: (named ``id``) is replicated as each trial's NWB ``id``. An empty table (or one missing the period columns the extent is derived from) is skipped. """ - required = (_NWB_START_COLUMN, *_NWB_STOP_COLUMNS) + required = (_NWB_START_COLUMN, _NWB_STOP_COLUMN) if trials.empty or any(col not in trials.columns for col in required): return descriptions = TrialConfig.column_descriptions() From 45f7f797f4abd1a0701cf6dea4513e41267ee652 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Mon, 17 Aug 2026 11:45:09 -0700 Subject: [PATCH 2/5] test: update tests --- tests/test_pipeline/test_pipeline.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py index c897ea7..d7650d3 100644 --- a/tests/test_pipeline/test_pipeline.py +++ b/tests/test_pipeline/test_pipeline.py @@ -297,7 +297,7 @@ 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. + time stays ``NaN`` rather than falling back to an earlier landmark. """ nwb_file = MagicMock() @@ -307,7 +307,8 @@ def test_add_trials_derives_native_start_and_stop_from_periods(): (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)] + assert extents[0] == (0.0, 1.0) + assert extents[1][0] == 1.0 and np.isnan(extents[1][1]) def test_add_trials_skips_frame_without_period_columns(): From b4f64f41606ea26531bf2709f0275f613585e3da Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 20 Aug 2026 10:24:47 -0700 Subject: [PATCH 3/5] refactor: use session end time for last trial for ITI_stop_time --- .../pipeline/_pipeline.py | 13 +++--- .../processing/_trial_table.py | 45 +++++++++++++++++-- .../processing/models/trial_config.py | 2 +- 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/dynamic_foraging_processing/pipeline/_pipeline.py b/src/dynamic_foraging_processing/pipeline/_pipeline.py index 00390f5..6c6f013 100644 --- a/src/dynamic_foraging_processing/pipeline/_pipeline.py +++ b/src/dynamic_foraging_processing/pipeline/_pipeline.py @@ -60,9 +60,10 @@ _NWB_START_COLUMN = "quiescent_start_time" #: Trials-table column NWB's required native ``stop_time`` is taken from: the -#: end of the ITI. This is ``NaN`` on the last trial of the session (whose ITI -#: end is unknown) and the ``NaN`` is propagated rather than substituted, so an -#: unknown trial end reads as unknown instead of as a shortened trial. +#: end of the ITI, which on the last trial of the session is the ``EndSession`` +#: timestamp. It is ``NaN`` only when that stream is unavailable, and the ``NaN`` +#: is propagated rather than substituted, so an unknown trial end reads as +#: unknown instead of as a shortened trial. _NWB_STOP_COLUMN = "ITI_stop_time" #: Source repository recorded in the ``processing.json`` data process. @@ -238,8 +239,10 @@ def _trial_extent(row: pd.Series) -> t.Tuple[float, float]: 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. The last trial of the session has - no ITI end, so its stop time is ``NaN`` — an unknown end is reported as - unknown rather than substituted with an earlier landmark. + no following quiescent period, so its ITI — and therefore its stop time — + ends at the ``EndSession`` timestamp. Should that be unavailable the stop + time is ``NaN``: an unknown end is reported as unknown rather than + substituted with an earlier landmark. Parameters ---------- diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index 52974d2..cf0070b 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -131,6 +131,28 @@ def _event_times(df: t.Optional[pd.DataFrame]) -> np.ndarray: return np.empty(0) return df.sort_index().index.to_numpy(dtype=float) + @classmethod + def _session_end_time(cls, end_session: t.Optional[pd.DataFrame]) -> float: + """Return the session's end timestamp from the ``EndSession`` stream. + + The stream carries a single event marking the end of the session; its + timestamp closes the last trial's ITI, which has no following + ``QuiescentPeriod`` event to end it. + + Parameters + ---------- + end_session : pandas.DataFrame or None + The ``EndSession`` software-event stream's data. + + Returns + ------- + float + The end-of-session timestamp, or ``NaN`` when the stream is absent + or empty. The last event is used if more than one is present. + """ + times = cls._event_times(end_session) + return float(times[-1]) if times.size else np.nan + @staticmethod def _time_at(times: np.ndarray, index: int) -> float: """Return ``times[index]``, or ``NaN`` when the stream is that much shorter. @@ -832,6 +854,7 @@ def _trial_periods( response_period_times: np.ndarray, consumption_times: np.ndarray, iti_times: np.ndarray, + session_end_time: float = np.nan, ) -> t.Dict[str, float]: """Return the start and stop time of every task period for one trial. @@ -839,7 +862,7 @@ def _trial_periods( 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``. + event, so it is closed by ``session_end_time`` instead. Parameters ---------- @@ -848,6 +871,10 @@ def _trial_periods( quiescent_times, response_period_times, consumption_times, iti_times : numpy.ndarray Per-trial timestamps of the ``QuiescentPeriod``, ``ResponsePeriod``, ``RewardConsumptionPeriod``, and ``ItiPeriod`` streams. + session_end_time : float, optional + Fallback ``ITI_stop_time`` for a trial with no following + ``QuiescentPeriod`` event — the ``EndSession`` timestamp, passed only + for the last trial. Defaults to ``NaN``, leaving the stop time unset. Returns ------- @@ -858,6 +885,9 @@ def _trial_periods( 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) + iti_stop = cls._time_at(quiescent_times, index + 1) + if np.isnan(iti_stop): + iti_stop = session_end_time return { "quiescent_start_time": cls._time_at(quiescent_times, index), "quiescent_stop_time": response_start, @@ -866,7 +896,7 @@ def _trial_periods( "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), + "ITI_stop_time": iti_stop, } def _build_row( @@ -977,7 +1007,9 @@ def build(self) -> pd.DataFrame: 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``). + times come from the next event in sequence (see ``_trial_periods``). The + last trial's ITI has no following event, so it is closed by the + ``EndSession`` timestamp. Hardware streams are handled per their nature: the go cue is an event each trial selects within its ``[quiescent_start_time, ITI_start_time)`` @@ -1003,6 +1035,7 @@ def build(self) -> pd.DataFrame: iti = self._load("Behavior", "SoftwareEvents", "ItiPeriod") responses = self._load("Behavior", "SoftwareEvents", "Response") metrics = self._load("Behavior", "SoftwareEvents", "TrialMetrics") + end_session = self._load("Behavior", "SoftwareEvents", "EndSession") pulse_supply_left = self._load("Behavior", "HarpBehavior", "PulseSupplyPort0") pulse_supply_right = self._load("Behavior", "HarpBehavior", "PulseSupplyPort1") @@ -1020,6 +1053,9 @@ def build(self) -> pd.DataFrame: response_payloads = self._event_payloads(responses) metric_payloads = self._event_payloads(metrics) + # Closes the last trial's ITI, which has no following QuiescentPeriod. + session_end_time = self._session_end_time(end_session) + # Guard the positional alignment before we pair streams by index. n_trials = len(outcome_payloads) @@ -1072,6 +1108,9 @@ def build(self) -> pd.DataFrame: response_period_times=response_period_times, consumption_times=consumption_times, iti_times=iti_times, + # Only the last trial's ITI is closed by the session end; an + # earlier gap means a short stream, which _check_aligned reports. + session_end_time=session_end_time if i == n_trials - 1 else np.nan, ) 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) diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index 94b0f47..fbfb27b 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -58,7 +58,7 @@ class TrialConfig(BaseModel): ) 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." + "End time of the inter-trial interval, i.e. the start of the next trial's quiescent period (the following QuiescentPeriod timestamp). The last trial of the session has no following quiescent period, so it ends at the EndSession timestamp; NaN if that stream is unavailable." ), ) From 81a1b0c53cd36d0fcc2dfd35088a4e2de6a2f227 Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 20 Aug 2026 10:25:00 -0700 Subject: [PATCH 4/5] test: update tests --- tests/test_pipeline/test_pipeline.py | 8 ++-- tests/test_processing/test_trial_table.py | 47 ++++++++++++++++++++++- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py index d7650d3..9e777d4 100644 --- a/tests/test_pipeline/test_pipeline.py +++ b/tests/test_pipeline/test_pipeline.py @@ -37,8 +37,8 @@ def _make_pipeline() -> Pipeline: def _trials_frame() -> pd.DataFrame: """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. + The second trial's ``ITI_stop_time`` is ``NaN``, as it is on the last trial + of a session whose ``EndSession`` stream was unavailable. """ return pd.DataFrame( { @@ -296,8 +296,8 @@ def test_add_trials_populates_columns_and_rows(): 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 stays ``NaN`` rather than falling back to an earlier landmark. + An unknown ITI end (no following quiescent period and no ``EndSession`` + timestamp) stays ``NaN`` rather than falling back to an earlier landmark. """ nwb_file = MagicMock() diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index f66c0da..5389034 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -261,6 +261,8 @@ def _full_dataset(): ) ), "TrialMetrics": _Stream(_events([10.2, 20.2], [{"bias": 0.3}, {"bias": None}])), + # Closes the last trial's ITI, which has no following QuiescentPeriod. + "EndSession": _Stream(_events([30.0], [None])), } ) behavior = _Node( @@ -323,13 +325,14 @@ def test_build_full_dataset(): first, second = table.iloc[0], table.iloc[1] # 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). + # ends at the next trial's quiescent period — or, on the last trial, at the + # EndSession timestamp. 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"]) + assert second["ITI_start_time"] == 25.0 and second["ITI_stop_time"] == 30.0 # delay_start_time is the legacy name for the quiescent period start. assert first["delay_start_time"] == first["quiescent_start_time"] == 10.0 @@ -479,6 +482,46 @@ def test_build_raises_on_misaligned_streams_when_configured(): TrialTableBuilder(_misaligned_dataset(), raise_on_error=True).build() +# --------------------------------------------------------------------------- # +# ITI_stop_time on the last trial — the EndSession timestamp +# --------------------------------------------------------------------------- # +def test_build_last_iti_stop_is_nan_without_end_session(): + """Without an ``EndSession`` stream the last trial's ITI end stays unknown.""" + dataset = _full_dataset() + del dataset.children["Behavior"].children["SoftwareEvents"].children["EndSession"] + + table = TrialTableBuilder(dataset).build() + + # Earlier trials are unaffected; only the last one lacks a closing event. + assert table.iloc[0]["ITI_stop_time"] == 20.0 + assert np.isnan(table.iloc[1]["ITI_stop_time"]) + + +def test_build_last_iti_stop_is_nan_when_end_session_is_empty(): + """An ``EndSession`` stream carrying no events leaves the last ITI end unknown.""" + dataset = _full_dataset() + dataset.children["Behavior"].children["SoftwareEvents"].children["EndSession"] = _Stream( + _events([], []) + ) + + table = TrialTableBuilder(dataset).build() + + assert np.isnan(table.iloc[1]["ITI_stop_time"]) + + +def test_build_end_session_does_not_close_a_mid_session_gap(): + """Only the last trial falls back to ``EndSession``. + + A short ``QuiescentPeriod`` stream leaves an earlier trial without a closing + event too, but attributing the session end to it would invent a trial + spanning the rest of the session, so it stays ``NaN``. + """ + table = TrialTableBuilder(_misaligned_dataset()).build() + + assert np.isnan(table.iloc[0]["ITI_stop_time"]) + assert table.iloc[1]["ITI_stop_time"] == 30.0 + + # --------------------------------------------------------------------------- # # _summary_generator — composite trial generators # --------------------------------------------------------------------------- # From d26fb89f0474f4141b3cd84ecd08a0470ba4aa7a Mon Sep 17 00:00:00 2001 From: arjunsridhar12345 Date: Thu, 20 Aug 2026 10:25:14 -0700 Subject: [PATCH 5/5] docs: update docs with session end time documentation --- docs/trials_table_mapping.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index f1725e7..a672528 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -150,14 +150,14 @@ durations track the configured ones (reward consumption ≈ | `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. | +| `ITI_stop_time` | The **next** trial's `QuiescentPeriod` `timestamp`. The last trial has no following quiescent period, so it takes the `EndSession` `timestamp`; `NaN` if that stream is unavailable. | | `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. +`stop_time` is `ITI_stop_time` — which on the last trial is the `EndSession` +timestamp. > **`delay` means `quiescent`.** The legacy `delay_*` columns describe the > acquisition software's *quiescence period* — the lick-free interval preceding