def heatmap(w, vmin=None, vmax=None, diverge_color=False,
            ncol=1,
            plot_name=None, vocab=["A", "C", "G", "T"], figsize=(6, 2)):
    """Plot a heatmap from weight matrix w

    vmin, vmax = z axis range
    diverge_color = Should we use diverging colors?
    plot_name = plot_title
    vocab = vocabulary (corresponds to the first axis)
    """
    # Generate y and x values from the dimension lengths
    assert len(vocab) == w.shape[0]
    plt_y = np.arange(w.shape[0] + 1) + 0.5
    plt_x = np.arange(w.shape[1] + 1) - 0.5
    z_min = w.min()
    z_max = w.max()

    if vmin is None:
        vmin = z_min
    if vmax is None:
        vmax = z_max

    if diverge_color:
        color_map = plt.cm.RdBu
    else:
        color_map = plt.cm.Blues

    fig = plt.figure(figsize=figsize)
    # multiple axis
    if len(w.shape) == 3:
        #
        n_plots = w.shape[2]
        nrow = math.ceil(n_plots / ncol)
    else:
        n_plots = 1
        nrow = 1
        ncol = 1

    for i in range(n_plots):
        if len(w.shape) == 3:
            w_cur = w[:, :, i]
        else:
            w_cur = w
        ax = plt.subplot(nrow, ncol, i + 1)
        plt.tight_layout()
        im = ax.pcolormesh(plt_x, plt_y, w_cur, cmap=color_map,
                           vmin=vmin, vmax=vmax, edgecolors="white")
        ax.grid(False)
        ax.set_yticklabels([""] + vocab, minor=False)
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        ax.set_xticks(np.arange(w_cur.shape[1] + 1))
        ax.set_xlim(plt_x.min(), plt_x.max())
        ax.set_ylim(plt_y.min(), plt_y.max())

        # nice scale location:
        # http://stackoverflow.com/questions/18195758/set-matplotlib-colorbar-size-to-match-graph
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax)

        if plot_name is not None:
            if n_plots > 0:
                pln = plot_name + " {0}".format(i)
            else:
                pln = plot_name
            ax.set_title(pln)
        ax.set_aspect('equal')
    return fig
Plot a heatmap from weight matrix w

    vmin, vmax = z axis range
    diverge_color = Should we use diverging colors?
    plot_name = plot_title
    vocab = vocabulary (corresponds to the first axis)