Merge pull request #965 from int-brain-lab/TimelineTrialsHabituation

Timeline trials habituation

Author: k1o0

ibllib/io/extractors/ephys_fpga.py

@@ -35,7 +35,6 @@
 from pathlib import Path
 import uuid
 import re
-from functools import partial
 
 import matplotlib.pyplot as plt
 from matplotlib.colors import TABLEAU_COLORS
@@ -794,6 +793,7 @@ def _extract(self, sync=None, chmap=None, sync_collection='raw_ephys_data',
         elif (protocol_number := kwargs.get('protocol_number')) is not None:  # look for spacer
             # The spacers are TTLs generated by Bpod at the start of each protocol
             tmin, tmax = get_protocol_period(self.session_path, protocol_number, bpod)
+            tmin += (Spacer().times[-1] + Spacer().tup + 0.05)  # exclude spacer itself
         else:
             # Older sessions don't have protocol spacers so we sync the Bpod intervals here to
             # find the approximate end time of the protocol (this will exclude the passive signals
@@ -1453,16 +1453,23 @@ def get_bpod_event_times(self, sync, chmap, bpod_event_ttls=None, display=False,
         # lengths are defined by the state machine of the task protocol and therefore vary.
         if bpod_event_ttls is None:
             # Currently (at least v8.12 and below) there is no trial start or end TTL, only an ITI pulse
-            bpod_event_ttls = {'trial_iti': (1, 1.1), 'valve_open': (0, 0.4)}
+            bpod_event_ttls = {'trial_iti': (.999, 1.1), 'valve_open': (0, 0.4)}
         bpod_event_intervals = self._assign_events(
             bpod['times'], bpod['polarities'], bpod_event_ttls, display=display)
 
         # The first trial pulse is shorter and assigned to valve_open. Here we remove the first
         # valve event, prepend a 0 to the trial_start events, and drop the last trial if it was
         # incomplete in Bpod.
-        bpod_event_intervals['trial_iti'] = np.r_[bpod_event_intervals['valve_open'][0:1, :],
-                                                  bpod_event_intervals['trial_iti']]
-        bpod_event_intervals['valve_open'] = bpod_event_intervals['valve_open'][1:, :]
+        t0 = bpod_event_intervals['trial_iti'][0, 0]  # expect 1st event to be trial_start
+        pretrial = [(k, v[0, 0]) for k, v in bpod_event_intervals.items() if v.size and v[0, 0] < t0]
+        if pretrial:
+            (pretrial, _) = sorted(pretrial, key=lambda x: x[1])[0]  # take the earliest event
+            dt = np.diff(bpod_event_intervals[pretrial][0, :]) * 1e3  # record TTL length to log
+            _logger.debug('Reassigning first %s to trial_start. TTL length = %.3g ms', pretrial, dt)
+            bpod_event_intervals['trial_iti'] = np.r_[
+                bpod_event_intervals[pretrial][0:1, :], bpod_event_intervals['trial_iti']
+            ]
+            bpod_event_intervals[pretrial] = bpod_event_intervals[pretrial][1:, :]
 
         return bpod, bpod_event_intervals
 
@@ -1514,13 +1521,16 @@ def build_trials(self, sync, chmap, display=False, **kwargs):
         out.update({k: self.bpod2fpga(self.bpod_trials[k][ibpod]) for k in self.bpod_rsync_fields})
 
         # Assigning each event to a trial ensures exactly one event per trial (missing events are NaN)
-        assign_to_trial = partial(_assign_events_to_trial, fpga_events['intervals_0'])
         trials = alfio.AlfBunch({
-            'goCue_times': assign_to_trial(fpga_events['goCue_times'], take='first'),
-            'feedback_times': assign_to_trial(fpga_events['feedback_times']),
-            'stimCenter_times': assign_to_trial(self.frame2ttl['times'], take=-2),
-            'stimOn_times': assign_to_trial(self.frame2ttl['times'], take='first'),
-            'stimOff_times': assign_to_trial(self.frame2ttl['times']),
+            'goCue_times': _assign_events_to_trial(out['goCueTrigger_times'], fpga_events['goCue_times'], take='first'),
+            'feedback_times': _assign_events_to_trial(fpga_events['intervals_0'], fpga_events['feedback_times']),
+            'stimCenter_times': _assign_events_to_trial(
+                out['stimCenterTrigger_times'], self.frame2ttl['times'], take='first', t_trial_end=out['stimOffTrigger_times']),
+            'stimOn_times': _assign_events_to_trial(
+                out['stimOnTrigger_times'], self.frame2ttl['times'], take='first', t_trial_end=out['stimCenterTrigger_times']),
+            'stimOff_times': _assign_events_to_trial(
+                out['stimOffTrigger_times'], self.frame2ttl['times'],
+                take='first', t_trial_end=np.r_[out['intervals'][1:, 0], np.inf])
         })
         out.update({k: trials[k][ifpga] for k in trials.keys()})
 

ibllib/io/extractors/mesoscope.py

@@ -11,10 +11,11 @@
 from packaging import version
 
 from ibllib.plots.misc import squares, vertical_lines
-from ibllib.io.raw_daq_loaders import (extract_sync_timeline, timeline_get_channel,
-                                       correct_counter_discontinuities, load_timeline_sync_and_chmap)
+from ibllib.io.raw_daq_loaders import (
+    extract_sync_timeline, timeline_get_channel, correct_counter_discontinuities, load_timeline_sync_and_chmap)
 import ibllib.io.extractors.base as extractors_base
-from ibllib.io.extractors.ephys_fpga import FpgaTrials, WHEEL_TICKS, WHEEL_RADIUS_CM, _assign_events_to_trial
+from ibllib.io.extractors.ephys_fpga import (
+    FpgaTrials, FpgaTrialsHabituation, WHEEL_TICKS, WHEEL_RADIUS_CM, _assign_events_to_trial)
 from ibllib.io.extractors.training_wheel import extract_wheel_moves
 from ibllib.io.extractors.camera import attribute_times
 from brainbox.behavior.wheel import velocity_filtered
@@ -157,7 +158,7 @@ def load_sync(self, sync_collection='raw_sync_data', chmap=None, **_):
         return sync, chmap
 
     def _extract(self, sync=None, chmap=None, sync_collection='raw_sync_data', **kwargs) -> dict:
-        trials = super()._extract(sync, chmap, sync_collection='raw_sync_data', **kwargs)
+        trials = super()._extract(sync, chmap, sync_collection=sync_collection, **kwargs)
         if kwargs.get('display', False):
             plot_timeline(self.timeline, channels=chmap.keys(), raw=True)
         return trials
@@ -267,7 +268,8 @@ def assign_to_trial(events, take='last', starts=start_times, **kwargs):
         # Extract valve open times from the DAQ
         valve_driver_ttls = bpod_event_intervals['valve_open']
         correct = self.bpod_trials['feedbackType'] == 1
-        # If there is a reward_valve channel, the valve has
+        # If there is a reward_valve channel, the voltage across the valve has been recorded and
+        # should give a more accurate readout of the valve's activity.
         if any(ch['name'] == 'reward_valve' for ch in self.timeline['meta']['inputs']):
             # TODO Let's look at the expected open length based on calibration and reward volume
             # import scipy.interpolate
@@ -600,6 +602,63 @@ def _assign_events_audio(self, audio_times, audio_polarities, display=False):
         return t_ready_tone_in, t_error_tone_in
 
 
+class TimelineTrialsHabituation(FpgaTrialsHabituation, TimelineTrials):
+    """Habituation trials extraction from Timeline DAQ data."""
+
+    sync_field = 'intervals_0'
+
+    def build_trials(self, sync=None, chmap=None, **kwargs):
+        """
+        Extract task related event times from the sync.
+
+        The valve used at the mesoscope has a way to record the raw voltage across the solenoid,
+        giving a more accurate readout of the valve's activity. If the reward_valve channel is
+        present on the DAQ, this is used to extract the valve open times.
+
+        Parameters
+        ----------
+        sync : dict
+            'polarities' of fronts detected on sync trace for all 16 chans and their 'times'
+        chmap : dict
+            Map of channel names and their corresponding index.  Default to constant.
+
+        Returns
+        -------
+        dict
+            A map of trial event timestamps.
+        """
+        out = super().build_trials(sync, chmap, **kwargs)
+
+        start_times = out['intervals'][:, 0]
+        last_trial_end = out['intervals'][-1, 1]
+
+        # Extract valve open times from the DAQ
+        _, bpod_event_intervals = self.get_bpod_event_times(sync, chmap, **kwargs)
+        bpod_feedback_times = self.bpod2fpga(self.bpod_trials['feedback_times'])
+        valve_driver_ttls = bpod_event_intervals['valve_open']
+        # If there is a reward_valve channel, the voltage across the valve has been recorded and
+        # should give a more accurate readout of the valve's activity.
+        if any(ch['name'] == 'reward_valve' for ch in self.timeline['meta']['inputs']):
+            # Use the driver TTLs to find the valve open times that correspond to the valve opening
+            valve_intervals, valve_open_times = self.get_valve_open_times(driver_ttls=valve_driver_ttls)
+            if valve_open_times.size != start_times.size:
+                _logger.warning(
+                    'Number of valve open times does not equal number of correct trials (%i != %i)',
+                    valve_open_times.size, start_times.size)
+        else:
+            # Use the valve controller TTLs recorded on the Bpod channel as the reward time
+            valve_open_times = valve_driver_ttls[:, 0]
+        # there may be an extra last trial that's not in the Bpod intervals as the extractor ignores the last trial
+        valve_open_times = valve_open_times[valve_open_times <= last_trial_end]
+        out['valveOpen_times'] = _assign_events_to_trial(
+            bpod_feedback_times, valve_open_times, take='first', t_trial_end=out['intervals'][:, 1])
+
+        # Feedback times
+        out['feedback_times'] = np.copy(out['valveOpen_times'])
+
+        return out
+
+
 class MesoscopeSyncTimeline(extractors_base.BaseExtractor):
     """Extraction of mesoscope imaging times."""
 

ibllib/pipes/behavior_tasks.py

@@ -16,7 +16,7 @@
 from ibllib.io.extractors.ephys_passive import PassiveChoiceWorld
 from ibllib.io.extractors.bpod_trials import get_bpod_extractor
 from ibllib.io.extractors.ephys_fpga import FpgaTrials, FpgaTrialsHabituation, get_sync_and_chn_map
-from ibllib.io.extractors.mesoscope import TimelineTrials
+from ibllib.io.extractors.mesoscope import TimelineTrials, TimelineTrialsHabituation
 from ibllib.pipes import training_status
 from ibllib.plots.figures import BehaviourPlots
 
@@ -157,6 +157,40 @@ def run_qc(self, trials_data=None, update=True, **_):
         return qc
 
 
+class HabituationTrialsTimeline(HabituationTrialsNidq):
+    """Behaviour task extractor with DAQdata.raw NPY datasets."""
+    @property
+    def signature(self):
+        signature = super().signature
+        signature['input_files'] = [
+            ('_iblrig_taskData.raw.*', self.collection, True),
+            ('_iblrig_taskSettings.raw.*', self.collection, True),
+            (f'_{self.sync_namespace}_DAQdata.raw.npy', self.sync_collection, True),
+            (f'_{self.sync_namespace}_DAQdata.timestamps.npy', self.sync_collection, True),
+            (f'_{self.sync_namespace}_DAQdata.meta.json', self.sync_collection, True),
+        ]
+        return signature
+
+    def extract_behaviour(self, save=True, **kwargs):
+        """Extract the Bpod trials data and Timeline acquired signals."""
+        # First determine the extractor from the task protocol
+        bpod_trials, _ = HabituationTrialsBpod.extract_behaviour(self, save=False, **kwargs)
+
+        # Sync Bpod trials to DAQ
+        self.extractor = TimelineTrialsHabituation(self.session_path, bpod_trials=bpod_trials, bpod_extractor=self.extractor)
+        save_path = self.session_path / self.output_collection
+        if not self._spacer_support(self.extractor.settings):
+            _logger.warning('Protocol spacers not supported; setting protocol_number to None')
+            self.protocol_number = None
+
+        # NB: The stimOff times are called stimCenter times for habituation choice world
+        dsets, out_files = self.extractor.extract(
+            save=save, path_out=save_path, sync_collection=self.sync_collection,
+            task_collection=self.collection, protocol_number=self.protocol_number, **kwargs)
+
+        return dsets, out_files
+
+
 class TrialRegisterRaw(base_tasks.RegisterRawDataTask, base_tasks.BehaviourTask):
     priority = 100
     job_size = 'small'

ibllib/qc/task_metrics.py

@@ -458,14 +458,14 @@ def compute(self, **kwargs):
                 sessions = sessions[1:]
             metric = ([0, data['intervals'][-1, 1] - data['intervals'][0, 0]] +
                       [(datetime.fromisoformat(x['end_time']) -
-                        datetime.fromisoformat(x['start_time'])).total_seconds() / 60
+                        datetime.fromisoformat(x['start_time'])).total_seconds()
                        for x in [self.one.alyx.get(s['url']) for s in sessions]])
 
             # The duration from raw trial data
             # duration = map(float, self.extractor.raw_data[-1]['elapsed_time'].split(':'))
             # duration = timedelta(**dict(zip(('hours', 'minutes', 'seconds'),
-            #                                 duration))).total_seconds() / 60
-            metrics[check] = np.array(metric)
+            #                                 duration))).total_seconds()
+            metrics[check] = np.array(metric) / 60
             passed[check] = np.diff(metric) >= 12
 
         # Check event orders: trial_start < stim on < stim center < feedback < stim off
@@ -515,9 +515,19 @@ def compute(self, **kwargs):
         passed[check][-1] = np.nan
         metrics[check] = metric
 
+        # Check stim on go cue delay
+        # Unlike in ChoiceWorld, the go cue is in a race condition with the stimulus onset and
+        # therefore usually presents slightly before the stimulus. For now we will keep the QC
+        # thresholds the same as in ChoiceWorld which means this will always fail
+        check = prefix + 'stimOn_goCue_delays'
+        # If either are NaN, the result will be Inf to ensure that it crosses the failure threshold.
+        threshold = 0.01 if audio_output.lower() == 'harp' else 0.053
+        metric = np.nan_to_num(data['goCue_times'] - data['stimOn_times'], nan=np.inf)
+        passed[check] = np.abs(metric) < threshold
+        metrics[check] = metric
+
         # Checks common to training QC
-        checks = [check_goCue_delays, check_stimOn_goCue_delays,
-                  check_stimOn_delays, check_stimOff_delays]
+        checks = [check_goCue_delays, check_stimOn_delays, check_stimOff_delays]
         for fcn in checks:
             check = prefix + fcn.__name__[6:]
             metrics[check], passed[check] = fcn(data, audio_output=audio_output)

ibllib/qc/task_qc_viewer/task_qc.py

@@ -106,8 +106,8 @@ def n_trials(self):
         return self.qc.extractor.data['intervals'].shape[0]
 
     def get_wheel_data(self):
-        return {'re_pos': self.qc.extractor.data.get('wheel_position', np.array([])),
-                're_ts': self.qc.extractor.data.get('wheel_timestamps', np.array([]))}
+        return {'re_pos': self.qc.extractor.data.get('wheel_position') or np.array([]),
+                're_ts': self.qc.extractor.data.get('wheel_timestamps') or np.array([])}
 
     def create_plots(self, axes,
                      wheel_axes=None, trial_events=None, color_map=None, linestyle=None):
@@ -280,9 +280,11 @@ def show_session_task_qc(qc_or_session=None, bpod_only=False, local=False, one=N
         events = map(lambda x: x.replace('stimFreeze', 'stimCenter'), events)
 
     # Run QC and plot
-    w = ViewEphysQC.viewqc(wheel=qc.get_wheel_data())
+    wheel_data = qc.get_wheel_data()
+    w = ViewEphysQC.viewqc(wheel=wheel_data if wheel_data['re_pos'].size else None)
+
     qc.create_plots(w.wplot.canvas.ax,
-                    wheel_axes=w.wplot.canvas.ax2,
+                    wheel_axes=getattr(w.wplot.canvas, 'ax2', None),
                     trial_events=list(events),
                     color_map=cm,
                     linestyle=ls)
@@ -292,7 +294,7 @@ def show_session_task_qc(qc_or_session=None, bpod_only=False, local=False, one=N
     if 'task_qc' in locals():
         df_trials = pd.DataFrame({
             k: v for k, v in task_qc.extractor.data.items()
-            if v.size == n_trials and not k.startswith('wheel')
+            if not k.startswith('wheel') and v.size == n_trials
         })
         df = df_trials.merge(qc.frame, left_index=True, right_index=True)
     else: