first iteration of snapshot raster with brain locations

Author: oliche

ibllib/plots/figures.py

@@ -16,8 +16,12 @@
 import one.alf.io as alfio
 from one.alf.exceptions import ALFObjectNotFound
 from ibllib.io.video import get_video_frame, url_from_eid
+from brainbox.plot import driftmap
 from brainbox.behavior.dlc import SAMPLING, plot_trace_on_frame, plot_wheel_position, plot_lick_hist, \
     plot_lick_raster, plot_motion_energy_hist, plot_speed_hist, plot_pupil_diameter_hist
+from brainbox.io.one import load_spike_sorting_fast
+from brainbox.ephys_plots import plot_brain_regions
+
 
 logger = logging.getLogger('ibllib')
 
@@ -32,7 +36,23 @@ class SpikeSorting(ReportSnapshotProbe):
 
     def _run(self):
         """runs for initiated PID, streams data, destripe and check bad channels"""
-        assert self.pid
+        all_here, output_files = self.assert_expected(self.output_files, silent=True)
+        if all_here:
+            return output_files
+        spikes, clusters, channels = load_spike_sorting_fast(
+            eid=self.eid, probe=self.pname, one=self.one, nested=False,
+            dataset_types=['spikes.depths'], brain_regions=self.brain_regions)
+        fig, axs = plt.subplots(1, 2, gridspec_kw={'width_ratios': [.95, .05]}, sharey=True, figsize=(16, 9))
+        driftmap(spikes.times, spikes.depths, t_bin=0.007, d_bin=10, vmax=0.5, ax=axs[0])
+        if 'atlas_id' in channels.keys():
+            plot_brain_regions(channels['atlas_id'], channel_depths=channels['axial_um'],
+                               brain_regions=None, display=True, ax=axs[1])
+        title_str = f"{self.pid_label}, {self.pid}, {spikes.clusters.size:_} spikes, {clusters.depths.size:_} clusters"
+        axs[0].set(ylim=[0, 3800], title=title_str)
+        output_files = [self.output_directory.joinpath("spike_sorting_raster.png")]
+        fig.savefig(output_files[0])
+        plt.close(fig)
+        return output_files
 
     def get_probe_signature(self):
         input_signature = [('spikes.times.npy', f'alf/{self.pname}', True),
@@ -55,7 +75,7 @@ def get_probe_signature(self):
         pname = self.pname
         input_signature = [('*ap.meta', f'raw_ephys_data/{pname}', True),
                            ('*ap.ch', f'raw_ephys_data/{pname}', False)]
-                           # ('*ap.cbin', f'raw_ephys_data/{pname}', False)]
+        # ('*ap.cbin', f'raw_ephys_data/{pname}', False)]
         output_signature = [('raw_ephys_bad_channels.png', f'snapshot/{pname}', True),
                             ('raw_ephys_bad_channels_highpass.png', f'snapshot/{pname}', True),
                             ('raw_ephys_bad_channels_highpass.png', f'snapshot/{pname}', True),
@@ -69,19 +89,18 @@ def _run(self):
         assert self.pid
         SNAPSHOT_LABEL = "raw_ephys_bad_channels"
         eid, pname = self.one.pid2eid(self.pid)
-        output_directory = self.session_path.joinpath('snapshot', pname)
-        output_files = list(output_directory.glob(f'{SNAPSHOT_LABEL}*'))
+        output_files = list(self.output_directory.glob(f'{SNAPSHOT_LABEL}*'))
         if len(output_files) == 4:
             return output_files
-        output_directory.mkdir(exist_ok=True, parents=True)
+        self.output_directory.mkdir(exist_ok=True, parents=True)
         from brainbox.io.spikeglx import stream
         T0 = 60 * 30
         sr, t0 = stream(self.pid, T0, nsecs=1, one=self.one)
         raw = sr[:, :-sr.nsync].T
         channel_labels, channel_features = voltage.detect_bad_channels(raw, sr.fs)
         _, _, output_files = ephys_bad_channels(
             raw=raw, fs=sr.fs, channel_labels=channel_labels, channel_features=channel_features,
-            title=SNAPSHOT_LABEL, destripe=True, save_dir=output_directory)
+            title=SNAPSHOT_LABEL, destripe=True, save_dir=self.output_directory)
         return output_files
 
 

ibllib/plots/snapshot.py

@@ -35,22 +35,34 @@ def register_images(self, widths=None, function=None):
 
 
 class ReportSnapshotProbe(ReportSnapshot):
-
     signature = {
         'input_files': [],  # see setUp method for declaration of inputs
         'output_files': []  # see setUp method for declaration of inputs
     }
 
-    def __init__(self, pid, one=None, **kwargs):
+    def __init__(self, pid, one=None, brain_regions=None, **kwargs):
+        """
+        :param pid: probe insertion UUID from Alyx
+        :param one: one instance
+        :param brain_regions: (optional) ibllib.atlas.BrainRegion object
+        :param kwargs:
+        """
         assert one
         self.one = one
+        self.brain_regions = brain_regions
         self.content_type = 'probeinsertion'
         self.pid = pid
         self.eid, self.pname = self.one.pid2eid(self.pid)
         self.session_path = self.one.eid2path(self.eid)
+        self.output_directory = self.session_path.joinpath('snapshot', self.pname)
         self.get_probe_signature()
         super(ReportSnapshotProbe, self).__init__(self.session_path, object_id=pid, content_type=self.content_type, **kwargs)
 
+    @property
+    def pid_label(self):
+        """returns a probe insertion stub to label titles, for example: 'SWC_054_2020-10-05_001_probe01'"""
+        return '_'.join(list(self.session_path.parts[-3:]) + [self.pname])
+
     @abc.abstractmethod
     def get_probe_signature(self):
         # method that gets input and output signatures from the probe name. The format is a dictionary as follows: