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: trans_ref = BboxTransformTo(ref.bbox) transform = trans_ref + 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, high_ax, low_ax, y=0.005, **kwargs): x = (low_ax.get_position().x0 + high_ax.get_position().x1) / 2 fig.supxlabel(label, x=x, y=y, **kwargs) return None def super_ylabel(label, fig, high_ax, low_ax, x=0.005, **kwargs): y = (low_ax.get_position().y0 + high_ax.get_position().y1) / 2 fig.supylabel(label, x=x, y=y, **kwargs) return None 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): t0, t1 = ax.get_xlim() offset = (t1 - t0 - dur) * xshift x0 = t0 + offset x1 = x0 + dur if parent is None: parent = ax if transform is None: transform = BboxTransformTo(parent.bbox) if transform is not ax.transData: trans = ax.transData + transform.inverted() x0 = trans.transform((x0, 0))[0] x1 = trans.transform((x1, 0))[0] parent.add_artist(plt.Rectangle((x0, y0), x1 - x0, y1 - y0, transform=transform, **kwargs)) return None