Merge pull request #1018 from int-brain-lab/align

split the dataset creation part from the FTP patcher

Author: oliche

ibllib/qc/alignment_qc.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from one.alf.spec import QC
 from datetime import date
+from typing import Tuple
 
 from neuropixel import trace_header
 import spikeglx
@@ -275,13 +276,8 @@ def compute_alignment_status(self):
 
         return results
 
-    def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
-        """
-        Upload channels to alyx and flatiron based on the alignment specified by the alignment key
-        """
-
-        feature = np.array(self.alignments[alignment_key][0])
-        track = np.array(self.alignments[alignment_key][1])
+    def get_channels_mlapdv(self, alignment_key: str) -> Tuple[np.ndarray, np.ndarray, np.array]:
+        feature, track = np.array(self.alignments[alignment_key][:2])
 
         try:
             meta_dset = self.one.list_datasets(self.insertion['session'], '*ap.meta',
@@ -291,25 +287,82 @@ def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
                                               collection=f'raw_ephys_data/{self.insertion["name"]}',
                                               download_only=True)
             geometry = spikeglx.read_geometry(meta_file)
-            chns = np.c_[geometry['x'], geometry['y']]
+            channels_xy = np.c_[geometry['x'], geometry['y']]
         except Exception as err:
             self.log.warning(f"Could not compute channel locations from meta file, errored with message: {err}. "
                              f"Will use default Neuropixel 1 channels")
             geometry = trace_header(version=1)
-            chns = np.c_[geometry['x'], geometry['y']]
+            channels_xy = np.c_[geometry['x'], geometry['y']]
 
-        ephysalign = EphysAlignment(self.xyz_picks, chns[:, 1],
+        ephysalign = EphysAlignment(self.xyz_picks, channels_xy[:, 1],
                                     track_prev=track,
                                     feature_prev=feature,
                                     brain_atlas=self.brain_atlas)
         channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6)
         channels_atlas_id = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id']
+        return channels_mlapdv, channels_atlas_id, channels_xy
+
+    def create_electrode_datasets(self, alignment_key: str) -> list:
+        """
+        When an insertion alignment is resolved, we write datasets containing electrode locations
+        :return:
+        """
+        channels_mlapdv, channels_atlas_id, channels_xy = self.get_channels_mlapdv(alignment_key)
+        files_to_register = []
+        alf_path = self.one.eid2path(self.insertion['session']).joinpath('alf', self.insertion["name"])
+        alf_path.mkdir(exist_ok=True, parents=True)
+
+        f_name = alf_path.joinpath('electrodeSites.mlapdv.npy')
+        np.save(f_name, channels_mlapdv)
+        files_to_register.append(f_name)
+
+        f_name = alf_path.joinpath('electrodeSites.brainLocationIds_ccf_2017.npy')
+        np.save(f_name, channels_atlas_id)
+        files_to_register.append(f_name)
+
+        f_name = alf_path.joinpath('electrodeSites.localCoordinates.npy')
+        np.save(f_name, channels_xy)
+        files_to_register.append(f_name)
+        probe_collections = self.one.list_collections(self.insertion['session'], filename='channels*',
+                                                      collection=f'alf/{self.insertion["name"]}*')
+        feature, track = np.array(self.alignments[alignment_key][:2])
+        for collection in probe_collections:
+            chns = self.one.load_dataset(self.insertion['session'], 'channels.localCoordinates', collection=collection)
+
+            ephysalign = EphysAlignment(self.xyz_picks, chns[:, 1],
+                                        track_prev=track,
+                                        feature_prev=feature,
+                                        brain_atlas=self.brain_atlas)
+            channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6)
+            channels_atlas_id = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id']
+
+            alf_path = self.one.eid2path(self.insertion['session']).joinpath(collection)
+            alf_path.mkdir(exist_ok=True, parents=True)
+
+            f_name = alf_path.joinpath('channels.mlapdv.npy')
+            np.save(f_name, channels_mlapdv)
+            files_to_register.append(f_name)
+
+            f_name = alf_path.joinpath('channels.brainLocationIds_ccf_2017.npy')
+            np.save(f_name, channels_atlas_id)
+            files_to_register.append(f_name)
+
+        self.log.info("The following files have been saved:")
+        for f in files_to_register:
+            self.log.info(f)
 
+        return files_to_register
+
+    def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
+        """
+        Upload channels to alyx and flatiron based on the alignment specified by the alignment key
+        """
+        channels_mlapdv, channels_atlas_id, channels_xy = self.get_channels_mlapdv(alignment_key)
         # Need to change channels stored on alyx as well as the stored key is not the same as the latest key
         if upload_alyx:
             if alignment_key != self.align_keys_sorted[0]:
                 histology.register_aligned_track(self.eid, channels_mlapdv / 1e6,
-                                                 chn_coords=chns, one=self.one,
+                                                 chn_coords=channels_xy, one=self.one,
                                                  overwrite=True, channels=self.channels_flag,
                                                  brain_atlas=self.brain_atlas)
 
@@ -322,47 +375,9 @@ def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
 
         files_to_register = []
         if upload_flatiron:
+            files_to_register = self.create_electrode_datasets()
+            self.log.info(f"Registering {len(files_to_register)} datasets on FTP patcher")
             ftp_patcher = FTPPatcher(one=self.one)
-
-            alf_path = self.one.eid2path(self.insertion['session']).joinpath('alf', self.insertion["name"])
-            alf_path.mkdir(exist_ok=True, parents=True)
-
-            f_name = alf_path.joinpath('electrodeSites.mlapdv.npy')
-            np.save(f_name, channels_mlapdv)
-            files_to_register.append(f_name)
-
-            f_name = alf_path.joinpath('electrodeSites.brainLocationIds_ccf_2017.npy')
-            np.save(f_name, channels_atlas_id)
-            files_to_register.append(f_name)
-
-            f_name = alf_path.joinpath('electrodeSites.localCoordinates.npy')
-            np.save(f_name, chns)
-            files_to_register.append(f_name)
-
-            probe_collections = self.one.list_collections(self.insertion['session'], filename='channels*',
-                                                          collection=f'alf/{self.insertion["name"]}*')
-
-            for collection in probe_collections:
-                chns = self.one.load_dataset(self.insertion['session'], 'channels.localCoordinates', collection=collection)
-                ephysalign = EphysAlignment(self.xyz_picks, chns[:, 1],
-                                            track_prev=track,
-                                            feature_prev=feature,
-                                            brain_atlas=self.brain_atlas)
-                channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6)
-                channels_atlas_id = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id']
-
-                alf_path = self.one.eid2path(self.insertion['session']).joinpath(collection)
-                alf_path.mkdir(exist_ok=True, parents=True)
-
-                f_name = alf_path.joinpath('channels.mlapdv.npy')
-                np.save(f_name, channels_mlapdv)
-                files_to_register.append(f_name)
-
-                f_name = alf_path.joinpath('channels.brainLocationIds_ccf_2017.npy')
-                np.save(f_name, channels_atlas_id)
-                files_to_register.append(f_name)
-
-            self.log.info("Writing datasets to FlatIron")
             ftp_patcher.create_dataset(path=files_to_register,
                                        created_by=self.one.alyx.user)
 

ibllib/tests/extractors/test_extractors_base.py

@@ -30,11 +30,12 @@ def import_mock(self, name, *args):
     def test_get_task_extractor_map(self):
         """Test ibllib.io.extractors.base._get_task_extractor_map function."""
         # Check the custom map is loaded
-        with patch('builtins.__import__', side_effect=self.import_mock):
+        with patch.dict('sys.modules', {'projects': self.projects}):
             extractors = base._get_task_extractor_map()
             self.assertTrue(self.custom_extractors.items() < extractors.items())
         # Test handles case where module not installed
