Repository: int-brain-lab/ibllib

Issue:
mesoscope extractor expects nr of FOVs == nr or ROIs in raw SI tif
For multi-depth imaging, the nr of FOVs we define in the metadata (one per unique scanfield*depth) is not equal to the nr of 'enabled' ROIs in the header of the scanimage tifs (one per unique scanfield). However, the mesoscope extractor currently asserts this and returns an error if it doesn't, causing MesoscopeSync task to fail. This assertion should either be removed or replaced with something more general referring to the unique roiUUIDs e.g. something like
`assert sum(x['enable'] for x in artist['RoiGroups']['imagingRoiGroup']['rois']) == len(np.unique(raw_imaging_meta['FOV']['roiUUID']))
`
https://github.com/int-brain-lab/ibllib/blob/ced2259847fb2663f460c762649e4317b67abb2a/ibllib/io/extractors/mesoscope.py#L770

Discussion:

Patch:
diff --git a/ibllib/io/extractors/mesoscope.py b/ibllib/io/extractors/mesoscope.py
index ade22240..8754869a 100644
--- a/ibllib/io/extractors/mesoscope.py
+++ b/ibllib/io/extractors/mesoscope.py
@@ -746,6 +746,15 @@ class MesoscopeSyncTimeline(extractors_base.BaseExtractor):
         Calculate the time shifts for each field of view (FOV) and the relative offsets for each
         scan line.
 
+        For a 2 scan field, 2 depth recording (so 4 FOVs):
+
+        Frame 1, lines 1-512 correspond to FOV_00
+        Frame 1, lines 551-1062 correspond to FOV_01
+        Frame 2, lines 1-512 correspond to FOV_02
+        Frame 2, lines 551-1062 correspond to FOV_03
+        Frame 3, lines 1-512 correspond to FOV_00
+        ...
+
         Parameters
         ----------
         raw_imaging_meta : dict
@@ -765,26 +774,27 @@ class MesoscopeSyncTimeline(extractors_base.BaseExtractor):
         FOVs = raw_imaging_meta['FOV']
 
         # Double-check meta extracted properly
-        raw_meta = raw_imaging_meta['rawScanImageMeta']
-        artist = raw_meta['Artist']
-        assert sum(x['enable'] for x in artist['RoiGroups']['imagingRoiGroup']['rois']) == len(FOVs)
-
+        # assert meta.FOV.Zs is ascending but use slice_id field. This may not be necessary but is expected.
+        slice_ids = np.array([fov['slice_id'] for fov in FOVs])
+        assert np.all(np.diff([x['Zs'] for x in FOVs]) >= 0), 'FOV depths not in ascending order'
+        assert np.all(np.diff(slice_ids) >= 0), 'slice IDs not ordered'
         # Number of scan lines per FOV, i.e. number of Y pixels / image height
         n_lines = np.array([x['nXnYnZ'][1] for x in FOVs])
-        n_valid_lines = np.sum(n_lines)  # Number of lines imaged excluding flybacks
-        # Number of lines during flyback
-        n_lines_per_gap = int((raw_meta['Height'] - n_valid_lines) / (len(FOVs) - 1))
-        # The start and end indices of each FOV in the raw images
-        fov_start_idx = np.insert(np.cumsum(n_lines[:-1] + n_lines_per_gap), 0, 0)
-        fov_end_idx = fov_start_idx + n_lines
-        line_period = raw_imaging_meta['scanImageParams']['hRoiManager']['linePeriod']
-
-        line_indices = []
-        fov_time_shifts = fov_start_idx * line_period
-        line_time_shifts = []
 
-        for ln, s, e in zip(n_lines, fov_start_idx, fov_end_idx):
-            line_indices.append(np.arange(s, e))
-            line_time_shifts.append(np.arange(0, ln) * line_period)
+        # We get indices from MATLAB extracted metadata so below two lines are no longer needed
+        # n_valid_lines = np.sum(n_lines)  # Number of lines imaged excluding flybacks
+        # n_lines_per_gap = int((raw_meta['Height'] - n_valid_lines) / (len(FOVs) - 1))  # N lines during flyback
+        line_period = raw_imaging_meta['scanImageParams']['hRoiManager']['linePeriod']
+        frame_time_shifts = slice_ids / raw_imaging_meta['scanImageParams']['hRoiManager']['scanFrameRate']
+
+        # Line indices are now extracted by the MATLAB function mesoscopeMetadataExtraction.m
+        # They are indexed from 1 so we subtract 1 to convert to zero-indexed
+        line_indices = [np.array(fov['lineIdx']) - 1 for fov in FOVs]  # Convert to zero-indexed from MATLAB 1-indexed
+        assert all(lns.size == n for lns, n in zip(line_indices, n_lines)), 'unexpected number of scan lines'
+        # The start indices of each FOV in the raw images
+        fov_start_idx = np.array([lns[0] for lns in line_indices])
+        roi_time_shifts = fov_start_idx * line_period   # The time offset for each FOV
+        fov_time_shifts = roi_time_shifts + frame_time_shifts
+        line_time_shifts = [(lns - ln0) * line_period for lns, ln0 in zip(line_indices, fov_start_idx)]
 
         return line_indices, fov_time_shifts, line_time_shifts
