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

user_11804186 0
Viewed 1989

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

Seaborn flights example with step-like iso-z lines

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?

seaborn matplotlib python
2 Answers

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

enter image description here

user_39443220

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

resulting plot

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)

plot using contourf

user_120464090
Related