Source code for ibllib.io.extractors.biased_trials

from pathlib import Path, PureWindowsPath

from packaging import version
import numpy as np
from one.alf.io import AlfBunch

from ibllib.io.extractors.base import BaseBpodTrialsExtractor, run_extractor_classes
import ibllib.io.raw_data_loaders as raw
from ibllib.io.extractors.training_trials import (
    Choice,
    FeedbackTimes,
    FeedbackType,
    GoCueTimes,
    GoCueTriggerTimes,
    IncludedTrials,
    Intervals,
    ProbabilityLeft,
    ResponseTimes,
    RewardVolume,
    StimOnTriggerTimes,
    StimOnOffFreezeTimes,
    ItiInTimes,
    StimOffTriggerTimes,
    StimFreezeTriggerTimes,
    ErrorCueTriggerTimes,
    PhasePosQuiescence,
)
from ibllib.io.extractors.training_wheel import Wheel

__all__ = ['BiasedTrials', 'EphysTrials']


class ContrastLR(BaseBpodTrialsExtractor):
    """Get left and right contrasts from raw datafile."""

    save_names = ('_ibl_trials.contrastLeft.npy', '_ibl_trials.contrastRight.npy')
    var_names = ('contrastLeft', 'contrastRight')

    def _extract(self, **kwargs):
        contrastLeft = np.array([t['contrast'] if np.sign(t['position']) < 0 else np.nan for t in self.bpod_trials])
        contrastRight = np.array([t['contrast'] if np.sign(t['position']) > 0 else np.nan for t in self.bpod_trials])
        return contrastLeft, contrastRight


class ProbaContrasts(BaseBpodTrialsExtractor):
    """Bpod pre-generated values for probabilityLeft, contrastLR, phase, quiescence."""

    save_names = (
        '_ibl_trials.contrastLeft.npy',
        '_ibl_trials.contrastRight.npy',
        None,
        None,
        '_ibl_trials.probabilityLeft.npy',
        '_ibl_trials.quiescencePeriod.npy',
    )
    var_names = ('contrastLeft', 'contrastRight', 'phase', 'position', 'probabilityLeft', 'quiescence')

    def _extract(self, **kwargs):
        """Extracts positions, contrasts, quiescent delay, stimulus phase and probability left
        from pregenerated session files.  Used in ephysChoiceWorld extractions.
        Optional: saves alf contrastLR and probabilityLeft npy files"""
        pe = self.get_pregenerated_events(self.bpod_trials, self.settings)
        return [pe[k] for k in sorted(pe.keys())]

    @staticmethod
    def get_pregenerated_events(bpod_trials, settings):
        for k in ['PRELOADED_SESSION_NUM', 'PREGENERATED_SESSION_NUM', 'SESSION_TEMPLATE_ID']:
            num = settings.get(k, None)
            if num is not None:
                break
        if num is None:
            fn = settings.get('SESSION_LOADED_FILE_PATH', '')
            fn = PureWindowsPath(fn).name
            num = ''.join([d for d in fn if d.isdigit()])
            if num == '':
                raise ValueError("Can't extract left probability behaviour.")
        # Load the pregenerated file
        ntrials = len(bpod_trials)
        sessions_folder = Path(raw.__file__).parent.joinpath('extractors', 'ephys_sessions')
        fname = f'session_{num}_ephys_pcqs.npy'
        pcqsp = np.load(sessions_folder.joinpath(fname))
        pos = pcqsp[:, 0]
        con = pcqsp[:, 1]
        pos = pos[:ntrials]
        con = con[:ntrials]
        contrastRight = con.copy()
        contrastLeft = con.copy()
        contrastRight[pos < 0] = np.nan
        contrastLeft[pos > 0] = np.nan
        qui = pcqsp[:, 2]
        qui = qui[:ntrials]
        phase = pcqsp[:, 3]
        phase = phase[:ntrials]
        pLeft = pcqsp[:, 4]
        pLeft = pLeft[:ntrials]

        phase_path = sessions_folder.joinpath(f'session_{num}_stim_phase.npy')
        is_patched_version = version.parse(settings.get('IBLRIG_VERSION') or '0') > version.parse('6.4.0')
        if phase_path.exists() and is_patched_version:
            phase = np.load(phase_path)[:ntrials]

        return {
            'position': pos,
            'quiescence': qui,
            'phase': phase,
            'probabilityLeft': pLeft,
            'contrastRight': contrastRight,
            'contrastLeft': contrastLeft,
        }


