Box outside my y axis in matplotlib/seaborn-CodePudding

I"m trying to add a descriptive box outside my y axis on my heatmap in matplotlib/seaborn. I tried an additional subplot and hspan, but they weren't working for me. Help please!

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

col1, col2 = np.random.rand(31), np.random.rand(31)

df = pd.DataFrame({'first': col1,
                   'second': col2}, index = range(70,101))

fig, ax = plt.subplots(1,figsize = (5,10))
sns.heatmap(data=df, annot=True, cmap='RdYlGn', cbar=False)
ax.invert_yaxis()
ax.tick_params('y', labelrotation=0)

Here is a picture of the box beside the y axis for reference.

CodePudding user response：

You could create two subplots, put the heatmap in the right subplot and the text in the left one.

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

df = pd.DataFrame({'first': np.random.rand(31),
                   'second': np.random.rand(31)}, index=range(70, 101))

sns.set()
fig, (ax0, ax) = plt.subplots(ncols=2, figsize=(5, 10), gridspec_kw={'width_ratios': [1, 3]})
sns.heatmap(data=df, annot=True, cmap='RdYlGn', cbar=False, ax=ax)
ax.invert_yaxis()
ax.tick_params('y', labelrotation=0)
y_break = 21 / 31
ax0.axhspan(0, y_break, color='turquoise')
ax0.axhspan(y_break, 1, color='tomato')
ax0.axis('off')
ax0.set_ylim(0, 1)
ax0.text(0.5, y_break / 2, 'good', fontsize=24, rotation=90, ha='center', va='center')
ax0.text(0.5, (1   y_break) / 2, 'bad', fontsize=24, rotation=90, ha='center', va='center')
plt.tight_layout()
plt.show()