diff --git a/ibllib/pipes/mesoscope_tasks.py b/ibllib/pipes/mesoscope_tasks.py
index fae64a1f..ba103337 100644
--- a/ibllib/pipes/mesoscope_tasks.py
+++ b/ibllib/pipes/mesoscope_tasks.py
@@ -626,10 +626,10 @@ class MesoscopeSync(base_tasks.MesoscopeTask):
                             (f'_{self.sync_namespace}_DAQdata.timestamps.npy', self.sync_collection, True),
                             (f'_{self.sync_namespace}_DAQdata.meta.json', self.sync_collection, True),
                             ('_ibl_rawImagingData.meta.json', self.device_collection, True),
-                            ('rawImagingData.times_scanImage.npy', self.device_collection, True),
+                            ('rawImagingData.times_scanImage.npy', self.device_collection, True, True),  # register raw
                             (f'_{self.sync_namespace}_softwareEvents.log.htsv', self.sync_collection, False), ],
             'output_files': [('mpci.times.npy', 'alf/mesoscope/FOV*', True),
-                             ('mpciStack.timeshift.npy', 'alf/mesoscope/FOV*', True), ]
+                             ('mpciStack.timeshift.npy', 'alf/mesoscope/FOV*', True),]
         }
         return signature
 
diff --git a/ibllib/pipes/tasks.py b/ibllib/pipes/tasks.py
index e03be0d5..a409ed1e 100644
--- a/ibllib/pipes/tasks.py
+++ b/ibllib/pipes/tasks.py
@@ -88,6 +88,7 @@ from ibllib.oneibl.registration import get_lab
 from iblutil.util import Bunch
 import one.params
 from one.api import ONE
+from one.util import ensure_list
 from one import webclient
 
 _logger = logging.getLogger(__name__)
@@ -107,7 +108,7 @@ class Task(abc.ABC):
     time_elapsed_secs = None
     time_out_secs = 3600 * 2  # time-out after which a task is considered dead
     version = ibllib.__version__
-    signature = {'input_files': [], 'output_files': []}  # list of tuples (filename, collection, required_flag)
+    signature = {'input_files': [], 'output_files': []}  # list of tuples (filename, collection, required_flag[, register])
     force = False  # whether to re-download missing input files on local server if not present
     job_size = 'small'  # either 'small' or 'large', defines whether task should be run as part of the large or small job services
     env = None  # the environment name within which to run the task (NB: the env is not activated automatically!)
@@ -173,6 +174,13 @@ class Task(abc.ABC):
             -1: Errored
             -2: Didn't run as a lock was encountered
             -3: Incomplete
+
+        Notes
+        -----
+        - The `run_alyx_task` will update the Alyx Task status depending on both status and outputs
+          (i.e. the output of subclassed `_run` method):
+          Assuming a return value of 0... if Task.outputs is None, the status will be Empty;
+          if Task.outputs is a list (empty or otherwise), the status will be Complete.
         """
         # if task id of one properties are not available, local run only without alyx
         use_alyx = self.one is not None and self.taskid is not None
@@ -201,12 +209,13 @@ class Task(abc.ABC):
         start_time = time.time()
         try:
             setup = self.setUp(**kwargs)
+            self.outputs = self._input_files_to_register()
             _logger.info(f'Setup value is: {setup}')
             self.status = 0
             if not setup:
                 # case where outputs are present but don't have input files locally to rerun task
                 # label task as complete
-                _, self.outputs = self.assert_expected_outputs()
+                _, outputs = self.assert_expected_outputs()
             else:
                 # run task
                 if self.gpu >= 1:
@@ -218,8 +227,15 @@ class Task(abc.ABC):
                         _logger.removeHandler(ch)
                         ch.close()
                         return self.status
-                self.outputs = self._run(**kwargs)
+                outputs = self._run(**kwargs)
                 _logger.info(f'Job {self.__class__} complete')
+            if outputs is None:
+                # If run method returns None and no raw input files were registered, self.outputs
+                # should be None, meaning task will have an 'Empty' status. If run method returns
+                # a list, the status will be 'Complete' regardless of whether there are output files.
+                self.outputs = outputs if not self.outputs else self.outputs  # ensure None if no inputs registered
+            else:
+                self.outputs.extend(ensure_list(outputs))  # Add output files to list of inputs to register
         except Exception:
             _logger.error(traceback.format_exc())
             _logger.info(f'Job {self.__class__} errored')
@@ -263,6 +279,41 @@ class Task(abc.ABC):
         _ = self.register_images()
         return self.data_handler.uploadData(self.outputs, self.version, **kwargs)
 
+    def _input_files_to_register(self, assert_all_exist=False):
+        """
+        Return input datasets to be registered to Alyx.
+
+        These datasets are typically raw data files and are registered even if the task fails to complete.
+
+        Parameters
+        ----------
+        assert_all_exist
+            Raise AssertionError if not all required input datasets exist on disk.
+
+        Returns
+        -------
+        list of pathlib.Path
+            A list of input files to register.
+        """
+        try:
+            input_files = self.input_files
+        except AttributeError:
+            raise RuntimeError('Task.setUp must be run before calling this method.')
+        to_register, missing = [], []
+        for filename, collection, required, _ in filter(lambda f: len(f) > 3 and f[3], input_files):
+            filepath = self.session_path.joinpath(collection, filename)
+            if filepath.exists():
+                to_register.append(filepath)
+            elif required:
+                missing.append(filepath)
+        if any(missing):
+            missing_str = ', '.join(map(lambda x: x.relative_to(self.session_path).as_posix(), missing))
+            if assert_all_exist:
+                raise AssertionError(f'Missing required input files: {missing_str}')
+            else:
+                _logger.error(f'Missing required input files: {missing_str}')
+        return list(set(to_register) - set(missing))
+
     def register_images(self, **kwargs):
         """
         Registers images to alyx database
