Issue #666

ibllib/io/extractors/habituation_trials.py

@@ -1,12 +1,11 @@
+"""Habituation ChoiceWorld Bpod trials extraction."""
 import logging
 import numpy as np
 
 import ibllib.io.raw_data_loaders as raw
 from ibllib.io.extractors.base import BaseBpodTrialsExtractor, run_extractor_classes
 from ibllib.io.extractors.biased_trials import ContrastLR
-from ibllib.io.extractors.training_trials import (
-    FeedbackTimes, StimOnTriggerTimes, Intervals, GoCueTimes
-)
+from ibllib.io.extractors.training_trials import FeedbackTimes, StimOnTriggerTimes, GoCueTimes
 
 _logger = logging.getLogger(__name__)
 
@@ -24,9 +23,24 @@ def __init__(self, *args, **kwargs):
         self.save_names = tuple(f'_ibl_trials.{x}.npy' if x not in exclude else None for x in self.var_names)
 
     def _extract(self) -> dict:
+        """
+        Extract the Bpod trial events.
+
+        The Bpod state machine for this task has extremely misleading names! The 'iti' state is
+        actually the delay between valve open and trial end (the stimulus is still present during
+        this period), and the 'trial_start' state is actually the ITI during which there is a 1s
+        Bpod TTL and gray screen period.
+
+        Returns
+        -------
+        dict
+            A dictionary of Bpod trial events. The keys are defined in the `var_names` attribute.
+        """
         # Extract all trials...
 
-        # Get all stim_sync events detected
+        # Get all detected TTLs. These are stored for QC purposes
+        self.frame2ttl, self.audio = raw.load_bpod_fronts(self.session_path, data=self.bpod_trials)
+        # These are the frame2TTL pulses as a list of lists, one per trial
         ttls = [raw.get_port_events(tr, 'BNC1') for tr in self.bpod_trials]
 
         # Report missing events
@@ -38,10 +52,45 @@ def _extract(self) -> dict:
             _logger.warning(f'{self.session_path}: Missing BNC1 TTLs on {n_missing} trial(s)')
 
         # Extract datasets common to trainingChoiceWorld
-        training = [ContrastLR, FeedbackTimes, Intervals, GoCueTimes, StimOnTriggerTimes]
+        training = [ContrastLR, FeedbackTimes, GoCueTimes, StimOnTriggerTimes]
         out, _ = run_extractor_classes(training, session_path=self.session_path, save=False,
                                        bpod_trials=self.bpod_trials, settings=self.settings, task_collection=self.task_collection)
 
+        """
+        The 'trial_start' state is in fact the 1s grey screen period, therefore the first timestamp
+        is really the end of the previous trial and also the stimOff trigger time. The second
+        timestamp is the true trial start time.
+        """
+        (_, *ends), starts = zip(*[
+            t['behavior_data']['States timestamps']['trial_start'][-1] for t in self.bpod_trials]
+        )
+
+        # StimOffTrigger times
+        out['stimOffTrigger_times'] = np.array(ends)
+
+        # StimOff times
+        """
+        There should be exactly three TTLs per trial.  stimOff_times should be the first TTL pulse.
+        If 1 or more pulses are missing, we can not be confident of assigning the correct one.
+        """
+        out['stimOff_times'] = np.array([sync[0] if len(sync) == 3 else np.nan
+                                         for sync, off in zip(ttls[1:], ends)])
+
+        # Trial intervals
+        """
+        In terms of TTLs, the intervals are defined by the 'trial_start' state, however the stim
+        off time often happens after the trial end TTL front, i.e. after the 'trial_start' start
+        begins.  For these trials, we set the trial end time as the stim off time.
+        """
+        # NB: We lose the last trial because the stim off event occurs at trial_num + 1
+        n_trials = out['stimOff_times'].size
+        out['intervals'] = np.c_[starts, np.r_[ends, np.nan]][:n_trials, :]
+        to_update = out['intervals'][:, 1] < out['stimOff_times']
+        out['intervals'][to_update, 1] = out['stimOff_times'][to_update]
+
+        # itiIn times
+        out['itiIn_times'] = np.r_[ends, np.nan]
+
         # GoCueTriggerTimes is the same event as StimOnTriggerTimes
         out['goCueTrigger_times'] = out['stimOnTrigger_times'].copy()
 
@@ -75,38 +124,22 @@ def _extract(self) -> dict:
         trial_volume = [x['reward_amount'] for x in self.bpod_trials]
         out['rewardVolume'] = np.array(trial_volume).astype(np.float64)
 
-        # StimOffTrigger times
-        # StimOff occurs at trial start (ignore the first trial's state update)
-        out['stimOffTrigger_times'] = np.array(
-            [tr["behavior_data"]["States timestamps"]
-             ["trial_start"][0][0] for tr in self.bpod_trials[1:]]
-        )
-
-        # StimOff times
-        """
-        There should be exactly three TTLs per trial.  stimOff_times should be the first TTL pulse.
-        If 1 or more pulses are missing, we can not be confident of assigning the correct one.
-        """
-        trigg = out['stimOffTrigger_times']
-        out['stimOff_times'] = np.array([sync[0] if len(sync) == 3 else np.nan
-                                         for sync, off in zip(ttls[1:], trigg)])
-
         # FeedbackType is always positive
         out['feedbackType'] = np.ones(len(out['feedback_times']), dtype=np.int8)
 
-        # ItiIn times
-        out['itiIn_times'] = np.array(
-            [tr["behavior_data"]["States timestamps"]
-             ["iti"][0][0] for tr in self.bpod_trials]
-        )
-
         # Phase and position
         out['position'] = np.array([t['position'] for t in self.bpod_trials])
         out['phase'] = np.array([t['stim_phase'] for t in self.bpod_trials])
 
-        # NB: We lose the last trial because the stim off event occurs at trial_num + 1
-        n_trials = out['stimOff_times'].size
-        # return [out[k][:n_trials] for k in self.var_names]
+        # Double-check that the early and late trial events occur within the trial intervals
+        idx = ~np.isnan(out['stimOn_times'][:n_trials])
+        if np.any(out['stimOn_times'][:n_trials][idx] < out['intervals'][idx, 0]):
+            _logger.warning('Stim on events occurring outside trial intervals')
+        idx = ~np.isnan(out['stimOff_times'])
+        if np.any(out['stimOff_times'][idx] > out['intervals'][idx, 1]):
+            _logger.warning('Stim off events occurring outside trial intervals')
+
+        # Truncate arrays and return in correct order
         return {k: out[k][:n_trials] for k in self.var_names}