add channel files to spikesorting task

Author: mayofaulkner

ibllib/pipes/ephys_preprocessing.py

@@ -20,6 +20,7 @@
 from ibllib.pipes import tasks
 from ibllib.pipes.training_preprocessing import TrainingRegisterRaw as EphysRegisterRaw
 from ibllib.pipes.misc import create_alyx_probe_insertions
+from ibllib.qc.alignment_qc import get_aligned_channels
 from ibllib.qc.task_extractors import TaskQCExtractor
 from ibllib.qc.task_metrics import TaskQC
 from ibllib.qc.camera import run_all_qc as run_camera_qc
@@ -406,6 +407,18 @@ def _run(self, probes=None):
                 tar_dir.mkdir(parents=True, exist_ok=True)
                 out = spikes.ks2_to_tar(ks2_dir, tar_dir, force=self.FORCE_RERUN)
                 out_files.extend(out)
+
+                if self.one:
+                    eid = self.one.path2eid(self.session_path, query_type='remote')
+                    ins = self.one.alyx.rest('insertions', 'list', session=eid, name=label)
+                    if len(ins) != 0:
+                        resolved = ins[0].get('json', {'temp': 0}).get('extended_qc', {'temp': 0}). \
+                            get('alignment_resolved', False)
+                        if resolved:
+                            chns = np.load(probe_out_path.joinpath('channels.localCoordinates.npy'))
+                            out = get_aligned_channels(ins[0], chns, one=self.one, save_dir=probe_out_path)
+                            out_files.extend(out)
+
             except BaseException:
                 _logger.error(traceback.format_exc())
                 self.status = -1

ibllib/qc/alignment_qc.py

@@ -1,6 +1,7 @@
 import logging
 
 import numpy as np
+from pathlib import Path
 
 from ibllib.atlas import AllenAtlas
 from ibllib.pipes import histology
@@ -387,3 +388,32 @@ def update_experimenter_evaluation(self, prev_alignments=None, override=False):
         else:
             self.log.warning(f'No experimenter qc found, qc field of probe insertion {self.eid} '
                              f'will not be updated')
+
+
+def get_aligned_channels(ins, chn_coords, one, ba=None, save_dir=None):
+
+    ba = ba or AllenAtlas(25)
+    depths = chn_coords[:, 1]
+    xyz = np.array(ins['json']['xyz_picks']) / 1e6
+    traj = one.alyx.rest('trajectories', 'list', probe_insertion=ins['id'],
+                         provenance='Ephys aligned histology track')[0]
+    align_key = ins['json']['extended_qc']['alignment_stored']
+    feature = traj['json'][align_key][0]
+    track = traj['json'][align_key][1]
+    ephysalign = EphysAlignment(xyz, depths, track_prev=track,
+                                feature_prev=feature,
+                                brain_atlas=ba, speedy=True)
+    channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6)
+    channels_atlas_ids = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id']
+
+    out_files = []
+    if save_dir is not None:
+        f_name = Path(save_dir).joinpath('channels.mlapdv.npy')
+        np.save(f_name, channels_mlapdv)
+        out_files.append(f_name)
+
+        f_name = Path(save_dir).joinpath('channels.brainLocationIds_ccf_2017.npy')
+        np.save(f_name, channels_atlas_ids)
+        out_files.append(f_name)
+
+    return out_files