seaborn 热图中的轮廓 (iso-z) 或阈值线

Question

有没有办法自动将轮廓（iso-z）线添加到具有具体 x 和 y 值的热图中？

请考虑官方的seaborn航班数据集：

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

我想象阶梯状的线条看起来像下面显示的（lhs），表示阈值（这里是 200 和 400）。它们不需要以任何方式内插或平滑，尽管如果更容易实现的话，也可以这样做。

如果水平线使解决方案进一步复杂化，它们也可以省略（rhs）。

到目前为止，我已经尝试手动添加 hlines 和 vlines，以覆盖 kdeplot 等，但没有达到预期的结果。有人可以暗示我走向正确的方向吗？

Answer 1

您可以使用LineCollection：

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

以下方法使用等高线图来添加等值线。ndimage.zoom创建一个精致的网格，有助于获得更平滑的轮廓线。

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

或者，可以绘制一个contourf图来填充图像，并且仅使用以下标签和注释sns.heatmap：

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