如何在dendogram中围绕分组簇绘制彩色矩形？

Question

我尝试将彩色矩形添加到树状图结果中，如下所示：

这是我的树状图代码：

from scipy.cluster.hierarchy import dendrogram
...
plt.figure(figsize=(250, 100))
labelsize=20
ticksize=15
plt.title(file_name.split(".")[0], fontsize=labelsize)
plt.xlabel('stock', fontsize=labelsize)
plt.ylabel('distance', fontsize=labelsize)
dendrogram(
   Z,
   leaf_rotation=90.,  # rotates the x axis labels
   leaf_font_size=8.,  # font size for the x axis labels
   labels = corr.columns
)
pylab.yticks(fontsize=ticksize)
pylab.xticks(rotation=-90, fontsize=ticksize)

但是，这只是添加彩色线，而不是上图中的矩形。我怎样才能创建这样的图像

谢谢

Answer 1

可以在生成的路径集合中循环并绘制边界框

或者，您可以将高度设置为color_threshold=参数，该参数默认为Z[:, 2].max() * 0.7

最后一个集合是未分类的行，因此下面的示例代码循环遍历所有早期集合

import matplotlib.pyplot as plt
from scipy.cluster import hierarchy
import numpy as np

N = 15
ytdist = np.random.randint(10, 1000, N * (N + 1) // 2)
Z = hierarchy.linkage(ytdist)

fig, ax = plt.subplots(1, 1, figsize=(8, 3))
dn1 = hierarchy.dendrogram(Z, ax=ax)

for coll in ax.collections[:-1]:  # the last collection is the ungrouped level
    xmin, xmax = np.inf, -np.inf
    ymax = -np.inf
    for p in coll.get_paths():
        box = p.get_extents()
        (x0, _), (x1, y1) = p.get_extents().get_points()
        xmin = min(xmin, x0)
        xmax = max(xmax, x1)
        ymax = max(ymax, y1)
    rec = plt.Rectangle((xmin - 4, 0), xmax - xmin + 8, ymax*1.05,
                        facecolor=coll.get_color()[0], alpha=0.2, edgecolor="none")
    ax.add_patch(rec)
plt.show()