jupyter_util.py

import matplotlib.pyplot as plt
import numpy as np
from IPython.display import HTML, display
from io import BytesIO
import base64
import os
import contextlib

def draw_eta_lines(ax, eta_range = (-3, 3), n=None, s=None, fmt="%.2f", **kwargs):
    assert (n is None and s is not None) or (n is not None and s is None)
    eta_min, eta_max = eta_range
    if n is not None:
        eta_vals = np.linspace(eta_min, eta_max, n)
    else:
        eta_vals = np.arange(eta_min, eta_max+s, s)
    thetas = 2*np.arctan(np.exp(-eta_vals))
    zmin, zmax = ax.get_xlim()
    rmin, rmax = ax.get_ylim()
    z_vals = np.array([zmin if theta > math.pi/2. else zmax for theta in thetas])
    r_vals = z_vals * np.tan(thetas)
    
    for eta, theta, z_out, r_out in zip(eta_vals, thetas, z_vals, r_vals):
        
        z_in = 0
        r_in = 0
        
        if eta == 0.:
            ax.text(0, rmax, s=r"$\eta=0$", ha="center", va="bottom")
            ax.plot([0, 0], [max(r_in, rmin), rmax], **kwargs)
            continue
        
        if r_out > rmax:
            # re-calc z based on r
            z_out = rmax / np.tan(theta)
            r_out = rmax
        
        if r_out < rmin:
            # would not show anyway
            continue
           
        if r_in < rmin:
            # re-calc z_in
            z_in = rmin / np.tan(theta)
            r_in = rmin
            
        ax.plot([z_in, z_out], [r_in, r_out], **kwargs)
        ax.text(z_out, r_out, s=r"$\eta=%s$"%(fmt%eta), ha="center", va="bottom") 
        
        
def embed_figure(fig, filename):
    
    _, ext = os.path.splitext(filename)
    
    if ext == ".pdf":
        mime = "application/pdf"
        fmt = "pdf"
    elif ext == ".png":
        mime = "image/png"
        fmt = "png"
        
    d = BytesIO()    
    fig.savefig(d, format=fmt)
    data = d.getvalue()
    b64 = base64.b64encode(data)
    payload = b64.decode()
    html = f'''
    <a download="{filename}" 
       href="data:{mime};base64,{payload}" 
       target="_blank">Download {filename}</a>'''.strip()
    display(HTML(html))
    
class FlowLayout(object):
    ''' A class / object to display plots in a horizontal / flow layout below a cell '''
    def __init__(self):
        # string buffer for the HTML: initially some CSS; images to be appended
        self.sHtml =  """
        <style>
        .floating-box {
        display: inline-block;
        background:#fff;
        margin: 10px;
        border: 0;  
        }
        </style>
        """
        
    @contextlib.contextmanager
    def subplots(self):
        fig, ax = plt.subplots()
        yield fig, ax
        self.add_plot(ax)

    def add_plot(self, oAxes):
        ''' Saves a PNG representation of a Matplotlib Axes object '''
        Bio= BytesIO() # bytes buffer for the plot
        fig = oAxes.get_figure()
        fig.canvas.print_png(Bio) # make a png of the plot in the buffer

        # encode the bytes as string using base 64 
        sB64Img = base64.b64encode(Bio.getvalue()).decode()
        self.sHtml+= (
            '<div class="floating-box">'+ 
            '<img src="data:image/png;base64,{}\n">'.format(sB64Img)+
            '</div>')
    
        plt.close()
        
    def add_plots(self, axs):
        for ax in axs:
            self.add_plot(ax)
            plt.close()

    def display(self):
        ''' Final step - display the accumulated HTML '''
        display(HTML(self.sHtml))
        
    def __enter__(self, *args):
        return self
        
    def __exit__(self, *args):
        self.display()
        
        
def plot_hist(ax, hist, bin_edges, yerr=np.sqrt, **kwargs):
    extra = {"linestyle":'none', "capsize": 0}
    extra.update(kwargs)
    ax.errorbar(bin_edges[1:]-(bin_edges[1:]-bin_edges[0:-1])/2, 
            hist, yerr=yerr(hist), xerr=(bin_edges[0]-bin_edges[1])/2,
            **extra)