Dendrograma en Python con scipy y matplotlib

Un dendrograma es un diagrama en forma de árbol que muestra el resultado de un clustering jerárquico: las observaciones que se unen más abajo en el árbol son más parecidas. En Python el clustering se calcula con linkage de scipy.cluster.hierarchy y el árbol se dibuja con dendrogram. Los ejemplos utilizan datos sobre accidentes de tráfico en 20 estados de Estados Unidos.

Dendrograma con linkage y dendrogram

linkage calcula el clustering jerárquico (en este caso con el método de Ward) y dendrogram lo dibuja sobre los ejes actuales de matplotlib.

import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram

crashes = sns.load_dataset("car_crashes").set_index("abbrev").head(20)
Z = linkage(crashes[["speeding", "alcohol", "not_distracted", "no_previous"]], method = "ward")

fig, ax = plt.subplots()
dendrogram(Z, labels = crashes.index.tolist(), ax = ax)
# plt.show()

Dendrograma horizontal

orientation coloca la raíz del árbol arriba ("top", por defecto), abajo ("bottom"), a la izquierda ("left") o a la derecha ("right"). Un dendrograma horizontal deja más espacio para las etiquetas.

import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram

crashes = sns.load_dataset("car_crashes").set_index("abbrev").head(20)
Z = linkage(crashes[["speeding", "alcohol", "not_distracted", "no_previous"]], method = "ward")

fig, ax = plt.subplots()
dendrogram(Z, labels = crashes.index.tolist(), orientation = "left", ax = ax)
# plt.show()

Colorear los clusters con un umbral

color_threshold establece la altura a la que se corta el árbol: cada cluster por debajo de esa altura recibe un color diferente. Dibuja el corte con axhline para que sea visible.

import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram

crashes = sns.load_dataset("car_crashes").set_index("abbrev").head(20)
Z = linkage(crashes[["speeding", "alcohol", "not_distracted", "no_previous"]], method = "ward")

fig, ax = plt.subplots()
dendrogram(Z, labels = crashes.index.tolist(), color_threshold = 15,
           above_threshold_color = "gray", ax = ax)
ax.axhline(15, color = "black", linestyle = "--")
# plt.show()

Truncar el dendrograma

Con muchas observaciones el árbol es difícil de leer. truncate_mode = "lastp" muestra solo los últimos p clusters unidos; el número de observaciones de cada uno se muestra entre paréntesis.

import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram

crashes = sns.load_dataset("car_crashes").set_index("abbrev")
Z = linkage(crashes[["speeding", "alcohol", "not_distracted", "no_previous"]], method = "ward")

fig, ax = plt.subplots()
dendrogram(Z, truncate_mode = "lastp", p = 8, ax = ax)
# plt.show()

Métodos de enlace

La forma del árbol depende del método utilizado para medir la distancia entre clusters: "ward", "complete", "average", "single"…

import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram

crashes = sns.load_dataset("car_crashes").set_index("abbrev").head(20)
X = crashes[["speeding", "alcohol", "not_distracted", "no_previous"]]

fig, axs = plt.subplots(1, 2, figsize = (8, 3.5))
for ax, method in zip(axs, ["ward", "single"]):
    dendrogram(linkage(X, method = method), labels = X.index.tolist(), ax = ax)
    ax.set_title(method)
# plt.show()

Para combinar el dendrograma con un mapa de calor utiliza el clustermap de seaborn. Consulta también el dendrograma en plotly.