Matplotlib Correlation Matrix Heatmap With Grouped Colors As Labels

I have a correlation matrix hat I am trying to visualize with matplotlib. I can create a heatmap style figure just fine, but I am running into problems with how I want the labels.

Solution 1:

You may create auxilary axes next to the plot and plot colored bar plots to them. Turning the axes spines off lets those bars look like labelboxes.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

# COREELATION MATRIX TEST #
corr = 2*np.random.rand(150,150)-1# labels [start,end]
labels = np.array([[0,15],[16,36],[37,82],[83,111],[112,149]])
colors = ["crimson", "limegreen","gold","orchid","turquoise"]

fig, ax = plt.subplots()

im = ax.imshow(corr, cmap="Blues")

ax.set_title('THIS IS MY TITLE')
fig.colorbar(im, ticks=[-1,-0.8,-0.6,-0.4,-0.2,0.0,0.2,0.4,0.6,0.8,1.0])

# create axes next to plot
divider = make_axes_locatable(ax)
axb = divider.append_axes("bottom", "10%", pad=0.06, sharex=ax)
axl = divider.append_axes("left", "10%", pad=0.06, sharey=ax)
axb.invert_yaxis()
axl.invert_xaxis()
axb.axis("off")
axl.axis("off")


# plot colored bar plots to the axes
barkw = dict( color=colors, linewidth=0.72, ec="k", clip_on=False, align='edge',)
axb.bar(labels[:,0],np.ones(len(labels)), 
        width=np.diff(labels, axis=1).flatten(), **barkw)
axl.barh(labels[:,0],np.ones(len(labels)), 
         height=np.diff(labels, axis=1).flatten(), **barkw)

# set margins to zero again
ax.margins(0)
ax.tick_params(axis="both", bottom=0, left=0, labelbottom=0,labelleft=0)
# Label the boxes
textkw = dict(ha="center", va="center", fontsize="small")
for k,l in labels:
    axb.text((k+l)/2.,0.5, "{}-{}".format(k,l), **textkw)
    axl.text(0.5,(k+l)/2., "{}-{}".format(k,l), rotation=-90,**textkw)

plt.show()