Merge pull request #724 from int-brain-lab/custom_task_qc

Custom task qc

Author: k1o0

ibllib/io/extractors/base.py

@@ -17,13 +17,16 @@
 
 class BaseExtractor(abc.ABC):
     """
-    Base extractor class
+    Base extractor class.
+
     Writing an extractor checklist:
-    -   on the child class, overload the _extract method
-    -   this method should output one or several numpy.arrays or dataframe with a consistent shape
-    -   save_names is a list or a string of filenames, there should be one per dataset
-    -   set save_names to None for a dataset that doesn't need saving (could be set dynamically
-    in the _extract method)
+
+    - on the child class, overload the _extract method
+    - this method should output one or several numpy.arrays or dataframe with a consistent shape
+    - save_names is a list or a string of filenames, there should be one per dataset
+    - set save_names to None for a dataset that doesn't need saving (could be set dynamically in
+      the _extract method)
+
     :param session_path: Absolute path of session folder
     :type session_path: str/Path
     """
@@ -122,10 +125,11 @@ def _extract(self):
 
 class BaseBpodTrialsExtractor(BaseExtractor):
     """
-    Base (abstract) extractor class for bpod jsonable data set
-    Wrps the _extract private method
+    Base (abstract) extractor class for bpod jsonable data set.
 
-    :param session_path: Absolute path of session folder
+    Wraps the _extract private method.
+
+    :param session_path: Absolute path of session folder.
     :type session_path: str
     :param bpod_trials
     :param settings
@@ -159,6 +163,12 @@ def extract(self, bpod_trials=None, settings=None, **kwargs):
             self.settings["IBLRIG_VERSION"] = "100.0.0"
         return super(BaseBpodTrialsExtractor, self).extract(**kwargs)
 
+    @property
+    def alf_path(self):
+        """pathlib.Path: The full task collection filepath."""
+        if self.session_path:
+            return self.session_path.joinpath(self.task_collection or '').absolute()
+
 
 def run_extractor_classes(classes, session_path=None, **kwargs):
     """

ibllib/io/extractors/training_trials.py

