How to make a heatmap using seaborn with colored lines as secondary x and y axes

Question

Dear Stackoverflow community;

recently I have been playing around with the Seaborn heatmap function and I was wondering, whether it is possible to use color lines as secondary axes? To be more specific, I would like to reproduce the following plot, which I know how to do in mathematica, in python:

I have written the following, much simplified example

vals = np.array(['A','B','C','D'])
corrs = np.array([[0, 1, 0.5, 3],[1,0,2,0],[0.5, 2,0, 3],[3,0,3,0]])
ax = sns.heatmap(corrs,xticklabels=vals, yticklabels=vals)
plt.show()

and then tinkered with jointplot trying to get a colored line to appear under the axes ticks, but to no avail. Any help would be very much appreciated!

Answer 1

You can get something similar via a clustermap, without dendrograms. Your test data don't explain how you want to do the coloring. (Also, correlation coefficients outside the range -1, 1 are a bit strange.)

The code below uses red, green, magenta and blues to indicate the rows and columns.

import seaborn as sns
import pandas as pd
import numpy as np

vals = np.array(['A', 'B', 'C', 'D'])
corrs = np.array([[0, 1, 0.5, 3], [1, 0, 2, 0], [0.5, 2, 0, 3], [3, 0, 3, 0]])
corrs_df = pd.DataFrame(corrs, columns=vals, index=vals)

g = sns.clustermap(corrs_df,
                   row_cluster=False, col_cluster=False,
                   row_colors=[*'rgmb'], col_colors=[*'rgmb'],
                   dendrogram_ratio=(.1, .02),
                   cbar_pos=(0.01, .2, .03, .5),
                   figsize=(10, 6))