plot matrix with dendogram

dendogram.py

import numpy as np
import copy
import random
import pylab
import scipy.cluster.hierarchy as sch


def create_linkage_matrix(num_tasks):
    """
    create stack of matrices for
    each inner node of binary tree

    """

    num_levels = int(np.log2(num_tasks))

    stack = []

    cumulative = 0

    for i in range(num_levels):
        block_size = 2 ** i
        num_blocks = int(num_tasks / block_size)

        for j in range(0, num_blocks, 2):
            row = [cumulative+j, cumulative+j+1, i+1, block_size*2]
            print row
            stack.append(row)

        cumulative += num_blocks
    stack = np.array(stack, dtype=np.float64)

    return stack


def plot_dendrogram(data, Y): 
    """ 
    http://stackoverflow.com/questions/2982929/plotting-results-of-hierarchical-clustering-ontop-of-a-matrix-of-data-in-python
    """

    D = copy.copy(data)
    D /= D[0,0]


    # Compute and plot first dendrogram.
    fig = pylab.figure(figsize=(8,8))
    ax1 = fig.add_axes([0.09,0.1,0.2,0.6])
    ax1.set_axis_off()
    Z1 = sch.dendrogram(Y, orientation='right', link_color_func=lambda k: "black")
    ax1.set_xticks([])
    ax1.set_yticks([])

    # Compute and plot second dendrogram.
    ax2 = fig.add_axes([0.3,0.71,0.6,0.2])
    ax2.set_axis_off()
    Z2 = sch.dendrogram(Y, link_color_func=lambda k: "black")
    ax2.set_xticks([])
    ax2.set_yticks([])

    # Plot distance matrix.
    axmatrix = fig.add_axes([0.3,0.1,0.6,0.6])
    idx1 = Z1['leaves']
    idx2 = Z2['leaves']
    D = D[idx1,:]
    D = D[:,idx2]
    im = axmatrix.matshow(D, aspect='auto', origin='lower') #, cmap=pylab.cm.YlGnBu)
    axmatrix.set_xticks([])
    axmatrix.set_yticks([])

    # Plot colorbar.
    axcolor = fig.add_axes([0.91,0.1,0.02,0.6])
    pylab.colorbar(im, cax=axcolor)
    fig.show()

    pylab.show()



def main():
    Y = create_linkage_matrix(32)
    plot_dendrogram(sim, Y)

    pylab.figure()
    Z1 = sch.dendrogram(Y, orientation='right')
    pylab.imshow(sim, interpolation="nearest")

    pylab.show()