1761696499260742

Committed 05 Oct 2023 09:46AM UTC coverage: 55.27% (-1.4%) from 56.628%

Build # 1761696499260742

Build Type

Pull #655

continuous-integration/UCL

Committed by

bimac

Commit Message

add @sleepless decorator

Pull Request Pull Request #655: add @sleepless decorator

Run Details

21 of 21 new or added lines in 1 file covered. (100.0%)

10330 of 18690 relevant lines covered (55.27%)

0.55 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

81.03

/ibllib/qc/task_extractors.py

import logging

import numpy as np
from scipy.interpolate import interp1d

from ibllib.io.extractors import bpod_trials
from ibllib.io.extractors.base import get_session_extractor_type
from ibllib.io.extractors.training_wheel import get_wheel_position
from ibllib.io.extractors import ephys_fpga
import ibllib.io.raw_data_loaders as raw
from one.alf.spec import is_session_path
import one.alf.io as alfio
from one.api import ONE


_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',
                   'wheel_moves_intervals', 'wheel_moves_peak_amplitude',
                   'wheel_position', 'wheel_timestamps']


class TaskQCExtractor(object):
    def __init__(self, session_path, lazy=False, one=None, download_data=False, bpod_only=False,
                 sync_collection=None, sync_type=None, task_collection=None):
        """
        A class for extracting the task data required to perform task quality control
        :param session_path: a valid session path
        :param lazy: if True, the data are not extracted immediately
        :param one: an instance of ONE, used to download the raw data if download_data is True
        :param download_data: if True, any missing raw data is downloaded via ONE
        :param bpod_only: extract from from raw Bpod data only, even for FPGA sessions
        """
        if not is_session_path(session_path):
            raise ValueError('Invalid session path')
        self.session_path = session_path
        self.one = one
        self.log = _logger

        self.data = None
        self.settings = None
        self.raw_data = None
        self.frame_ttls = self.audio_ttls = self.bpod_ttls = None
        self.type = None
        self.wheel_encoding = None
        self.bpod_only = bpod_only
        self.sync_collection = sync_collection or 'raw_ephys_data'
        self.sync_type = sync_type
        self.task_collection = task_collection or 'raw_behavior_data'

        if download_data:
            self.one = one or ONE()
            self._ensure_required_data()

        if not lazy:
            self.load_raw_data()
            self.extract_data()

    def _ensure_required_data(self):
        """
        Attempt to download any required raw data if missing, and raise exception if any data are
        missing.
        :return:
        """
        dstypes = [
            '_iblrig_taskData.raw',
            '_iblrig_taskSettings.raw',
            '_iblrig_encoderPositions.raw',
            '_iblrig_encoderEvents.raw',
            '_iblrig_stimPositionScreen.raw',
            '_iblrig_syncSquareUpdate.raw',
            '_iblrig_encoderTrialInfo.raw',
            '_iblrig_ambientSensorData.raw',
        ]
        eid = self.one.path2eid(self.session_path)
        self.log.info(f'Downloading data for session {eid}')
        # Ensure we have the settings
        settings, _ = self.one.load_datasets(eid, ['_iblrig_taskSettings.raw.json'],
                                             collections=[self.task_collection],
                                             download_only=True, assert_present=False)

        is_ephys = get_session_extractor_type(self.session_path, task_collection=self.task_collection) == 'ephys'
        self.sync_type = self.sync_type or 'nidq' if is_ephys else 'bpod'
        is_fpga = 'bpod' not in self.sync_type

        if settings and is_ephys:

            dstypes.extend(['_spikeglx_sync.channels',
                            '_spikeglx_sync.polarities',
                            '_spikeglx_sync.times',
                            'ephysData.raw.meta',
                            'ephysData.raw.wiring'])
        elif settings and is_fpga:

            dstypes.extend(['_spikeglx_sync.channels',
                            '_spikeglx_sync.polarities',
                            '_spikeglx_sync.times',
                            'DAQData.raw.meta',
                            'DAQData.wiring'])

        dataset = self.one.type2datasets(eid, dstypes, details=True)
        files = self.one._check_filesystem(dataset)

        missing = [True] * len(dstypes) if not files else [x is None for x in files]
        if self.session_path is None or all(missing):
            self.lazy = True
            self.log.error('Data not found on server, can\'t calculate QC.')
        elif any(missing):
            self.log.warning(
                f'Missing some datasets for session {eid} in path {self.session_path}'
            )

    def load_raw_data(self):
        """
        Loads the TTLs, raw task data and task settings
        :return:
        """
        self.log.info(f'Loading raw data from {self.session_path}')
        self.type = self.type or get_session_extractor_type(self.session_path, task_collection=self.task_collection)
        # Finds the sync type when it isn't explicitly set, if ephys we assume nidq otherwise bpod
        self.sync_type = self.sync_type or 'nidq' if self.type == 'ephys' else 'bpod'

        self.settings, self.raw_data = raw.load_bpod(self.session_path, task_collection=self.task_collection)
        # Fetch the TTLs for the photodiode and audio
        if self.sync_type == 'bpod' or self.bpod_only is True:  # Extract from Bpod
            self.frame_ttls, self.audio_ttls = raw.load_bpod_fronts(
                self.session_path, data=self.raw_data, task_collection=self.task_collection)
        else:  # Extract from FPGA
            sync, chmap = ephys_fpga.get_sync_and_chn_map(self.session_path, self.sync_collection)

            def channel_events(name):
                """Fetches the polarities and times for a given channel"""
                keys = ('polarities', 'times')
                mask = sync['channels'] == chmap[name]
                return dict(zip(keys, (sync[k][mask] for k in keys)))

            ttls = [ephys_fpga._clean_frame2ttl(channel_events('frame2ttl')),
                    ephys_fpga._clean_audio(channel_events('audio')),
                    channel_events('bpod')]
            self.frame_ttls, self.audio_ttls, self.bpod_ttls = ttls

    def extract_data(self):
        """Extracts and loads behaviour data for QC
        NB: partial extraction when bpod_only attribute is False requires intervals and
        intervals_bpod to be assigned to the data attribute before calling this function.
        :return:
        """
        self.log.info(f'Extracting session: {self.session_path}')
        self.type = self.type or get_session_extractor_type(self.session_path, task_collection=self.task_collection)
        # Finds the sync type when it isn't explicitly set, if ephys we assume nidq otherwise bpod
        self.sync_type = self.sync_type or 'nidq' if self.type == 'ephys' else 'bpod'

        self.wheel_encoding = 'X4' if (self.sync_type != 'bpod' and not self.bpod_only) else 'X1'

        if not self.raw_data:
            self.load_raw_data()
        # Run extractors
        if self.sync_type != 'bpod' and not self.bpod_only:
            data, _ = ephys_fpga.extract_all(self.session_path, save=False, task_collection=self.task_collection)
            bpod2fpga = interp1d(data['intervals_bpod'][:, 0], data['table']['intervals_0'],
                                 fill_value='extrapolate')
            # Add Bpod wheel data
            re_ts, pos = get_wheel_position(self.session_path, self.raw_data, task_collection=self.task_collection)
            data['wheel_timestamps_bpod'] = bpod2fpga(re_ts)
            data['wheel_position_bpod'] = pos
        else:
            kwargs = dict(save=False, bpod_trials=self.raw_data, settings=self.settings, task_collection=self.task_collection)
            trials, wheel, _ = bpod_trials.extract_all(self.session_path, **kwargs)
            n_trials = np.unique(list(map(lambda k: trials[k].shape[0], trials)))[0]
            if self.type == 'habituation':
                data = trials
                data['position'] = np.array([t['position'] for t in self.raw_data])
                data['phase'] = np.array([t['stim_phase'] for t in self.raw_data])
                # Nasty hack to trim last trial due to stim off events happening at trial num + 1
                data = {k: v[:n_trials] for k, v in data.items()}
            else:
                data = {**trials, **wheel}
        # Update the data attribute with extracted data
        self.data = self.rename_data(data)

    @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
        if 'wheel_moves_intervals' not in data and 'wheelMoves_intervals' in data:
            data['wheel_moves_intervals'] = data.pop('wheelMoves_intervals')
        if 'wheel_moves_peak_amplitude' not in data and 'wheelMoves_peakAmplitude' in data:
            data['wheel_moves_peak_amplitude'] = data.pop('wheelMoves_peakAmplitude')
        data['correct'] = correct
        diff_fields = list(set(REQUIRED_FIELDS).difference(set(data.keys())))
        for miss_field in diff_fields:
            data[miss_field] = data['feedback_times'] * np.nan
        if len(diff_fields):
            _logger.warning(f'QC extractor, missing fields filled with NaNs: {diff_fields}')
        return data

