heat-maps-dendrogram.py

import scipy
import matplotlib.pyplot as plt
import scipy.cluster.hierarchy as sch
# %matplotlib inline
# uncomment the above line if using Jupyter notebook

# creating dummy variable for plotting
# 'D' is a 2-d array that has our data
# for demonstartion purpose we are creating random variables
x = scipy.rand(40)
D = scipy.zeros([40,40])
for i in range(40):
    for j in range(40):
        D[i, j] = abs(x[i] - x[j])
        
# Dendrogram that comes to the left
fig = plt.figure(figsize=(8,8))
# Add an axes at position rect [left, bottom, width, height]
ax1 = fig.add_axes([0.09, 0.1, 0.2, 0.6])
Y = sch.linkage(D, method='centroid')
# orientation='left' is reponsible for making the 
# dendrogram appear to the left
Z1 = sch.dendrogram(Y, orientation='left')
ax1.set_xticks([])
ax1.set_yticks([])

# top side dendogram
ax2 = fig.add_axes([0.3, 0.71, 0.6, 0.2])
Y = sch.linkage(D, method='single')
Z2 = sch.dendrogram(Y)
ax2.set_xticks([])
ax2.set_yticks([])

# main heat-map
axmatrix = fig.add_axes([0.3, 0.1, 0.6, 0.6])
idx1 = Z1['leaves']
idx2 = Z2['leaves']
D = D[idx1, :]
D = D[:, idx2]
# the actual heat-map
im = axmatrix.matshow(D, aspect='auto', origin='lower', cmap="YlGnBu")
axmatrix.set_xticks([])
axmatrix.set_yticks([])

# xticks to the right (x-axis)
axmatrix.set_xticks(range(40))
axmatrix.set_xticklabels(idx1, minor=False)
axmatrix.xaxis.set_label_position('bottom')
axmatrix.xaxis.tick_bottom()

plt.xticks(rotation=-90, fontsize=8)

# xticks to the right (y-axis)
axmatrix.set_yticks(range(40))
axmatrix.set_yticklabels(idx2, minor=False)
axmatrix.yaxis.set_label_position('right')
axmatrix.yaxis.tick_right()

# to add the color bar
# axcolor = fig.add_axes([0.94, 0.1, 0.02, 0.6])

plt.show()