Contour (iso-z) or threshold lines in seaborn heatmap

Question

Is there a way to automatically add contour (iso-z) lines to a heatmap with concrete x and y values?

Please consider the official seaborn flights dataset:

import seaborn as sns
flights = sns.load_dataset("flights")
flights = flights.pivot("month", "year", "passengers")
sns.heatmap(flights, annot=True, fmt='d')

I imagine the step-like lines to look something like shown below (lhs), indicating thresholds (here 200 and 400). They do not need to be interpolated or smoothed in any way, although that would do as well, if easier to realize.

If the horizontal lines complicate the solution further, they too could be omitted (rhs).

So far, I have tried to add hlines and vlines manually, to overlay a kdeplot etc. without the desired result. Could somebody hint me into the right direction?

Answer 1

You can use aLineCollection:

import seaborn as sns
import numpy as np
from matplotlib.collections import LineCollection

flights = sns.load_dataset("flights")
flights = flights.pivot("month", "year", "passengers")
ax = sns.heatmap(flights, annot=True, fmt='d')

def add_iso_line(ax, value, color):
    v = flights.gt(value).diff(axis=1).fillna(False).to_numpy()
    h = flights.gt(value).diff(axis=0).fillna(False).to_numpy()
    
    try:
        l = np.argwhere(v.T)    
        vlines = np.array(list(zip(l, np.stack((l[:,0], l[:,1]+1)).T)))
        
        l = np.argwhere(h.T)    
        hlines = np.array(list(zip(l, np.stack((l[:,0]+1, l[:,1])).T)))
        
        lines = np.vstack((vlines, hlines))
        ax.add_collection(LineCollection(lines, lw=3, colors=color ))
    except:
        pass
    
add_iso_line(ax, 200, 'b')
add_iso_line(ax, 400, 'y')

Answer 2

The following approach uses a contour plot for to add the isolines. ndimage.zoom creates a refined grid which helps to obtain much smoother contour lines.

import seaborn as sns
import numpy as np
from matplotlib import pyplot as plt
from scipy import ndimage

flights = sns.load_dataset("flights")
flights = flights.pivot("month", "year", "passengers")
fig, ax = plt.subplots()

smooth_scale = 5
z = ndimage.zoom(flights.to_numpy(), smooth_scale)
cntr = ax.contour(np.linspace(0, len(flights.columns), len(flights.columns) * smooth_scale),
                  np.linspace(0, len(flights.index), len(flights.index) * smooth_scale),
                  z, levels=(200, 400), colors='yellow')
ax = sns.heatmap(flights, annot=True, fmt='d', cbar=True, ax=ax)

plt.tight_layout()
plt.show()

Alternatively, one could draw a contourf plot for filling the image, and only use the labels and annotations from sns.heatmap:

smooth_scale = 5
z = ndimage.zoom(flights.to_numpy(), smooth_scale)

cntr = ax.contourf(np.linspace(0, len(flights.columns), len(flights.columns) * smooth_scale),
                   np.linspace(0, len(flights.index), len(flights.index) * smooth_scale),
                   z, levels=np.arange(100, 701, 100), cmap='inferno')
ax = sns.heatmap(flights, annot=True, fmt='d', alpha=0, cbar=False, ax=ax)
plt.colorbar(cntr, ax=ax)