Heatmap with multi-color y-axis and correspondend colorbar

Question

I want to create a heatmap with seaborn, similar to this (with the following code):

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np

# Create data
df = pd.DataFrame(np.random.random((5,5)), columns=["a","b","c","d","e"])

# Default heatmap
ax = sns.heatmap(df)
plt.show()

I'd also like to add a new variable (lets say new_var = pd.DataFrame(np.random.random((5,1)), columns=["new variable"])), such as that the values (and possibly the spine and ticks as well) of the y-axis are colored according to the new variable and a second color bar plotted in the same plot to represent the colors of the y-axis values. How can I do that?

Answer 1

This uses the new values to color the y-ticks and the y-tick labels and adds the associated colorbar.

import matplotlib.pyplot as plt
import matplotlib
import seaborn as sns
import pandas as pd
import numpy as np

# Create data
df = pd.DataFrame(np.random.random((5,5)), columns=["a","b","c","d","e"])

# Default heatmap
ax = sns.heatmap(df)

new_var = pd.DataFrame(np.random.random((5,1)), columns=["new variable"])

# Create the colorbar for y-ticks and labels
norm = plt.Normalize(new_var.min(), new_var.max())
cmap = matplotlib.cm.get_cmap('turbo')

yticks_locations = ax.get_yticks()
yticks_labels = df.index.values
#hide original ticks
ax.tick_params(axis='y', left=False)
ax.set_yticklabels([])

for var, ytick_loc, ytick_label in zip(new_var.values, yticks_locations, yticks_labels):
    color = cmap(norm(float(var)))
    ax.annotate(ytick_label, xy=(1, ytick_loc), xycoords='data', xytext=(-0.4, ytick_loc),
    arrowprops=dict(arrowstyle="-", color=color, lw=1), zorder=0, rotation=90, color=color)

# Add colorbar for y-tick colors
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
cb = ax.figure.colorbar(sm)
# Match the seaborn style
cb.outline.set_visible(False)