@@ -570,9 +570,9 @@ def get_stimOn_times_ge5(session_path, data=False, task_collection='raw_behavior
     @staticmethod
     def get_stimOn_times_lt5(session_path, data=False, task_collection='raw_behavior_data'):
         """
-        Find the time of the statemachine command to turn on hte stim
+        Find the time of the statemachine command to turn on the stim
         (state stim_on start or rotary_encoder_event2)
-        Find the next frame change from the photodiodeafter that TS.
+        Find the next frame change from the photodiode after that TS.
         Screen is not displaying anything until then.
         (Frame changes are in BNC1High and BNC1Low)
         """

ibllib/io/globus.py

@@ -90,16 +90,98 @@ def __init__(self, session_path, **kwargs):
             self.output_collection += f'/task_{self.protocol_number:02}'
 
     def get_protocol(self, protocol=None, task_collection=None):
-        return protocol if protocol else sess_params.get_task_protocol(self.session_params, task_collection)
+        """
+        Return the task protocol name.
+
+        This returns the task protocol based on the task collection. If `protocol` is not None, this
+        acts as an identity function. If both `task_collection` and `protocol` are None, returns
+        the protocol defined in the experiment description file only if a single protocol was run.
+        If the `task_collection` is not None, the associated protocol name is returned.
+
+
+        Parameters
+        ----------
+        protocol : str
+            A task protocol name. If not None, the same value is returned.
+        task_collection : str
+            The task collection whose protocol name to return. May be None if only one protocol run.
+
+        Returns
+        -------
+        str, None
+            The task protocol name, or None, if no protocol found.
+
+        Raises
+        ------
+        ValueError
+            For session with multiple task protocols, a task collection must be passed.
+        """
+        if protocol:
+            return protocol
+        protocol = sess_params.get_task_protocol(self.session_params, task_collection) or None
+        if isinstance(protocol, set):
+            if len(protocol) == 1:
+                protocol = next(iter(protocol))
+            else:
+                raise ValueError('Multiple task protocols for session. Task collection must be explicitly defined.')
+        return protocol
 
     def get_task_collection(self, collection=None):
+        """
+        Return the task collection.
+
+        If `collection` is not None, this acts as an identity function. Otherwise loads it from
+        the experiment description if only one protocol was run.
+
+        Parameters
+        ----------
+        collection : str
+            A task collection. If not None, the same value is returned.
+
+        Returns
+        -------
+        str, None
+            The task collection, or None if no task protocols were run.
+
+        Raises
+        ------
+        AssertionError
+            Raised if multiple protocols were run and collection is None, or if experiment
+            description file is improperly formatted.
+
+        """
         if not collection:
             collection = sess_params.get_task_collection(self.session_params)
         # If inferring the collection from the experiment description, assert only one returned
         assert collection is None or isinstance(collection, str) or len(collection) == 1
         return collection
 
     def get_protocol_number(self, number=None, task_protocol=None):
+        """
+        Return the task protocol number.
+
+        Numbering starts from 0. If the 'protocol_number' field is missing from the experiment
+        description, None is returned. If `task_protocol` is None, the first protocol number if n
+        protocols == 1, otherwise returns None.
+
+        NB: :func:`ibllib.pipes.dynamic_pipeline.make_pipeline` will determine the protocol number
+        from the order of the tasks in the experiment description if the task collection follows
+        the pattern 'raw_task_data_XX'. If the task protocol does not follow this pattern, the
+        experiment description file should explicitly define the number with the 'protocol_number'
+        field.
+
+        Parameters
+        ----------
+        number : int
+            The protocol number. If not None, the same value is returned.
+        task_protocol : str
+            The task protocol name.
+
+        Returns
+        -------
+        int, None
+            The task protocol number, if defined.
+        """
         if number is None:  # Do not use "if not number" as that will return True if number is 0
             number = sess_params.get_task_protocol_number(self.session_params, task_protocol)
         # If inferring the number from the experiment description, assert only one returned (or something went wrong)

ibllib/pipes/base_tasks.py

@@ -277,6 +277,7 @@ class ChoiceWorldTrialsBpod(base_tasks.BehaviourTask):
     priority = 90
     job_size = 'small'
     extractor = None
+    """ibllib.io.extractors.base.BaseBpodTrialsExtractor: An instance of the Bpod trials extractor."""
 
     @property
     def signature(self):
@@ -318,7 +319,24 @@ def _extract_behaviour(self, **kwargs):
         self.extractor.default_path = self.output_collection
         return self.extractor.extract(task_collection=self.collection, **kwargs)
 
-    def _run_qc(self, trials_data=None, update=True):
+    def _run_qc(self, trials_data=None, update=True, QC=None):
+        """
+        Run the task QC.
+
+        Parameters
+        ----------
+        trials_data : dict
+            The complete extracted task data.
+        update : bool
+            If True, updates the session QC fields on Alyx.
+        QC : ibllib.qc.task_metrics.TaskQC
+            An optional QC class to instantiate.
+
+        Returns
+        -------
+        ibllib.qc.task_metrics.TaskQC
+            The task QC object.
+        """
         if not self.extractor or trials_data is None:
             trials_data, _ = self._extract_behaviour(save=False)
         if not trials_data:
@@ -328,10 +346,11 @@ def _run_qc(self, trials_data=None, update=True):
         qc_extractor = TaskQCExtractor(self.session_path, lazy=True, sync_collection=self.sync_collection, one=self.one,
                                        sync_type=self.sync, task_collection=self.collection)
         qc_extractor.data = qc_extractor.rename_data(trials_data)
-        if type(self.extractor).__name__ == 'HabituationTrials':
-            qc = HabituationQC(self.session_path, one=self.one, log=_logger)
-        else:
-            qc = TaskQC(self.session_path, one=self.one, log=_logger)
+        if not QC:
+            QC = HabituationQC if type(self.extractor).__name__ == 'HabituationTrials' else TaskQC
+            _logger.debug('Running QC with %s.%s', QC.__module__, QC.__name__)
+        qc = QC(self.session_path, one=self.one, log=_logger)
+        if QC is not HabituationQC:
             qc_extractor.wheel_encoding = 'X1'
         qc_extractor.settings = self.extractor.settings
         qc_extractor.frame_ttls, qc_extractor.audio_ttls = load_bpod_fronts(
@@ -412,7 +431,7 @@ def _extract_behaviour(self, save=True, **kwargs):
             task_collection=self.collection, protocol_number=self.protocol_number, **kwargs)
         return outputs, files
 
-    def _run_qc(self, trials_data=None, update=False, plot_qc=False):
+    def _run_qc(self, trials_data=None, update=False, plot_qc=False, QC=None):
         if not self.extractor or trials_data is None:
             trials_data, _ = self._extract_behaviour(save=False)
         if not trials_data:
@@ -422,10 +441,11 @@ def _run_qc(self, trials_data=None, update=False, plot_qc=False):
         qc_extractor = TaskQCExtractor(self.session_path, lazy=True, sync_collection=self.sync_collection, one=self.one,
                                        sync_type=self.sync, task_collection=self.collection)
         qc_extractor.data = qc_extractor.rename_data(trials_data.copy())
-        if type(self.extractor).__name__ == 'HabituationTrials':
-            qc = HabituationQC(self.session_path, one=self.one, log=_logger)
-        else:
-            qc = TaskQC(self.session_path, one=self.one, log=_logger)
+        if not QC:
+            QC = HabituationQC if type(self.extractor).__name__ == 'HabituationTrials' else TaskQC
+        _logger.debug('Running QC with %s.%s', QC.__module__, QC.__name__)
+        qc = QC(self.session_path, one=self.one, log=_logger)
+        if QC is not HabituationQC:
             # Add Bpod wheel data
             wheel_ts_bpod = self.extractor.bpod2fpga(self.extractor.bpod_trials['wheel_timestamps'])
             qc_extractor.data['wheel_timestamps_bpod'] = wheel_ts_bpod

ibllib/pipes/behavior_tasks.py

@@ -9,6 +9,7 @@
 import sys
 import time
 import logging
+import warnings
 from functools import wraps
 from pathlib import Path
 from typing import Union, List, Callable, Any
@@ -365,6 +366,8 @@ def load_params_dict(params_fname: str) -> dict:
 
 
 def load_videopc_params():
+    """(DEPRECATED) This will be removed in favour of iblrigv8 functions."""
+    warnings.warn('load_videopc_params will be removed in favour of iblrigv8', FutureWarning)
     if not load_params_dict("videopc_params"):
         create_videopc_params()
     return load_params_dict("videopc_params")
@@ -472,6 +475,9 @@ def create_basic_transfer_params(param_str='transfer_params', local_data_path=No
 
 
 def create_videopc_params(force=False, silent=False):
+    """(DEPRECATED) This will be removed in favour of iblrigv8 functions."""
+    url = 'https://github.com/int-brain-lab/iblrig/blob/videopc/docs/source/video.rst'
+    warnings.warn(f'create_videopc_params is deprecated, see {url}', DeprecationWarning)
     if Path(params.getfile("videopc_params")).exists() and not force:
         print(f"{params.getfile('videopc_params')} exists already, exiting...")
         print(Path(params.getfile("videopc_params")).exists())