Heatmap functions for Pandas dataframes.

pandas heatmaps

"""Plots a Pandas dataframe as a heatmap"""
import matplotlib as mpl
import matplotlib.pyplot as plt

def heatmap(df,
            edgecolors='w',
            cmap=mpl.cm.RdBu,
            log=False,vmin=0,vmax=500):    
    width = len(df.columns)/4
    height = len(df.index)/4
    
    fig, ax = plt.subplots(figsize=(width,height))
      
    heatmap = ax.pcolor(df,
                        edgecolors=edgecolors,  # put white lines between squares in heatmap
                        cmap=cmap,
                        vmin=vmin, # defaults to 0
                        vmax=vmax, # defaults to 500
                        norm=mpl.colors.LogNorm() if log else None)
    
    ax.autoscale(tight=True)  # get rid of whitespace in margins of heatmap
    ax.set_aspect('equal')  # ensure heatmap cells are square
    ax.xaxis.set_ticks_position('top')  # put column labels at the top
    ax.tick_params(bottom='off', top='off', left='off', right='off')  # turn off ticks
    
    plt.yticks(np.arange(len(df.index)) + 0.5, df.index)
    plt.xticks(np.arange(len(df.columns)) + 0.5, df.columns, rotation=90)
    
    # ugliness from http://matplotlib.org/users/tight_layout_guide.html
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", "3%", pad="1%")
    plt.colorbar(heatmap, cax=cax)
    return fig


"""Binary Heatmap"""

import matplotlib as mpl
import matplotlib.pyplot as plt
 
def binary_heatmap(df):
    df = df[::-1]  # reverse df to put first row at top (last row at origin)

    width = len(df.columns)/5
    height = len(df.index)/5
    
    fig, ax = plt.subplots(figsize=(width,height))
      
    heatmap = ax.pcolor(df,
                        edgecolors='k',  # put black lines between squares in heatmap
                        cmap=mpl.cm.binary)  # black/white colomarp
    
    ax.autoscale(tight=True)  # get rid of whitespace in margins of heatmap
    ax.set_aspect('equal')  # ensure heatmap cells are square
    ax.xaxis.set_ticks_position('top')  # put column labels at the top
    ax.tick_params(bottom='off', top='off', left='off', right='off')  # turn off ticks
    
    plt.yticks(np.arange(len(df.index)) + 0.5, df.index)
    plt.xticks(np.arange(len(df.columns)) + 0.5, df.columns, rotation=90)

    
    plt.tight_layout()