Merge remote-tracking branch 'origin/develop' into manual_curation_ssloader

Author: mayofaulkner

ibllib/io/extractors/mesoscope.py

@@ -672,8 +672,8 @@ def _extract(self, sync=None, chmap=None, device_collection='raw_imaging_data',
             _, fov_time_shifts, line_time_shifts = self.get_timeshifts(imaging_data['meta'])
             assert len(fov_time_shifts) == self.n_FOVs, f'unexpected number of FOVs for {collection}'
             ts = frame_times[np.logical_and(frame_times >= tmin, frame_times <= tmax)]
-            assert ts.size >= imaging_data['times_scanImage'].size, (f'fewer DAQ timestamps for {collection} '
-                                                                     f'than expected')
+            assert ts.size >= imaging_data[
+                'times_scanImage'].size, f"fewer DAQ timestamps for {collection} than expected: DAQ/frames = {ts.size}/{imaging_data['times_scanImage'].size}"
             if ts.size > imaging_data['times_scanImage'].size:
                 _logger.warning(
                     'More DAQ frame times detected for %s than were found in the raw image data.\n'

ibllib/pipes/dynamic_pipeline.py

@@ -232,15 +232,21 @@ def _get_trials_tasks(session_path, acquisition_description=None, sync_tasks=Non
     sync_kwargs = {'sync': sync, **sync_args}
 
     # Behavior tasks
+    protocol_numbers = set()
     task_protocols = acquisition_description.get('tasks', [])
     for i, (protocol, task_info) in enumerate(chain(*map(dict.items, task_protocols))):
         collection = task_info.get('collection', f'raw_task_data_{i:02}')
         task_kwargs = {'protocol': protocol, 'collection': collection}
-        # For now the order of protocols in the list will take precedence. If collections are numbered,
-        # check that the numbers match the order.  This may change in the future.
+        # The order of protocols in the list will take precedence unless protocol_number is present.
+        # If collections are numbered, check that the numbers match the order.
+        n = i
         if re.match(r'^raw_task_data_\d{2}$', collection):
-            task_kwargs['protocol_number'] = i
-            if int(collection.split('_')[-1]) != i:
+            # Protocol number may be overridden by the protocol_number key
+            if task_info.get('protocol_number') is not None:
+                n = task_info['protocol_number']
+            task_kwargs['protocol_number'] = n
+            assert not (n in protocol_numbers or protocol_numbers.add(n)), 'protocol numbers must be unique'
+            if int(collection.split('_')[-1]) != n:
                 _logger.warning('Number in collection name does not match task order')
         if extractors := task_info.get('extractors', False):
             extractors = (extractors,) if isinstance(extractors, str) else extractors

ibllib/tests/test_dynamic_pipeline.py

@@ -3,6 +3,7 @@
 import unittest
 from unittest import mock
 from itertools import chain
+from copy import deepcopy
 
 import yaml
 
@@ -79,6 +80,20 @@ def test_get_trials_tasks(self):
         self.assertIsInstance(tasks[0], btasks.ChoiceWorldTrialsNidq)
         one.load_datasets.assert_called()  # check that description file is checked on disk
 
+        # A session where protocol number is specified
+        description = deepcopy(self.description)
+        description['tasks'][1]['passiveChoiceWorld']['protocol_number'] = 2
+        fcn_to_patch = 'ibllib.pipes.dynamic_pipeline._load_acquisition_description'
+        with mock.patch(fcn_to_patch, return_value=description), self.assertLogs(dyn.__name__, 'WARNING') as cm:
+            tasks = dyn.get_trials_tasks(self.session_path_dynamic)
+        self.assertEqual(2, tasks[1].protocol_number)
+        self.assertRegex(cm.records[0].getMessage(), 'Number in collection name does not match task order')
+        # If protocol number is not unique, an assertion error should be raised
+        description = deepcopy(self.description)
+        description['tasks'][1]['passiveChoiceWorld']['protocol_number'] = 0
+        with mock.patch(fcn_to_patch, return_value=description):
+            self.assertRaises(AssertionError, dyn.get_trials_tasks, self.session_path_dynamic)
+
         # A session with timeline acquisition
         self.description['sync']['nidq']['acquisition_software'] = 'timeline'
         with open(self.session_path_dynamic / '_ibl_experiment.description.yaml', 'w') as fp: