remove wheel qc checks if habituation protocol

Author: mayofaulkner

ibllib/pipes/video_tasks.py

@@ -252,7 +252,8 @@ def run_qc(self, camera_data=None, update=True):
         if camera_data is None:
             camera_data, _ = self.extract_camera(save=False)
         qc = CameraQC(
-            self.session_path, 'left', sync_type='bpod', sync_collection=self.collection, one=self.one)
+            self.session_path, 'left', sync_type='bpod', sync_collection=self.collection, one=self.one,
+            protocol=self.protocol)
         qc.run(update=update)
         return qc
 

ibllib/qc/camera.py

@@ -135,6 +135,7 @@ def __init__(self, session_path_or_eid, camera, **kwargs):
         self.n_samples = kwargs.pop('n_samples', 100)
         self.sync_collection = kwargs.pop('sync_collection', None)
         self.sync = kwargs.pop('sync_type', None)
+        self.protocol = kwargs.pop('protocol', None)
         super().__init__(session_path_or_eid, **kwargs)
 
         # Data
@@ -163,7 +164,10 @@ def __init__(self, session_path_or_eid, camera, **kwargs):
         self.outcome = spec.QC.NOT_SET
 
         # Specify any checks to remove
-        self.checks_to_remove = []
+        if self.protocol is not None and 'habituation' in self.protocol:
+            self.checks_to_remove = ['check_wheel_alignment']
+        else:
+            self.checks_to_remove = []
         self._type = None
 
     @property
@@ -271,8 +275,12 @@ def load_data(self, extract_times: bool = False, load_video: bool = True) -> Non
                 else:
                     raise NotImplementedError(f'Unknown namespace "{ns}"')
             else:
-                wheel_data = training_wheel.get_wheel_position(
-                    self.session_path, task_collection=task_collection)
+                if self.protocol is not None and 'habituation' in self.protocol:
+                    wheel_data = training_wheel.get_wheel_position(
+                        self.session_path, task_collection=task_collection)
+                else:
+                    wheel_data = [None, None]
+
             self.data['wheel'] = Bunch(zip(wheel_keys, wheel_data))
 
         # Find short period of wheel motion for motion correlation.