如何在树状图（Dendrogram）中为分组聚类绘制彩色矩形框？

To add colored rectangular boxes around clusters in your dendrogram, you'll need to identify clusters from your linkage matrix, map them to the ordered leaves of the dendrogram, and then draw rectangles around each group using matplotlib patches. Here's how to modify your existing code to achieve this:

Modified Code with Cluster Boxes #

from scipy.cluster.hierarchy import dendrogram, fcluster
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Assume Z is your precomputed linkage matrix, corr is your DataFrame
file_name = "your_stock_data.csv"  # Replace with your actual file name

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)

# Capture dendrogram output to get leaf order and structure
dendro_out = dendrogram(
    Z,
    leaf_rotation=90.,
    leaf_font_size=8.,
    labels=corr.columns,
    color_threshold=None  # Optional: disable default line coloring to match box colors
)

# Step 1: Define clusters (choose one of the two options below)
# Option A: Specify number of clusters
num_clusters = 5  # Adjust this based on your desired cluster count
clusters = fcluster(Z, num_clusters, criterion='maxclust')

# Option B: Specify a distance threshold (uncomment to use instead)
# distance_threshold = 12  # Adjust based on your dendrogram's y-axis scale
# clusters = fcluster(Z, distance_threshold, criterion='distance')

# Step 2: Map clusters to the dendrogram's leaf order
cluster_labels_in_order = [clusters[i] for i in dendro_out['leaves']]
unique_clusters = sorted(set(cluster_labels_in_order))

# Step 3: Draw colored rectangles around each cluster
ax = plt.gca()
cmap = plt.cm.get_cmap('tab10', len(unique_clusters))  # Use a colormap with enough distinct colors

# Adjust box height relative to the plot's y-axis
y_min, y_max = ax.get_ylim()
box_height = y_max * 0.05  # Increase/decrease this value to make boxes taller/shorter

for cluster_num in unique_clusters:
    # Find all leaf indices belonging to the current cluster
    cluster_leaf_indices = [idx for idx, label in enumerate(cluster_labels_in_order) if label == cluster_num]
    x_start = min(cluster_leaf_indices)
    x_end = max(cluster_leaf_indices)
    
    # Calculate rectangle boundaries (each leaf occupies 1 unit on the x-axis)
    rect_left = x_start - 0.5
    rect_width = (x_end + 0.5) - rect_left
    
    # Create and add the rectangle patch
    rect = patches.Rectangle(
        (rect_left, y_min),
        rect_width,
        box_height,
        facecolor=cmap(cluster_num - 1),  # Cluster numbers start at 1, colormap indices start at 0
        alpha=0.3,  # Transparency ensures labels remain visible
        edgecolor='black'  # Optional: add a border for better definition
    )
    ax.add_patch(rect)

# Adjust tick formatting as before
plt.yticks(fontsize=ticksize)
plt.xticks(rotation=-90, fontsize=ticksize)

plt.show()

Key Notes: #

• Cluster Definition: Choose between specifying a fixed number of clusters (maxclust) or a distance threshold (distance) based on your analysis needs.
• Customization: Tweak box_height, alpha, cmap, or edgecolor to match the style of your example figure.
• Leaf Order: The dendrogram rearranges leaves based on clustering, so we map cluster labels to this reordered list to ensure boxes align correctly with the visual clusters.

内容的提问来源于stack exchange，提问作者Bedrick Kiq