-        with patch('builtins.__import__', side_effect=ModuleNotFoundError):
+        with patch.dict('sys.modules', {'projects': None}):
+            # When trying to import a module that's None in sys.modules, it will raise ImportError
             extractors = base._get_task_extractor_map()
             self.assertFalse(set(self.custom_extractors.items()).issubset(set(extractors.items())))
         # Remove the file and check exception is caught

ibllib/tests/qc/test_alignment_qc.py

@@ -11,6 +11,7 @@
 
 from one.api import ONE
 from neuropixel import trace_header
+import one.alf.io as alfio
 
 from ibllib.tests import TEST_DB
 from ibllib.tests.fixtures.utils import register_new_session
@@ -398,7 +399,6 @@ class TestUploadToFlatIron(unittest.TestCase):
     xyz_picks = None
     trajectory = None
 
-    @unittest.skip("Skip FTP upload test")
     @classmethod
     def setUpClass(cls) -> None:
         data = np.load(Path(Path(__file__).parent.parent.
@@ -413,6 +413,9 @@ def setUpClass(cls) -> None:
                        allow_pickle=True)
         insertion = data['insertion'].tolist()
         insertion['json'] = {'xyz_picks': cls.xyz_picks}
+        # makes sure there is no existing probe insertion before creating a new one
+        for ins in one.alyx.rest('insertions', 'list', session=insertion['session'], name=insertion['name'], no_cache=True):
+            one.alyx.rest('insertions', 'delete', id=ins['id'])
         probe_insertion = one.alyx.rest('insertions', 'create', data=insertion)
         cls.probe_id = probe_insertion['id']
         cls.probe_name = probe_insertion['name']
@@ -422,23 +425,23 @@ def setUpClass(cls) -> None:
         cls.trajectory.update({'json': cls.alignments})
         cls.traj = one.alyx.rest('trajectories', 'create', data=cls.trajectory)
 
-        align_qc = AlignmentQC(cls.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
-        align_qc.load_data(prev_alignments=cls.traj['json'],
-                           xyz_picks=np.array(cls.xyz_picks) / 1e6,
-                           cluster_chns=cls.cluster_chns,
-                           depths=SITES_COORDINATES[:, 1],
-                           chn_coords=SITES_COORDINATES)
-        cls.file_paths = align_qc.resolve_manual('2020-09-28T15:57:25_mayo', update=True,
-                                                 upload_alyx=True, upload_flatiron=True)
+        cls.align_qc = AlignmentQC(cls.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
+        cls.align_qc.load_data(prev_alignments=cls.traj['json'],
+                               xyz_picks=np.array(cls.xyz_picks) / 1e6,
+                               cluster_chns=cls.cluster_chns,
+                               depths=SITES_COORDINATES[:, 1],
+                               chn_coords=SITES_COORDINATES)
+        cls.file_paths = cls.align_qc.resolve_manual('2020-09-28T15:57:25_mayo', update=True,
+                                                     upload_alyx=True, upload_flatiron=False)
         print(cls.file_paths)
 
     def test_data_content(self):
-        alf_path = one.eid2path(EPHYS_SESSION).joinpath('alf', self.probe_name)
-        channels_mlapdv = np.load(alf_path.joinpath('channels.mlapdv.npy'))
-        self.assertTrue(np.all(np.abs(channels_mlapdv) > 0))
-        channels_id = np.load(alf_path.joinpath('channels.brainLocationIds_ccf_2017.npy'))
-        self.assertEqual(channels_mlapdv.shape[0], channels_id.shape[0])
+        files_to_register = self.align_qc.create_electrode_datasets(alignment_key='2020-09-28T15:57:25_mayo')
+        esites = alfio.load_object(one.eid2path(EPHYS_SESSION).joinpath('alf', self.probe_name), 'electrodeSites')
+        np.testing.assert_array_equal([v.shape[0] for v in esites.values()], 384)
+        self.assertEqual(len(files_to_register), 3)
 
+    @unittest.skip
     def test_upload_to_flatiron(self):
         for file in self.file_paths:
             file_registered = one.alyx.get(f'/datasets?&session={EPHYS_SESSION}'