207 lines
8.0 KiB
Python
207 lines
8.0 KiB
Python
import string
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.transforms import BboxTransformTo
|
|
from itertools import product
|
|
|
|
def prepare_fig(nrows, ncols, width=8, height=None, rheight=2, unit=1/2.54,
|
|
left=0.01, right=0.95, bottom=0.01, top=0.95,
|
|
wspace=0.4, hspace=0.4):
|
|
if height is None:
|
|
height = rheight * nrows
|
|
fig = plt.figure(figsize=(width * unit, height * unit))
|
|
grid = fig.add_gridspec(nrows=nrows, ncols=ncols, wspace=wspace, hspace=hspace,
|
|
left=left, right=right, top=top, bottom=bottom)
|
|
axes = np.zeros((nrows, ncols), dtype=object)
|
|
for i, j in product(range(nrows), range(ncols)):
|
|
axes[i, j] = fig.add_subplot(grid[i, j])
|
|
axes[i, j].set_facecolor('none')
|
|
return fig, axes
|
|
|
|
def hide_ticks(ax, side='bottom', ticks=True):
|
|
axis = 'x' if side in ['top', 'bottom'] else 'y'
|
|
params = {side: ticks, 'label' + side: False}
|
|
ax.tick_params(axis=axis, which='both', **params)
|
|
return None
|
|
|
|
def hide_axis(ax, side='bottom'):
|
|
ax.spines[side].set_visible(False)
|
|
params = {side: False, 'label' + side: False}
|
|
ax.tick_params(axis='x' if side in ['top', 'bottom'] else 'y',
|
|
which='both', **params)
|
|
return None
|
|
|
|
def letter_subplot(artist, label, x=None, y=None, xref=None, yref=None, ref=None,
|
|
ha='left', va='bottom', fontsize=16, fontweight='bold', **kwargs):
|
|
trans_artist = BboxTransformTo(artist.bbox)
|
|
if x is None or y is None:
|
|
transform = BboxTransformTo(ref.bbox) + trans_artist.inverted()
|
|
if x is None:
|
|
x = transform.transform([xref, 0])[0]
|
|
if y is None:
|
|
y = transform.transform([0, yref])[1]
|
|
artist.text(x, y, label, transform=trans_artist, ha=ha, va=va,
|
|
fontsize=fontsize, fontweight=fontweight, **kwargs)
|
|
return None
|
|
|
|
def letter_subplots(artists, labels=None, x=None, y=None, xref=None, yref=None, ref=None,
|
|
ha='left', va='bottom', fontsize=16, fontweight='bold', **kwargs):
|
|
if labels is None:
|
|
labels = string.ascii_lowercase
|
|
for artist, label in zip(artists, labels):
|
|
letter_subplot(artist, label, x, y, xref, yref, ref=ref, ha=ha, va=va,
|
|
fontsize=fontsize, fontweight=fontweight, **kwargs)
|
|
return None
|
|
|
|
def xlimits(time, ax=None, minval=None, maxval=None, pad=0.05):
|
|
limits = [minval, maxval]
|
|
if minval is None:
|
|
limits[0] = time[0]
|
|
if maxval is None:
|
|
limits[1] = time[-1]
|
|
span = limits[1] - limits[0]
|
|
if pad and minval is None:
|
|
limits[0] -= span * pad
|
|
if pad and maxval is None:
|
|
limits[1] += span * pad
|
|
if ax is not None:
|
|
return ax.set_xlim(limits)
|
|
return limits
|
|
|
|
def ylimits(signal, ax=None, minval=None, maxval=None, pad=0.05):
|
|
limits = [minval, maxval]
|
|
if minval is None:
|
|
limits[0] = signal.min()
|
|
if maxval is None:
|
|
limits[1] = signal.max()
|
|
span = limits[1] - limits[0]
|
|
if pad and minval is None:
|
|
limits[0] -= span * pad
|
|
if pad and maxval is None:
|
|
limits[1] += span * pad
|
|
if ax is not None:
|
|
return ax.set_ylim(limits)
|
|
return limits
|
|
|
|
def xlabel(ax, label, x=None, y=-0.1, fontsize=20, transform=None, **kwargs):
|
|
ax.set_xlabel(label, fontsize=fontsize, **kwargs)
|
|
if x is None:
|
|
x = 0.5
|
|
if transform is not None:
|
|
x = (ax.transAxes + transform.inverted()).transform((x, 0))[0]
|
|
ax.xaxis.set_label_coords(x, y, transform=transform)
|
|
return None
|
|
|
|
def ylabel(ax, label, x=-0.2, y=None, fontsize=20, transform=None, **kwargs):
|
|
ax.set_ylabel(label, fontsize=fontsize, **kwargs)
|
|
if y is None:
|
|
y = 0.5
|
|
if transform is not None:
|
|
y = (ax.transAxes + transform.inverted()).transform((0, y))[1]
|
|
ax.yaxis.set_label_coords(x, y, transform=transform)
|
|
return None
|
|
|
|
def super_xlabel(label, fig, left_ax, right_ax, y=0.005,
|
|
left_fig=None, right_fig=None, **kwargs):
|
|
left_x = left_ax.get_position().x0
|
|
right_x = right_ax.get_position().x1
|
|
if left_fig is not None or right_fig is not None:
|
|
trans_fig = BboxTransformTo(fig.bbox)
|
|
if left_fig is not None:
|
|
transform = BboxTransformTo(left_fig.bbox) + trans_fig.inverted()
|
|
left_x = transform.transform((left_x, 0))[0]
|
|
if right_fig is not None:
|
|
transform = BboxTransformTo(right_fig.bbox) + trans_fig.inverted()
|
|
right_x = transform.transform((right_x, 0))[0]
|
|
return fig.supxlabel(label, x=(left_x + right_x) / 2, y=y, **kwargs)
|
|
|
|
def super_ylabel(label, fig, low_ax, high_ax, x=0.005,
|
|
high_fig=None, low_fig=None, **kwargs):
|
|
low_y = high_ax.get_position().y0
|
|
high_y = low_ax.get_position().y1
|
|
if low_fig is not None or high_fig is not None:
|
|
trans_fig = BboxTransformTo(fig.bbox)
|
|
if low_fig is not None:
|
|
transform = BboxTransformTo(low_fig.bbox) + trans_fig.inverted()
|
|
low_y = transform.transform((0, low_y))[1]
|
|
if high_fig is not None:
|
|
transform = BboxTransformTo(high_fig.bbox) + trans_fig.inverted()
|
|
high_y = transform.transform((0, high_y))[1]
|
|
return fig.supylabel(label, x=x, y=(low_y + high_y) / 2, **kwargs)
|
|
|
|
def plot_line(ax, time, signal, ymin=None, ymax=None, xmin=None, xmax=None,
|
|
xpad=None, ypad=0.05, yloc=None, xloc=None, **kwargs):
|
|
handles = ax.plot(time, signal, **kwargs)
|
|
xlimits(time, ax=ax, minval=xmin, maxval=xmax, pad=xpad)
|
|
ylimits(signal, ax=ax, minval=ymin, maxval=ymax, pad=ypad)
|
|
if xloc is not None:
|
|
ax.xaxis.set_major_locator(plt.MultipleLocator(xloc))
|
|
if yloc is not None:
|
|
ax.yaxis.set_major_locator(plt.MultipleLocator(yloc))
|
|
return handles
|
|
|
|
def plot_barcode(ax, time, binary, offset=0.5, xmin=None, xmax=None, **kwargs):
|
|
lower, upper, handles = 0, 1, []
|
|
for i in range(binary.shape[1]):
|
|
h = ax.fill_between(time, lower, upper, where=binary[:, i], **kwargs)
|
|
handles.append(h)
|
|
if i < binary.shape[1] - 1:
|
|
lower += offset + 1
|
|
upper += offset + 1
|
|
xlimits(time, ax=ax, minval=xmin, maxval=xmax, pad=0)
|
|
ax.set_ylim(0, upper)
|
|
hide_axis(ax, 'bottom')
|
|
hide_axis(ax, 'left')
|
|
return handles
|
|
|
|
def indicate_zoom(fig, high_ax, low_ax, zoom_abs, **kwargs):
|
|
y0 = low_ax.get_position().y0
|
|
y1 = high_ax.get_position().y1
|
|
transform = low_ax.transData + fig.transFigure.inverted()
|
|
x0 = transform.transform((zoom_abs[0], 0))[0]
|
|
x1 = transform.transform((zoom_abs[1], 0))[0]
|
|
fig.add_artist(plt.Rectangle((x0, y0), x1 - x0, y1 - y0,
|
|
transform=fig.transFigure, **kwargs))
|
|
return None
|
|
|
|
def assign_colors(handles, types, colors):
|
|
for handle, type_id in zip(handles, types):
|
|
handle.set_color(colors[str(int(type_id))])
|
|
return None
|
|
|
|
def reorder_traces(handles, signal, zlow=2, zhigh=2.5):
|
|
inds = np.argsort(signal.std(axis=0))
|
|
zorders = np.linspace(zlow, zhigh, len(inds))[::-1]
|
|
for ind, z in zip(inds, zorders):
|
|
handles[ind].set_zorder(z)
|
|
return None
|
|
|
|
def strip_zeros(num, right_digits=5):
|
|
if isinstance(num, int):
|
|
return num
|
|
num = f'{num:.{right_digits}f}'
|
|
left, right = num.split('.')
|
|
right = right.rstrip('0')
|
|
if right:
|
|
return f'{left}.{right}'
|
|
return left
|
|
|
|
def time_bar(ax, dur, y0=0.9, y1=0.95, xshift=0.5, parent=None, transform=None, **kwargs):
|
|
t_lims = ax.get_xlim()
|
|
span = t_lims[1] - t_lims[0]
|
|
if parent is not None or transform is not None:
|
|
if transform is None:
|
|
transform = BboxTransformTo(parent.bbox)
|
|
kwargs['transform'] = transform
|
|
transform = ax.transData + transform.inverted()
|
|
x0 = transform.transform((t_lims[0], 0))[0]
|
|
x1 = transform.transform((t_lims[0] + dur, 0))[0]
|
|
dur = x1 - x0
|
|
span = 1
|
|
elif parent is None:
|
|
parent = ax
|
|
x0 = (span - dur) * xshift
|
|
x1 = x0 + dur
|
|
parent.add_artist(plt.Rectangle((x0, y0), dur, y1 - y0, **kwargs))
|
|
return None
|