fix: last trial has iti stop time as nan but stop time has a value

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

src/dynamic_foraging_processing/pipeline/_pipeline.py

@@ -59,10 +59,12 @@
 #: ``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, 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.
 _CODE_URL = "https://github.com/AllenNeuralDynamics/dynamic-foraging-processing"
@@ -236,8 +238,11 @@ 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 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
         ----------
@@ -247,11 +252,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 +269,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()

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,14 +854,15 @@ 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.
 
         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``.
+        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)

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."
         ),
     )
 

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 falls back to the ITI start.
+    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()
 
@@ -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():

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
 # --------------------------------------------------------------------------- #