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?

Leon

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')

Silsbye

If all values of np.argwhere(v.T) are False, then np.vstack((vlines, hlines)) would return a dimension error – Saire

What should we do in that case? – Saire

@RajeshC thanks for the comment: in this case the requested contour level is outside the value range of the underlying array and no contour lines are to be added, so we just wrap the whole logic in a try/except block (see updated code) – Silsbye

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)

Jez