1	import logging	1✔
2
3	import numpy as np	1✔
4	from scipy.interpolate import interp1d	1✔
5
6	from ibllib.io.extractors import bpod_trials	1✔
7	from ibllib.io.extractors.base import get_session_extractor_type	1✔
8	from ibllib.io.extractors.training_wheel import get_wheel_position	1✔
9	from ibllib.io.extractors import ephys_fpga	1✔
10	import ibllib.io.raw_data_loaders as raw	1✔
11	from one.alf.spec import is_session_path	1✔
12	import one.alf.io as alfio	1✔
13	from one.api import ONE	1✔
14
15
16	_logger = logging.getLogger('ibllib')	1✔
17
18	REQUIRED_FIELDS = ['choice', 'contrastLeft', 'contrastRight', 'correct',	1✔
19	'errorCueTrigger_times', 'errorCue_times', 'feedbackType', 'feedback_times',
20	'firstMovement_times', 'goCueTrigger_times', 'goCue_times', 'intervals',
21	'itiIn_times', 'phase', 'position', 'probabilityLeft', 'quiescence',
22	'response_times', 'rewardVolume', 'stimFreezeTrigger_times',
23	'stimFreeze_times', 'stimOffTrigger_times', 'stimOff_times',
24	'stimOnTrigger_times', 'stimOn_times', 'valveOpen_times',
25	'wheel_moves_intervals', 'wheel_moves_peak_amplitude',
26	'wheel_position', 'wheel_timestamps']
27
28
29	class TaskQCExtractor(object):	1✔
30	def __init__(self, session_path, lazy=False, one=None, download_data=False, bpod_only=False,	1✔
31	sync_collection=None, sync_type=None, task_collection=None):
32	"""
33	A class for extracting the task data required to perform task quality control
34	:param session_path: a valid session path
35	:param lazy: if True, the data are not extracted immediately
36	:param one: an instance of ONE, used to download the raw data if download_data is True
37	:param download_data: if True, any missing raw data is downloaded via ONE
38	:param bpod_only: extract from from raw Bpod data only, even for FPGA sessions
39	"""
40	if not is_session_path(session_path):	1✔
41	raise ValueError('Invalid session path')	×
42	self.session_path = session_path	1✔
43	self.one = one	1✔
44	self.log = _logger	1✔
45
46	self.data = None	1✔
47	self.settings = None	1✔
48	self.raw_data = None	1✔
49	self.frame_ttls = self.audio_ttls = self.bpod_ttls = None	1✔
50	self.type = None	1✔
51	self.wheel_encoding = None	1✔
52	self.bpod_only = bpod_only	1✔
53	self.sync_collection = sync_collection or 'raw_ephys_data'	1✔
54	self.sync_type = sync_type	1✔
55	self.task_collection = task_collection or 'raw_behavior_data'	1✔
56
57	if download_data:	1✔
58	self.one = one or ONE()	×
59	self._ensure_required_data()	×
60
61	if not lazy:	1✔
62	self.load_raw_data()	1✔
63	self.extract_data()	1✔
64
65	def _ensure_required_data(self):	1✔
66	"""
67	Attempt to download any required raw data if missing, and raise exception if any data are
68	missing.
69	:return:
70	"""
71	dstypes = [	×
72	'_iblrig_taskData.raw',
73	'_iblrig_taskSettings.raw',
74	'_iblrig_encoderPositions.raw',
75	'_iblrig_encoderEvents.raw',
76	'_iblrig_stimPositionScreen.raw',
77	'_iblrig_syncSquareUpdate.raw',
78	'_iblrig_encoderTrialInfo.raw',
79	'_iblrig_ambientSensorData.raw',
80	]
81	eid = self.one.path2eid(self.session_path)	×
82	self.log.info(f'Downloading data for session {eid}')	×
83	# Ensure we have the settings
84	settings, _ = self.one.load_datasets(eid, ['_iblrig_taskSettings.raw.json'],	×
85	collections=[self.task_collection],
86	download_only=True, assert_present=False)
87
88	is_ephys = get_session_extractor_type(self.session_path, task_collection=self.task_collection) == 'ephys'	×
89	self.sync_type = self.sync_type or 'nidq' if is_ephys else 'bpod'	×
90	is_fpga = 'bpod' not in self.sync_type	×
91
92	if settings and is_ephys:	×
93
94	dstypes.extend(['_spikeglx_sync.channels',	×
95	'_spikeglx_sync.polarities',
96	'_spikeglx_sync.times',
97	'ephysData.raw.meta',
98	'ephysData.raw.wiring'])
99	elif settings and is_fpga:	×
100
101	dstypes.extend(['_spikeglx_sync.channels',	×
102	'_spikeglx_sync.polarities',
103	'_spikeglx_sync.times',
104	'DAQData.raw.meta',
105	'DAQData.wiring'])
106
107	dataset = self.one.type2datasets(eid, dstypes, details=True)	×
108	files = self.one._check_filesystem(dataset)	×
109
110	missing = [True] * len(dstypes) if not files else [x is None for x in files]	×
111	if self.session_path is None or all(missing):	×
112	self.lazy = True	×
113	self.log.error('Data not found on server, can\'t calculate QC.')	×
114	elif any(missing):	×
115	self.log.warning(	×
116	f'Missing some datasets for session {eid} in path {self.session_path}'
117	)
118
119	def load_raw_data(self):	1✔
120	"""
121	Loads the TTLs, raw task data and task settings
122	:return:
123	"""
124	self.log.info(f'Loading raw data from {self.session_path}')	1✔
125	self.type = self.type or get_session_extractor_type(self.session_path, task_collection=self.task_collection)	1✔
126	# Finds the sync type when it isn't explicitly set, if ephys we assume nidq otherwise bpod
127	self.sync_type = self.sync_type or 'nidq' if self.type == 'ephys' else 'bpod'	1✔
128
129	self.settings, self.raw_data = raw.load_bpod(self.session_path, task_collection=self.task_collection)	1✔
130	# Fetch the TTLs for the photodiode and audio
131	if self.sync_type == 'bpod' or self.bpod_only is True: # Extract from Bpod	1✔
132	self.frame_ttls, self.audio_ttls = raw.load_bpod_fronts(	1✔
133	self.session_path, data=self.raw_data, task_collection=self.task_collection)
134	else: # Extract from FPGA
135	sync, chmap = ephys_fpga.get_sync_and_chn_map(self.session_path, self.sync_collection)	1✔
136
137	def channel_events(name):	1✔
138	"""Fetches the polarities and times for a given channel"""
139	keys = ('polarities', 'times')	1✔
140	mask = sync['channels'] == chmap[name]	1✔
141	return dict(zip(keys, (sync[k][mask] for k in keys)))	1✔
142
143	ttls = [ephys_fpga._clean_frame2ttl(channel_events('frame2ttl')),	1✔
144	ephys_fpga._clean_audio(channel_events('audio')),
145	channel_events('bpod')]
146	self.frame_ttls, self.audio_ttls, self.bpod_ttls = ttls	1✔
147
148	def extract_data(self):	1✔
149	"""Extracts and loads behaviour data for QC
150	NB: partial extraction when bpod_only attribute is False requires intervals and
151	intervals_bpod to be assigned to the data attribute before calling this function.
152	:return:
153	"""
154	self.log.info(f'Extracting session: {self.session_path}')	1✔
155	self.type = self.type or get_session_extractor_type(self.session_path, task_collection=self.task_collection)	1✔
156	# Finds the sync type when it isn't explicitly set, if ephys we assume nidq otherwise bpod
157	self.sync_type = self.sync_type or 'nidq' if self.type == 'ephys' else 'bpod'	1✔
158
159	self.wheel_encoding = 'X4' if (self.sync_type != 'bpod' and not self.bpod_only) else 'X1'	1✔
160
161	if not self.raw_data:	1✔
162	self.load_raw_data()	1✔
163	# Run extractors
164	if self.sync_type != 'bpod' and not self.bpod_only:	1✔
165	data, _ = ephys_fpga.extract_all(self.session_path, save=False, task_collection=self.task_collection)	1✔
166	bpod2fpga = interp1d(data['intervals_bpod'][:, 0], data['table']['intervals_0'],	1✔
167	fill_value='extrapolate')
168	# Add Bpod wheel data
169	re_ts, pos = get_wheel_position(self.session_path, self.raw_data, task_collection=self.task_collection)	1✔
170	data['wheel_timestamps_bpod'] = bpod2fpga(re_ts)	1✔
171	data['wheel_position_bpod'] = pos	1✔
172	else:
173	kwargs = dict(save=False, bpod_trials=self.raw_data, settings=self.settings, task_collection=self.task_collection)	1✔
174	trials, wheel, _ = bpod_trials.extract_all(self.session_path, **kwargs)	1✔
175	n_trials = np.unique(list(map(lambda k: trials[k].shape[0], trials)))[0]	1✔
176	if self.type == 'habituation':	1✔
177	data = trials	1✔
178	data['position'] = np.array([t['position'] for t in self.raw_data])	1✔
179	data['phase'] = np.array([t['stim_phase'] for t in self.raw_data])	1✔
180	# Nasty hack to trim last trial due to stim off events happening at trial num + 1
181	data = {k: v[:n_trials] for k, v in data.items()}	1✔
182	else:
183	data = {trials, wheel}	1✔
184	# Update the data attribute with extracted data
185	self.data = self.rename_data(data)	1✔
186
187	@staticmethod	1✔
188	def rename_data(data):	1✔
189	"""Rename the extracted data dict for use with TaskQC
190	Splits 'feedback_times' to 'errorCue_times' and 'valveOpen_times'.
191	NB: The data is not copied before making changes
192	:param data: A dict of task data returned by the task extractors
193	:return: the same dict after modifying the keys
194	"""
195	# Expand trials dataframe into key value pairs
196	trials_table = data.pop('table', None)	1✔
197	if trials_table is not None:	1✔
198	data = {data, alfio.AlfBunch.from_df(trials_table)}	1✔
199	correct = data['feedbackType'] > 0	1✔
200	# get valve_time and errorCue_times from feedback_times
201	if 'errorCue_times' not in data:	1✔
202	data['errorCue_times'] = data['feedback_times'].copy()	1✔
203	data['errorCue_times'][correct] = np.nan	1✔
204	if 'valveOpen_times' not in data:	1✔
205	data['valveOpen_times'] = data['feedback_times'].copy()	1✔
206	data['valveOpen_times'][~correct] = np.nan	1✔
207	if 'wheel_moves_intervals' not in data and 'wheelMoves_intervals' in data:	1✔
208	data['wheel_moves_intervals'] = data.pop('wheelMoves_intervals')	1✔
209	if 'wheel_moves_peak_amplitude' not in data and 'wheelMoves_peakAmplitude' in data:	1✔
210	data['wheel_moves_peak_amplitude'] = data.pop('wheelMoves_peakAmplitude')	1✔
211	data['correct'] = correct	1✔
212	diff_fields = list(set(REQUIRED_FIELDS).difference(set(data.keys())))	1✔
213	for miss_field in diff_fields:	1✔
214	data[miss_field] = data['feedback_times'] * np.nan	1✔
215	if len(diff_fields):	1✔
216	_logger.warning(f'QC extractor, missing fields filled with NaNs: {diff_fields}')	1✔
217	return data	1✔

int-brain-lab / ibllib / 1761696499260742

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous