diff --git a/atlaselectrophysiology/plot_data.py b/atlaselectrophysiology/plot_data.py
index db51324f..f4a358f9 100644
--- a/atlaselectrophysiology/plot_data.py
+++ b/atlaselectrophysiology/plot_data.py
@@ -407,6 +407,9 @@ def gain2level(gain):
                 h = neuropixel.split_trace_header(h, shank=int(sr.meta.get('NP2.4_shank')))
             else:
                 h = neuropixel.trace_header(sr.major_version, nshank=np.unique(th['shank']).size)
+                idx = np.isin(h['ind'], th['ind'])
+                for k in h.keys():
+                    h[k] = h[k][idx]
 
             s0 = t * sr.fs
             tsel = slice(int(s0), int(s0) + int(1 * sr.fs))