Dendrograms and heatmaps

heatmap.py

import pandas as pd
from pandas import DataFrame as DF
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from IPython.display import display
import IPython.html.widgets as widgets
from scipy.cluster.hierarchy import \
    linkage,\
    leaves_list,\
    dendrogram
from toolz import \
    pipe,\
    thread_first,\
    curry,\
    frequencies

def plot_matrix_tree(matrix,xlink,ylink,xlabels,ylabels,color):
    plt.figure(figsize=(20,10))
    plt.subplots_adjust(hspace=0,wspace=0)
    gs = gridspec.GridSpec(2, 2,height_ratios=[1,5],width_ratios=[5,1]) 
    
    def subplot_dendrogram(subplot,link,orientation):
        plt.subplot(subplot)
        dendrogram(link, orientation = orientation, link_color_func = lambda x: 'k')
        plt.axis('off')
        if orientation == 'left':
            plt.gca().invert_yaxis() # Need to flip dendrogram, since the default seems incorrect. 

    subplot_dendrogram(gs[0],xlink,'top')
    subplot_dendrogram(gs[3],ylink,'left')
    
    plt.subplot(gs[2])
    plt.imshow(matrix,interpolation='nearest',cmap=color,aspect='auto');
    [plt.gca().spines[loc].set_visible(False) for loc in ['top','bottom','left','right']]
    
    plt.xticks(range(len(xlabels)),xlabels,rotation=30,ha='right')
    plt.yticks(range(len(ylabels)),ylabels)

def matrix_tree(data,color):
    normed_data = data.values

    condition_link = linkage(normed_data)
    feature_link = linkage(normed_data.T)

    condition_order = leaves_list(condition_link)
    feature_order = leaves_list(feature_link)

    conditions = data.index.values[condition_order]
    features = data.columns.values[feature_order]

    color_matrix = normed_data.T[feature_order,:][:,condition_order]
    
    plot_matrix_tree(color_matrix,
                     condition_link,
                     feature_link,
                     conditions,
                     features,
                     color)