viewspikes: co-locate driftmap and PSD in one figure

Author: oliche

viewspikes/examples.py

@@ -1,10 +1,14 @@
 import numpy as np
 import scipy
+import matplotlib.pyplot as plt
 from easyqc.gui import viewseis
 
+
 from oneibl.one import ONE
 from ibllib.ephys import neuropixel
 from ibllib.dsp import voltage
+from ibllib.plots import color_cycle
+from brainbox.plot import driftmap
 
 from iblapps.needles2 import run_needles2
 from iblapps.viewspikes.data import stream, get_ks2, get_spikes
@@ -22,7 +26,9 @@
 
 ## Example 3: Show the PSD
 raw = sr[:, :-1].T
-show_psd(raw, sr.fs)
+fig, axes = plt.subplots(1, 2)
+fig.set_size_inches(18, 7)
+show_psd(raw, sr.fs, ax=axes[0])
 
 ## Example 4: Display the raw / pre-proc and KS2 parts -
 h = neuropixel.trace_header()
@@ -41,6 +47,16 @@
 overlay_spikes(eqc_dest, spikes, clusters, channels)
 overlay_spikes(eqc_ks2, spikes, clusters, channels)
 
-# hhh = {k: np.tile(h[k], 3) for k in h}
+# Do the driftmap
+driftmap(spikes['times'], spikes['depths'], t_bin=0.1, d_bin=5, ax=axes[1])
+
+
 # eqc_concat = viewseis(np.r_[butt, destripe, ks2], si=1 / sr.fs, h=hhh, t0=t0, title='concat')
 # overlay_spikes(eqc_concat, spikes, clusters, channels)
+from easyqc import qt
+import datetime
+qtapp = qt.create_app()
+screenshot = qtapp.primaryScreen().grabWindow(eqc_butt.winId())
+
+fn = datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S')
+screenshot.save(f'/home/olivier/Pictures/{fn}.png', 'png')

viewspikes/plots.py

@@ -7,25 +7,25 @@
 import ibllib.atlas as atlas
 from ibllib.ephys.neuropixel import SITES_COORDINATES
 from ibllib.pipes.ephys_alignment import EphysAlignment
-from ibllib.plots import wiggle
+from ibllib.plots import wiggle, color_cycle
 
 brain_atlas = atlas.AllenAtlas()
 # Instantiate brain atlas and one
 
 
-def show_psd(data, fs):
+def show_psd(data, fs, ax=None):
     psd = np.zeros((data.shape[0], 129))
     for tr in np.arange(data.shape[0]):
         f, psd[tr, :] = scipy.signal.welch(data[tr, :], fs=fs)
 
-    plt.figure()
-    plt.plot(f, 10 * np.log10(psd.T), color='gray', alpha=0.1)
-    plt.plot(f, 10 * np.log10(np.mean(psd, axis=0).T), color='red')
-    plt.xlabel('Frequency (Hz)')
-    plt.ylabel('PSD (dB rel V/Hz)')
-    plt.gca()
-    plt.gca().set_ylim(-150, -110)
-    plt.gca().set_xlim(0, fs / 2)
+    if ax is None:
+        fig, ax = plt.subplots()
+    ax.plot(f, 10 * np.log10(psd.T), color='gray', alpha=0.1)
+    ax.plot(f, 10 * np.log10(np.mean(psd, axis=0).T), color='red')
+    ax.set_xlabel('Frequency (Hz)')
+    ax.set_ylabel('PSD (dB rel V/Hz)')
+    ax.set_ylim(-150, -110)
+    ax.set_xlim(0, fs / 2)
     plt.show()
 
 
@@ -146,7 +146,11 @@ def overlay_spikes(self, spikes, clusters, channels):
     sc = self.layers['default']['layer']
     sc.setSize(8)
     sc.setSymbol('x')
-    sc.setPen(pg.mkPen((0, 255, 0, 155), width=1))
+    # sc.setPen(pg.mkPen((0, 255, 0, 155), width=1))
+    rgbs = [list((rgb * 255).astype(np.uint8)) for rgb in color_cycle(spikes['clusters'][ifirst:ilast])]
+    sc.setBrush([pg.mkBrush(rgb) for rgb in rgbs])
+    sc.setPen([pg.mkPen(rgb) for rgb in rgbs])
+    return sc
 
     # sc.setData(x=xspi, y=tspi, brush=pg.mkBrush((255, 0, 0)))
     def callback(sc, points, evt):