class TrialsTableBiased(BaseBpodTrialsExtractor):
    """
    Extracts the following into a table from Bpod raw data:
        intervals, goCue_times, response_times, choice, stimOn_times, contrastLeft, contrastRight,
        feedback_times, feedbackType, rewardVolume, probabilityLeft, firstMovement_times
    Additionally extracts the following wheel data:
        wheel_timestamps, wheel_position, wheelMoves_intervals, wheelMoves_peakAmplitude
    """

    save_names = (
        '_ibl_trials.table.pqt',
        None,
        None,
        '_ibl_wheel.timestamps.npy',
        '_ibl_wheel.position.npy',
        '_ibl_wheelMoves.intervals.npy',
        '_ibl_wheelMoves.peakAmplitude.npy',
        None,
        None,
    )
    var_names = (
        'table',
        'stimOff_times',
        'stimFreeze_times',
        'wheel_timestamps',
        'wheel_position',
        'wheelMoves_intervals',
        'wheelMoves_peakAmplitude',
        'wheelMoves_peakVelocity_times',
        'is_final_movement',
    )

    def _extract(self, extractor_classes=None, **kwargs):
        extractor_classes = extractor_classes or []
        base = [
            Intervals,
            GoCueTimes,
            ResponseTimes,
            Choice,
            StimOnOffFreezeTimes,
            ContrastLR,
            FeedbackTimes,
            FeedbackType,
            RewardVolume,
            ProbabilityLeft,
            Wheel,
        ]
        out, _ = run_extractor_classes(
            base + extractor_classes,
            session_path=self.session_path,
            bpod_trials=self.bpod_trials,
            settings=self.settings,
            save=False,
            task_collection=self.task_collection,
        )

        table = AlfBunch({k: out.pop(k) for k in list(out.keys()) if k not in self.var_names})
        assert len(table.keys()) == 12

        return table.to_df(), *(out.pop(x) for x in self.var_names if x != 'table')


class TrialsTableEphys(BaseBpodTrialsExtractor):
    """
    Extracts the following into a table from Bpod raw data:
        intervals, goCue_times, response_times, choice, stimOn_times, contrastLeft, contrastRight,
        feedback_times, feedbackType, rewardVolume, probabilityLeft, firstMovement_times
    Additionally extracts the following wheel data:
        wheel_timestamps, wheel_position, wheelMoves_intervals, wheelMoves_peakAmplitude
    """

    save_names = (
        '_ibl_trials.table.pqt',
        None,
        None,
        '_ibl_wheel.timestamps.npy',
        '_ibl_wheel.position.npy',
        '_ibl_wheelMoves.intervals.npy',
        '_ibl_wheelMoves.peakAmplitude.npy',
        None,
        None,
        None,
        None,
        '_ibl_trials.quiescencePeriod.npy',
    )
    var_names = (
        'table',
        'stimOff_times',
        'stimFreeze_times',
        'wheel_timestamps',
        'wheel_position',
        'wheelMoves_intervals',
        'wheelMoves_peakAmplitude',
        'wheelMoves_peakVelocity_times',
        'is_final_movement',
        'phase',
        'position',
        'quiescence',
    )

    def _extract(self, extractor_classes=None, **kwargs):
        extractor_classes = extractor_classes or []
        base = [
            Intervals,
            GoCueTimes,
            ResponseTimes,
            Choice,
            StimOnOffFreezeTimes,
            ProbaContrasts,
            FeedbackTimes,
            FeedbackType,
            RewardVolume,
            Wheel,
        ]
        # Exclude from trials table
        out, _ = run_extractor_classes(
            base + extractor_classes,
            session_path=self.session_path,
            bpod_trials=self.bpod_trials,
            settings=self.settings,
            save=False,
            task_collection=self.task_collection,
        )
        table = AlfBunch({k: v for k, v in out.items() if k not in self.var_names})
        assert len(table.keys()) == 12

        return table.to_df(), *(out.pop(x) for x in self.var_names if x != 'table')


[docs] class BiasedTrials(BaseBpodTrialsExtractor): """ Same as training_trials.TrainingTrials except... - there is no RepNum - ContrastLR is extracted differently - IncludedTrials is only extracted for 5.0.0 or greater """ save_names = ( '_ibl_trials.goCueTrigger_times.npy', '_ibl_trials.stimOnTrigger_times.npy', None, '_ibl_trials.stimOffTrigger_times.npy', None, None, '_ibl_trials.table.pqt', '_ibl_trials.stimOff_times.npy', None, '_ibl_wheel.timestamps.npy', '_ibl_wheel.position.npy', '_ibl_wheelMoves.intervals.npy', '_ibl_wheelMoves.peakAmplitude.npy', None, None, '_ibl_trials.included.npy', None, None, '_ibl_trials.quiescencePeriod.npy', ) var_names = ( 'goCueTrigger_times', 'stimOnTrigger_times', 'itiIn_times', 'stimOffTrigger_times', 'stimFreezeTrigger_times', 'errorCueTrigger_times', 'table', 'stimOff_times', 'stimFreeze_times', 'wheel_timestamps', 'wheel_position', 'wheelMoves_intervals', 'wheelMoves_peakAmplitude', 'wheelMoves_peakVelocity_times', 'is_final_movement', 'included', 'phase', 'position', 'quiescence', ) def _extract(self, extractor_classes=None, **kwargs) -> dict: extractor_classes = extractor_classes or [] base = [ GoCueTriggerTimes, StimOnTriggerTimes, ItiInTimes, StimOffTriggerTimes, StimFreezeTriggerTimes, ErrorCueTriggerTimes, TrialsTableBiased, IncludedTrials, PhasePosQuiescence, ] # Exclude from trials table out, _ = run_extractor_classes( base + extractor_classes, session_path=self.session_path, bpod_trials=self.bpod_trials, settings=self.settings, save=False, task_collection=self.task_collection, ) return {k: out[k] for k in self.var_names}
[docs] class EphysTrials(BaseBpodTrialsExtractor): """ Same as BiasedTrials except... - Contrast, phase, position, probabilityLeft and quiescence is extracted differently """ save_names = ( '_ibl_trials.goCueTrigger_times.npy', '_ibl_trials.stimOnTrigger_times.npy', None, '_ibl_trials.stimOffTrigger_times.npy', None, None, '_ibl_trials.table.pqt', '_ibl_trials.stimOff_times.npy', None, '_ibl_wheel.timestamps.npy', '_ibl_wheel.position.npy', '_ibl_wheelMoves.intervals.npy', '_ibl_wheelMoves.peakAmplitude.npy', None, None, '_ibl_trials.included.npy', None, None, '_ibl_trials.quiescencePeriod.npy', ) var_names = ( 'goCueTrigger_times', 'stimOnTrigger_times', 'itiIn_times', 'stimOffTrigger_times', 'stimFreezeTrigger_times', 'errorCueTrigger_times', 'table', 'stimOff_times', 'stimFreeze_times', 'wheel_timestamps', 'wheel_position', 'wheelMoves_intervals', 'wheelMoves_peakAmplitude', 'wheelMoves_peakVelocity_times', 'is_final_movement', 'included', 'phase', 'position', 'quiescence', ) def _extract(self, extractor_classes=None, **kwargs) -> dict: extractor_classes = extractor_classes or [] # For iblrig v8 we use the biased trials table instead. ContrastLeft, ContrastRight and ProbabilityLeft are # filled from the values in the bpod data itself rather than using the pregenerated session number iblrig_version = self.settings.get('IBLRIG_VERSION', self.settings.get('IBLRIG_VERSION_TAG', '0')) if version.parse(iblrig_version) >= version.parse('8.0.0'): TrialsTable = TrialsTableBiased else: TrialsTable = TrialsTableEphys base = [ GoCueTriggerTimes, StimOnTriggerTimes, ItiInTimes, StimOffTriggerTimes, StimFreezeTriggerTimes, ErrorCueTriggerTimes, TrialsTable, IncludedTrials, PhasePosQuiescence, ] # 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) # Exclude from trials table out, _ = run_extractor_classes( base + extractor_classes, session_path=self.session_path, bpod_trials=self.bpod_trials, settings=self.settings, save=False, task_collection=self.task_collection, ) return {k: out[k] for k in self.var_names}