Source code for ibllib.plots

#!/usr/bin/env python
# -*- coding:utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
import scipy

import ibllib.dsp as dsp


[docs]def wiggle(w, fs=1, gain=0.71, color='k', ax=None, fill=True, linewidth=0.5, t0=0, clip=2, **kwargs): """ Matplotlib display of wiggle traces :param w: 2D array (numpy array dimension nsamples, ntraces) :param fs: sampling frequency :param gain: display gain :param color: ('k') color of traces :param ax: (None) matplotlib axes object :param fill: (True) fill variable area above 0 :param t0: (0) timestamp of the first sample :return: None """ nech, ntr = w.shape tscale = np.arange(nech) / fs sf = gain / np.sqrt(dsp.rms(w.flatten())) def insert_zeros(trace): # Insert zero locations in data trace and tt vector based on linear fit # Find zeros zc_idx = np.where(np.diff(np.signbit(trace)))[0] x1 = tscale[zc_idx] x2 = tscale[zc_idx + 1] y1 = trace[zc_idx] y2 = trace[zc_idx + 1] a = (y2 - y1) / (x2 - x1) tt_zero = x1 - y1 / a # split tt and trace tt_split = np.split(tscale, zc_idx + 1) trace_split = np.split(trace, zc_idx + 1) tt_zi = tt_split[0] trace_zi = trace_split[0] # insert zeros in tt and trace for i in range(len(tt_zero)): tt_zi = np.hstack( (tt_zi, np.array([tt_zero[i]]), tt_split[i + 1])) trace_zi = np.hstack( (trace_zi, np.zeros(1), trace_split[i + 1])) return trace_zi, tt_zi if not ax: ax = plt.gca() for ntr in range(ntr): if fill: trace, t_trace = insert_zeros(w[:, ntr] * sf) if clip: trace = np.maximum(np.minimum(trace, clip), -clip) ax.fill_betweenx(t_trace + t0, ntr, trace + ntr, where=trace >= 0, facecolor=color, linewidth=linewidth) wplot = np.minimum(np.maximum(w[:, ntr] * sf, -clip), clip) ax.plot(wplot + ntr, tscale + t0, color, linewidth=linewidth, **kwargs) ax.set_xlim(-1, ntr + 1) ax.set_ylim(tscale[0] + t0, tscale[-1] + t0) ax.set_ylabel('Time (s)') ax.set_xlabel('Trace') ax.invert_yaxis() return ax
[docs]class Density: def __init__(self, w, fs=1, cmap='bone', ax=None, **kwargs): """ Matplotlib display of traces as a density display :param w: 2D array (numpy array dimension nsamples, ntraces) :param fs: sampling frequency (Hz) :param ax: axis to plot in :return: None """ w = w.reshape(w.shape[0], -1) nech, ntr = w.shape tscale = np.array([0, nech - 1]) / fs * 1e3 if ax is None: self.figure, ax = plt.subplots() else: self.figure = ax.get_figure() extent = [-0.5, ntr - 0.5, tscale[1], tscale[0]] self.im = ax.imshow(w, aspect='auto', cmap=cmap, extent=extent, origin='upper', **kwargs) ax.set_ylabel('Time (ms)') ax.set_xlabel('Trace') self.cid_key = self.figure.canvas.mpl_connect('key_press_event', self.on_key_press) self.ax = ax
[docs] def on_key_press(self, event): if event.key == 'ctrl+a': self.im.set_data(self.im.get_array() * np.sqrt(2)) elif event.key == 'ctrl+z': self.im.set_data(self.im.get_array() / np.sqrt(2)) else: return self.figure.canvas.draw()
[docs]class Traces: def __init__(self, w, fs=1, gain=0.71, color='k', ax=None, linewidth=0.5, t0=0, **kwargs): """ Matplotlib display of traces as a density display :param w: 2D array (numpy array dimension nsamples, ntraces) :param fs: sampling frequency (Hz) :param ax: axis to plot in :return: None """ w = w.reshape(w.shape[0], -1) nech, ntr = w.shape tscale = np.arange(nech) / fs * 1e3 sf = gain / dsp.rms(w.flatten()) / 2 if ax is None: self.figure, ax = plt.subplots() else: self.figure = ax.get_figure() self.plot = ax.plot(w * sf + np.arange(ntr), tscale + t0, color, linewidth=linewidth, **kwargs) ax.set_xlim(-1, ntr + 1) ax.set_ylim(tscale[0] + t0, tscale[-1] + t0) ax.set_ylabel('Time (ms)') ax.set_xlabel('Trace') ax.invert_yaxis() self.cid_key = self.figure.canvas.mpl_connect('key_press_event', self.on_key_press) self.ax = ax
[docs] def on_key_press(self, event): if event.key == 'ctrl+a': for i, l in enumerate(self.plot): l.set_xdata((l.get_xdata() - i) * np.sqrt(2) + i) elif event.key == 'ctrl+z': for i, l in enumerate(self.plot): l.set_xdata((l.get_xdata() - i) / np.sqrt(2) + i) else: return self.figure.canvas.draw()
[docs]def squares(tscale, polarity, ax=None, yrange=[-1, 1], **kwargs): """ Matplotlib display of rising and falling fronts in a square-wave pattern :param tscale: time of indices of fronts :param polarity: polarity of front (1: rising, -1:falling) :param ax: matplotlib axes object :return: None """ if not ax: ax = plt.gca() isort = np.argsort(tscale) tscale = tscale[isort] polarity = polarity[isort] f = np.tile(polarity, (2, 1)) t = np.concatenate((tscale, np.r_[tscale[1:], tscale[-1]])).reshape(2, f.shape[1]) ydata = f.transpose().ravel() ydata = (ydata + 1) / 2 * (yrange[1] - yrange[0]) + yrange[0] ax.plot(t.transpose().ravel(), ydata, **kwargs)
[docs]def vertical_lines(x, ymin=0, ymax=1, ax=None, **kwargs): """ From a x vector, draw separate vertical lines at each x location ranging from ymin to ymax :param x: numpy array vector of x values where to display lnes :param ymin: lower end of the lines (scalar) :param ymax: higher end of the lines (scalar) :param ax: (optional) matplotlib axis instance :return: None """ x = np.tile(x, (3, 1)) x[2, :] = np.nan y = np.zeros_like(x) y[0, :] = ymin y[1, :] = ymax y[2, :] = np.nan if not ax: ax = plt.gca() ax.plot(x.T.flatten(), y.T.flatten(), **kwargs)
[docs]def spectrum(w, fs, smooth=None, unwrap=True, axis=0, **kwargs): """ Display spectral density of a signal along a given dimension spectrum(w, fs) :param w: signal :param fs: sampling frequency (Hz) :param smooth: (None) frequency samples to smooth over :param unwrap: (True) unwraps the phase specrum :param axis: axis on which to compute the FFT :param kwargs: plot arguments to be passed to matplotlib :return: matplotlib axes """ axis = 0 smooth = None unwrap = True ns = w.shape[axis] fscale = dsp.fscale(ns, 1 / fs, one_sided=True) W = scipy.fft.rfft(w, axis=axis) amp = 20 * np.log10(np.abs(W)) phi = np.angle(W) if unwrap: phi = np.unwrap(phi) if smooth: nf = np.round(smooth / fscale[1] / 2) * 2 + 1 amp = dsp.smooth.mwa(amp, nf) phi = dsp.smooth.mwa(phi, nf) fig, ax = plt.subplots(2, 1, sharex=True) ax[0].plot(fscale, amp, **kwargs) ax[1].plot(fscale, phi, **kwargs) ax[0].set_title('Spectral Density (dB rel to amplitude.Hz^-0.5)') ax[0].set_ylabel('Amp (dB)') ax[1].set_ylabel('Phase (rad)') ax[1].set_xlabel('Frequency (Hz)') return ax
[docs]def color_cycle(ind=None): """ Gets the matplotlib color-cycle as RGB numpy array of floats between 0 and 1 :return: """ # import matplotlib as mpl # c = np.uint32(np.array([int(c['color'][1:], 16) for c in mpl.rcParams['axes.prop_cycle']])) # c = np.double(np.flip(np.reshape(c.view(np.uint8), (c.size, 4))[:, :3], 1)) / 255 c = np.array([[0.12156863, 0.46666667, 0.70588235], [1., 0.49803922, 0.05490196], [0.17254902, 0.62745098, 0.17254902], [0.83921569, 0.15294118, 0.15686275], [0.58039216, 0.40392157, 0.74117647], [0.54901961, 0.3372549, 0.29411765], [0.89019608, 0.46666667, 0.76078431], [0.49803922, 0.49803922, 0.49803922], [0.7372549, 0.74117647, 0.13333333], [0.09019608, 0.74509804, 0.81176471]]) if ind is None: return c else: return tuple(c[ind % c.shape[0], :])
if __name__ == "__main__": w = np.random.rand(500, 40) - 0.5 wiggle(w, fs=30000) Traces(w, fs=30000, color='r')