Merge pull request #374 from int-brain-lab/alignment_qc_pykilosort

Alignment qc pykilosort

Author: mayofaulkner

ibllib/qc/alignment_qc.py

@@ -5,7 +5,6 @@
 from ibllib.atlas import AllenAtlas
 from ibllib.atlas.regions import BrainRegions
 from ibllib.pipes import histology
-from ibllib.ephys.neuropixel import SITES_COORDINATES
 from ibllib.pipes.ephys_alignment import EphysAlignment
 from ibllib.qc import base
 from ibllib.oneibl.patcher import FTPPatcher
@@ -20,14 +19,15 @@ class AlignmentQC(base.QC):
     Class that is used to update the extended_qc of the probe insertion fields with the results
     from the ephys alignment procedure
     """
-    def __init__(self, probe_id, one=None, brain_atlas=None, channels=True):
+    def __init__(self, probe_id, one=None, brain_atlas=None, channels=True, collection=None):
         super().__init__(probe_id, one=one, log=_log, endpoint='insertions')
 
         # Data
         self.alignments = None
         self.xyz_picks = None
         self.depths = None
         self.cluster_chns = None
+        self.chn_coords = None
         self.align_keys_sorted = None
 
         # Metrics and passed trials
@@ -43,7 +43,10 @@ def __init__(self, probe_id, one=None, brain_atlas=None, channels=True):
         self.resolved = (self.insertion.get('json', {'temp': 0}).get('extended_qc').
                          get('alignment_resolved', False))
 
-    def load_data(self, prev_alignments=None, xyz_picks=None, depths=None, cluster_chns=None):
+        self.probe_collection = collection
+
+    def load_data(self, prev_alignments=None, xyz_picks=None, depths=None, cluster_chns=None,
+                  chn_coords=None):
         """"
         Load data required to assess alignment qc and compute similarity matrix. If no arguments
         are given load_data will fetch all the relevant data required
@@ -71,16 +74,29 @@ def load_data(self, prev_alignments=None, xyz_picks=None, depths=None, cluster_c
         else:
             self.xyz_picks = xyz_picks
 
+        if not self.probe_collection:
+            all_collections = self.one.list_collections(self.insertion['session'])
+            if f'alf/{self.insertion["name"]}/pykilosort' in all_collections:
+                self.probe_collection = f'alf/{self.insertion["name"]}/pykilosort'
+            else:
+                self.probe_collection = f'alf/{self.insertion["name"]}'
+
+        if not np.any(chn_coords):
+            self.chn_coords = self.one.load_dataset(self.insertion['session'],
+                                                    'channels.localCoordinates.npy',
+                                                    collection=self.probe_collection)
+        else:
+            self.chn_coords = chn_coords
+
         if not np.any(depths):
-            self.depths = SITES_COORDINATES[:, 1]
+            self.depths = self.chn_coords[:, 1]
         else:
             self.depths = depths
 
         if not np.any(cluster_chns):
             self.cluster_chns = self.one.load_dataset(self.insertion['session'],
                                                       'clusters.channels.npy',
-                                                      collection=f'alf/{self.insertion["name"]}',
-                                                      download_only=True)
+                                                      collection=self.probe_collection)
         else:
             self.cluster_chns = cluster_chns
 
@@ -142,7 +158,7 @@ def run(self, update=True, upload_alyx=True, upload_flatiron=True):
             if update:
                 self.update_extended_qc(results)
 
-            if results['alignment_resolved'] and (upload_alyx or upload_flatiron):
+            if results['alignment_resolved'] and np.bitwise_or(upload_alyx, upload_flatiron):
                 self.upload_channels(results['alignment_stored'], upload_alyx, upload_flatiron)
 
         return results
@@ -281,7 +297,7 @@ def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
         if upload_flatiron:
             ftp_patcher = FTPPatcher(one=self.one)
             insertion = self.one.alyx.get(f'/insertions/{self.eid}', clobber=True)
-            alf_path = self.one.eid2path(insertion['session']).joinpath('alf', insertion['name'])
+            alf_path = self.one.eid2path(insertion['session']).joinpath(self.probe_collection)
             alf_path.mkdir(exist_ok=True, parents=True)
 
             # Make the channels.mlapdv dataset
