Lots of stuff. Syncing to home.

python/plot_functions.py

@@ -2,21 +2,6 @@ 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'
@@ -31,6 +16,20 @@ def hide_axis(ax, side='bottom'):
                    which='both', **params)
     return None
 
+def title_subplot(artist, title, x=0.5, y=1.0, xref=None, yref=None, ref=None,
+                  ha='center', va='bottom', fontsize=16, fontweight='normal', **kwargs):
+
+    trans_artist = BboxTransformTo(artist.bbox)
+    if xref is not None or yref is not None:
+        transform = BboxTransformTo(ref.bbox) + trans_artist.inverted()
+        if xref is not None:
+            x = transform.transform((xref, 0))[0]
+        if yref is not None:
+            y = transform.transform((0, yref))[1]
+    artist.text(x, y, title, transform=trans_artist, ha=ha, va=va,
+                fontsize=fontsize, fontweight=fontweight, **kwargs)
+    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)
@@ -141,6 +140,8 @@ def plot_line(ax, time, signal, ymin=None, ymax=None, xmin=None, xmax=None,
     return handles
 
 def plot_barcode(ax, time, binary, offset=0.5, xmin=None, xmax=None, **kwargs):
+    if binary.ndim == 1:
+        binary = binary[:, None]
     lower, upper, handles = 0, 1, []
     for i in range(binary.shape[1]):
         h = ax.fill_between(time, lower, upper, where=binary[:, i], **kwargs)
@@ -186,21 +187,63 @@ def strip_zeros(num, right_digits=5):
         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:
+def time_bar(ax, dur, y0=0.9, y1=0.95, xshift=0.5, parent=None, **kwargs):
+    if parent is None:
         parent = ax
-    x0 = (span - dur) * xshift
-    x1 = x0 + dur
+    trans_parent = BboxTransformTo(parent.bbox)
+    kwargs['transform'] = trans_parent
+    transform = ax.transData + trans_parent.inverted()
+    t0 = ax.get_xlim()[0]
+    x0 = transform.transform((t0, 0))[0]
+    x1 = transform.transform((t0 + dur, 0))[0]
+    dur = x1 - x0
+    x0 = (1 - dur) * xshift
     parent.add_artist(plt.Rectangle((x0, y0), dur, y1 - y0, **kwargs))
     return None
+
+def zoom_inset(ax, inset, handle, x0=None, x1=None, y0=None, y1=None, ref='x',
+               transform = None, 
+               low_left=False, up_left=False, low_right=False, up_right=False,
+               props=['c', 'lw', 'ls', 'zorder', 'alpha'], **kwargs):
+    if not kwargs:
+        kwargs = dict(edgecolor='k', alpha=1, lw=2)
+
+    if transform is not None:
+        transform = transform + ax.transData.inverted()
+    xlims = ax.get_xlim()
+    ylims = ax.get_ylim()  
+    if x0 is None:
+        x0 = xlims[0]
+    elif transform is not None:
+        x0 = transform.transform((x0, 0))[0]
+    if x1 is None:
+        x1 = xlims[1]
+    elif transform is not None:
+        x1 = transform.transform((x1, 0))[0]
+    if y0 is None:
+        y0 = ylims[0]
+    elif transform is not None:
+        y0 = transform.transform((0, y0))[1]
+    if y1 is None:
+        y1 = ylims[1]
+    elif transform is not None:
+        y1 = transform.transform((0, y1))[1]
+    inset.set_xlim(x0, x1)
+    inset.set_ylim(y0, y1)
+
+    x = handle.get_xdata()
+    y = handle.get_ydata()
+    if ref == 'x':
+        zoom_inds = (x >= x0) & (x <= x1)
+    elif ref == 'y':
+        zoom_inds = (y >= y0) & (y <= y1)
+    x = x[zoom_inds]
+    y = y[zoom_inds]
+
+    inset_handle = inset.plot(x, y)[0]
+    inset_handle.set(**{prop: plt.getp(handle, prop) for prop in props})
+    elements = ax.indicate_inset_zoom(inset, **kwargs)
+
+    visibility = low_left, up_left, low_right, up_right
+    [l.set_visible(v) for l, v in zip(elements.connectors, visibility)]
+    return inset_handle, elements.rectangle, elements.connectors