@@ -766,7 +817,7 @@ def run_alyx_task(tdict=None, session_path=None, one=None, job_deck=None,
     patch_data = {'time_elapsed_secs': task.time_elapsed_secs, 'log': task.log,
                   'version': task.version}
     # if there is no data to register, set status to Empty
-    if task.outputs is None:
+    if task.outputs is None:  # NB: an empty list is still considered Complete.
         patch_data['status'] = 'Empty'
     # otherwise register data and set (provisional) status to Complete
     else:


Test patch:
diff --git a/ibllib/tests/extractors/test_mesoscope.py b/ibllib/tests/extractors/test_mesoscope.py
new file mode 100644
index 00000000..b9447b63
--- /dev/null
+++ b/ibllib/tests/extractors/test_mesoscope.py
@@ -0,0 +1,64 @@
+"""Tests for ibllib.io.extractors.mesoscope module."""
+import unittest
+from itertools import repeat, chain
+
+import numpy as np
+
+from ibllib.io.extractors import mesoscope
+
+
+class TestMesoscopeSyncTimeline(unittest.TestCase):
+    """Tests for MesoscopeSyncTimeline extractor class."""
+    def setUp(self) -> None:
+        """Simulate for meta data for 9 FOVs at 3 different depths.
+
+        These simulated values match those from SP048/2024-02-05/001.
+        """
+        n_lines_flyback = 75
+        self.n_lines = 512
+        self.n_FOV = 9
+        n_depths = 3
+        assert self.n_FOV > n_depths and self.n_FOV % n_depths == 0
+        reps = int(self.n_FOV / n_depths)
+        start_depth = 60
+        delta_depth = 40
+        self.line_period = 4.158e-05
+
+        self.meta = {
+            'scanImageParams': {'hRoiManager': {'linePeriod': self.line_period, 'scanFrameRate': 13.6803}},
+            'FOV': []
+        }
+        nXnYnZ = [self.n_lines, self.n_lines, 1]
+        for i, slice_id in enumerate(chain.from_iterable(map(lambda x: list(repeat(x, reps)), range(n_depths)))):
+            offset = (i % n_depths) * (self.n_lines + n_lines_flyback) - ((i % n_depths) - 1)
+            offset = offset or 1  # start at 1 for MATLAB indexing
+            fov = {'slice_id': slice_id, 'Zs': start_depth + (delta_depth * slice_id),
+                   'nXnYnZ': nXnYnZ, 'lineIdx': list(range(offset, self.n_lines + offset))}
+            self.meta['FOV'].append(fov)
+
+    def test_get_timeshifts_multidepth(self):
+        """Test MescopeSyncTimeline.get_timeshifts method.
+
+        This tests output when given multiple FOVs at different depths. The tasks/mesoscope_tasks.py
+        module in iblscripts more thoroughly tests single-depth imaging with real data.
+        """
+        line_indices, fov_time_shifts, line_time_shifts = mesoscope.MesoscopeSyncTimeline.get_timeshifts(self.meta)
+        expected = [np.array(x['lineIdx']) for x in self.meta['FOV']]
+        self.assertTrue(np.all(x == y) for x, y in zip(expected, line_indices))
+        self.assertEqual(self.n_FOV, len(fov_time_shifts))
+        self.assertEqual(self.n_FOV, len(line_time_shifts))
+        self.assertTrue(all(len(x) == self.n_lines for x in line_time_shifts))
+
+        expected = self.line_period * np.arange(self.n_lines)
+        for i, line_shifts in enumerate(line_time_shifts):
+            with self.subTest(f'FOV == {i}'):
+                self.assertEqual(self.n_lines, len(line_shifts))
+                np.testing.assert_almost_equal(expected, line_shifts)
+
+        # NB: The following values are fixed for the setup parameters
+        expected = [0., 0.02436588, 0.04873176, 0.07309781, 0.09746369, 0.12182957, 0.14619562, 0.1705615, 0.19492738]
+        np.testing.assert_almost_equal(expected, fov_time_shifts)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/ibllib/tests/test_mesoscope.py b/ibllib/tests/test_mesoscope.py
index 7ae3e5cd..f62d42fa 100644
--- a/ibllib/tests/test_mesoscope.py
+++ b/ibllib/tests/test_mesoscope.py
@@ -1,4 +1,4 @@
-"""Tests for ibllib.pipes.mesoscope_tasks"""
+"""Tests for ibllib.pipes.mesoscope_tasks."""
 import sys
 import unittest
 from unittest import mock
diff --git a/ibllib/tests/test_tasks.py b/ibllib/tests/test_tasks.py
index a3f5cff0..482697b3 100644
--- a/ibllib/tests/test_tasks.py
+++ b/ibllib/tests/test_tasks.py
@@ -113,6 +113,7 @@ class TaskGpuLock(ibllib.pipes.tasks.Task):
     def setUp(self):
         self.make_lock_file()
         self.data_handler = self.get_data_handler()
+        self.input_files = []
         return True
 
     def _run(self, overwrite=False):
@@ -357,6 +358,42 @@ class TestDynamicTask(unittest.TestCase):
         self.assertEqual('raw_ephys_data/probe00', collection)
 
 
+class TestTask(unittest.TestCase):
+    def setUp(self):
+        tmpdir = tempfile.TemporaryDirectory()
+        self.addCleanup(tmpdir.cleanup)
+        self.tmpdir = Path(tmpdir.name)
+        self.session_path = self.tmpdir.joinpath('subject', '1900-01-01', '001')
+        self.session_path.mkdir(parents=True)
+
+    def test_input_files_to_register(self):
+        """Test for Task._input_files_to_register method."""
+        task = Task00(self.session_path)
+        self.assertRaises(RuntimeError, task._input_files_to_register)
+        task.input_files = [('register.optional.ext', 'alf', False, True),
+                            ('register.optional_foo.ext', 'alf', False, True),
+                            ('register.required.ext', 'alf', True, True),
+                            ('ignore.required.ext', 'alf', True),
+                            ('ignore.optional.ext', 'alf', False)]
+        self.session_path.joinpath('alf').mkdir()
+        for f in task.input_files:
+            self.session_path.joinpath(f[1], f[0]).touch()
+        files = task._input_files_to_register(assert_all_exist=True)
+        expected = [self.session_path.joinpath('alf', 'register.required.ext'),
+                    self.session_path.joinpath('alf', 'register.optional.ext'),
+                    self.session_path.joinpath('alf', 'register.optional_foo.ext')]
+        self.assertCountEqual(files, expected)
+        expected[2].unlink()
+        with self.assertNoLogs(ibllib.pipes.tasks.__name__, level='ERROR'):
+            files = task._input_files_to_register(assert_all_exist=True)
+            self.assertCountEqual(files, expected[:2])
+        expected[0].unlink()
+        with self.assertLogs(ibllib.pipes.tasks.__name__, level='ERROR'):
+            files = task._input_files_to_register(assert_all_exist=False)
+        self.assertEqual(files, expected[1:2])
+        self.assertRaises(AssertionError, task._input_files_to_register, assert_all_exist=True)
+
+
 class TestMisc(unittest.TestCase):
     """Tests for misc functions in ibllib.pipes.tasks module."""