@@ -318,7 +334,7 @@ def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
         if upload_alyx:
             if alignment_key != self.align_keys_sorted[0]:
                 histology.register_aligned_track(self.eid, channels_mlapdv / 1e6,
-                                                 chn_coords=SITES_COORDINATES, one=self.one,
+                                                 chn_coords=self.chn_coords, one=self.one,
                                                  overwrite=True, channels=self.channels)
 
                 ephys_traj = self.one.alyx.get(f'/trajectories?&probe_insertion={self.eid}'

ibllib/tests/qc/test_alignment_qc.py

@@ -12,6 +12,7 @@
 from ibllib.pipes.misc import create_alyx_probe_insertions
 from ibllib.qc.alignment_qc import AlignmentQC
 from ibllib.pipes.histology import register_track
+from ibllib.ephys.neuropixel import SITES_COORDINATES
 
 
 EPHYS_SESSION = 'b1c968ad-4874-468d-b2e4-5ffa9b9964e9'
@@ -135,7 +136,8 @@ def _03_alignments_disagree(self):
         traj = one.alyx.rest('trajectories', 'update', id=self.prev_traj_id, data=trajectory)
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
         align_qc.load_data(prev_alignments=traj['json'], xyz_picks=np.array(self.xyz_picks) / 1e6,
-                           cluster_chns=self.cluster_chns)
+                           cluster_chns=self.cluster_chns, depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.run(update=True, upload_alyx=True, upload_flatiron=False)
 
         _verify(self, alignment_qc=0.782216, alignment_resolved=False,
@@ -151,7 +153,8 @@ def _04_alignments_agree(self):
         traj = one.alyx.rest('trajectories', 'update', id=self.prev_traj_id, data=trajectory)
         assert(self.prev_traj_id == traj['id'])
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
-        align_qc.load_data(cluster_chns=self.cluster_chns)
+        align_qc.load_data(cluster_chns=self.cluster_chns, depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.run(update=True, upload_alyx=True, upload_flatiron=False)
 
         _verify(self, alignment_resolved='qc', alignment_qc=0.952319, trajectory_created=False,
@@ -165,7 +168,8 @@ def _05_not_latest_alignments_agree(self):
         self.assertEqual(self.prev_traj_id, traj['id'])
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
         align_qc.load_data(prev_alignments=traj['json'], xyz_picks=np.array(self.xyz_picks) / 1e6,
-                           cluster_chns=self.cluster_chns)
+                           cluster_chns=self.cluster_chns, depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.resolved = 0
         align_qc.run(update=True, upload_alyx=True, upload_flatiron=False)
 
@@ -220,7 +224,9 @@ def _01_normal_computation(self):
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
         align_qc.load_data(prev_alignments=self.traj['json'],
                            xyz_picks=np.array(self.xyz_picks) / 1e6,
-                           cluster_chns=self.cluster_chns)
+                           cluster_chns=self.cluster_chns,
+                           depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.run(update=True, upload_alyx=True, upload_flatiron=False)
         _verify(self,
                 alignment_resolved=False,
@@ -233,7 +239,9 @@ def _02_manual_resolution_latest(self):
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
         align_qc.load_data(prev_alignments=self.traj['json'],
                            xyz_picks=np.array(self.xyz_picks) / 1e6,
-                           cluster_chns=self.cluster_chns)
+                           cluster_chns=self.cluster_chns,
+                           depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.resolve_manual('2020-09-28T15:57:25_mayo', update=True, upload_alyx=True,
                                 upload_flatiron=False)
         _verify(self,
@@ -247,7 +255,9 @@ def _03_manual_resolution_not_latest(self):
         align_qc = AlignmentQC(self.probe_id, one=one, brain_atlas=brain_atlas, channels=False)
         align_qc.load_data(prev_alignments=self.traj['json'],
                            xyz_picks=np.array(self.xyz_picks) / 1e6,
-                           cluster_chns=self.cluster_chns)
+                           cluster_chns=self.cluster_chns,
+                           depths=SITES_COORDINATES[:, 1],
+                           chn_coords=SITES_COORDINATES)
         align_qc.resolve_manual('2020-09-28T10:03:06_alejandro', update=True, upload_alyx=True,
                                 upload_flatiron=False, force=True)
         _verify(self,
@@ -329,7 +339,9 @@ def setUpClass(cls) -> None:
         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)
+                           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)
         print(cls.file_paths)