Source code for ibllib.qc.task_extractors

import logging

import numpy as np

from one.alf.spec import is_session_path
import one.alf.io as alfio


_logger = logging.getLogger('ibllib')

REQUIRED_FIELDS = ['choice', 'contrastLeft', 'contrastRight', 'correct',
                   'errorCueTrigger_times', 'errorCue_times', 'feedbackType', 'feedback_times',
                   'firstMovement_times', 'goCueTrigger_times', 'goCue_times', 'intervals',
                   'itiIn_times', 'phase', 'position', 'probabilityLeft', 'quiescence',
                   'response_times', 'rewardVolume', 'stimFreezeTrigger_times',
                   'stimFreeze_times', 'stimOffTrigger_times', 'stimOff_times',
                   'stimOnTrigger_times', 'stimOn_times', 'valveOpen_times',
                   'wheelMoves_intervals', 'wheelMoves_peakAmplitude', 'wheelMoves_peakVelocity_times',
                   'wheel_position', 'wheel_timestamps']




class TaskQCExtractor:
    def __init__(self, session_path):
        """
        A class for holding the task data required to perform task quality control.
        :param session_path: a valid session path
        """
        if not is_session_path(session_path):
            raise ValueError('Invalid session path')
        self.session_path = session_path
        self.log = _logger

        self.data = None
        self.settings = None
        self.raw_data = None
        self.frame_ttls = self.audio_ttls = self.bpod_ttls = None
        self.wheel_encoding = None



    @staticmethod
    def rename_data(data):
        """Rename the extracted data dict for use with TaskQC
        Splits 'feedback_times' to 'errorCue_times' and 'valveOpen_times'.
        NB: The data is not copied before making changes
        :param data: A dict of task data returned by the task extractors
        :return: the same dict after modifying the keys
        """
        # Expand trials dataframe into key value pairs
        trials_table = data.pop('table', None)
        if trials_table is not None:
            data = {**data, **alfio.AlfBunch.from_df(trials_table)}
        correct = data['feedbackType'] > 0
        # get valve_time and errorCue_times from feedback_times
        if 'errorCue_times' not in data:
            data['errorCue_times'] = data['feedback_times'].copy()
            data['errorCue_times'][correct] = np.nan
        if 'valveOpen_times' not in data:
            data['valveOpen_times'] = data['feedback_times'].copy()
            data['valveOpen_times'][~correct] = np.nan
        data['correct'] = correct
        diff_fields = list(set(REQUIRED_FIELDS).difference(set(data.keys())))
        for miss_field in diff_fields:
            data[miss_field] = None if miss_field.startswith('wheel') else data['feedback_times'] * np.nan
        if len(diff_fields):
            _logger.warning(f'QC extractor, missing fields filled with NaNs: {diff_fields}')
        return data