如何在 matplotlib 中调整树状图的分​​支长度（如在 astrodendro 中）？ [Python]

Question

下面是我的结果图，但我希望它看起来像 astrodendro 中截断的树状图，例如 this：

还有一个来自this paper 的非常酷的树状图，我想在matplotlib 中重新创建它。

下面是生成带有噪声变量的iris 数据集并在matplotlib 中绘制树状图的代码。

有谁知道如何：（1）像示例图中那样截断分支；和/或 (2) 将 astrodendro 与自定义链接矩阵和标签一起使用？

import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
import astrodendro
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial import distance

def iris_data(noise=None, palette="hls", desat=1):
    # Iris dataset
    X = pd.DataFrame(load_iris().data,
                     index = [*map(lambda x:f"iris_{x}", range(150))],
                     columns = [*map(lambda x: x.split(" (cm)")[0].replace(" ","_"), load_iris().feature_names)])

    y = pd.Series(load_iris().target,
                           index = X.index,
                           name = "Species")
    c = map_colors(y, mode=1, palette=palette, desat=desat)#y.map(lambda x:{0:"red",1:"green",2:"blue"}[x])

    if noise is not None:
        X_noise = pd.DataFrame(
            np.random.RandomState(0).normal(size=(X.shape[0], noise)),
            index=X_iris.index,
            columns=[*map(lambda x:f"noise_{x}", range(noise))]
        )
        X = pd.concat([X, X_noise], axis=1)
    return (X, y, c)

def dism2linkage(DF_dism, method="ward"):
    """
    Input: A (m x m) dissimalrity Pandas DataFrame object where the diagonal is 0
    Output: Hierarchical clustering encoded as a linkage matrix

    Further reading:
    http://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.cluster.hierarchy.linkage.html
    https://pypi.python.org/pypi/fastcluster
    """
    #Linkage Matrix
    Ar_dist = distance.squareform(DF_dism.as_matrix())
    return linkage(Ar_dist,method=method)


# Get data
X_iris_with_noise, y_iris, c_iris = iris_data(50)
# Get distance matrix
df_dism = 1- X_iris_with_noise.corr().abs()
# Get linkage matrix
Z = dism2linkage(df_dism)

#Create dendrogram
with plt.style.context("seaborn-white"):
    fig, ax = plt.subplots(figsize=(13,3))
    D_dendro = dendrogram(
             Z, 
             labels=df_dism.index,
             color_threshold=3.5,
             count_sort = "ascending",
             #link_color_func=lambda k: colors[k]
             ax=ax
    )
    ax.set_ylabel("Distance")

Answer 1

我不确定这是否真的构成了一个实用的答案，但它确实允许您生成带有截断悬挂线的树状图。诀窍是正常生成绘图，然后操纵生成的 matplotlib 绘图以重新创建线条。

我无法让您的示例在本地工作，所以我刚刚创建了一个虚拟数据集。

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

a = np.random.multivariate_normal([0, 10], [[3, 1], [1, 4]], size=[5,])
b = np.random.multivariate_normal([0, 10], [[3, 1], [1, 4]], size=[5,])
X = np.concatenate((a, b),)

Z = linkage(X, 'ward')

fig = plt.figure()
ax = fig.add_subplot(1,1,1)

dendrogram(Z, ax=ax)

生成的图是通常的长臂树状图。

现在来看更有趣的部分。一个树状图由多个LineCollection 对象组成（每种颜色一个）。为了更新行，我们遍历这些行，提取有关其组成路径的详细信息，修改这些以删除任何达到 y 为零的行，然后为这些修改后的路径重新创建 LineCollection。

然后将更新后的路径添加到坐标区，并删除原始路径。

一个棘手的部分是确定要绘制的高度而不是零。由于我们正在遍历每个树状图路径，因此我们不知道之前是哪一点——我们基本上不知道我们在哪里。但是，我们可以利用悬挂线垂直悬挂的事实。假设在同一个 x 上没有行，我们可以为给定的 x 寻找已知的其他 y 值，并在计算时将其用作新 y 的基础。缺点是为了确保我们有这个数字，我们必须预先扫描数据。

注意：如果您可以在同一 x 上获得树状图挂线，则需要包含 y 并搜索 此 x 上方最近的 y这样做。

import numpy as np
from matplotlib.path import Path
from matplotlib.collections import LineCollection

fig = plt.figure()
ax = fig.add_subplot(1,1,1)

dendrogram(Z, ax=ax);

for c in ax.collections[:]: # use [:] to get a copy, since we're adding to the same list
    paths = []
    for path in c.get_paths():
        segments = []
        y_at_x = {}
        # Pre-pass over all elements, to find the lowest y value at each x value.
        # we can use this to caculate where to cut our lines.
        for n, seg in enumerate(path.iter_segments()):
            x, y = seg[0]
            # Don't store if the y is zero, or if it's higher than the current low.
            if y > 0 and y < y_at_x.get(x, np.inf):
                y_at_x[x] = y

        for n, seg in enumerate(path.iter_segments()):
            x, y = seg[0]

            if y == 0:
                # If we know the last y at this x, use it - 0.5, limit > 0
                y = max(0, y_at_x.get(x, 0) - 0.5)

            segments.append([x,y])

        paths.append(segments)

    lc = LineCollection(paths, colors=c.get_colors())  # Recreate a LineCollection with the same params
    ax.add_collection(lc)
    ax.collections.remove(c) # Remove the original LineCollection

生成的树状图如下所示：