基于NetworkX调整边长度展示相关矩阵及可视化优化需求

完整实现代码 #

import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize

# 生成示例数据（真实数据可替换此处）
ranstate = np.random.RandomState(0)
df = pd.DataFrame(ranstate.rand(10, 10))
df[0] = df[0] + df[1]
df[2] = df[0] + df[1]
df[7] = df[9] + df[8]

corr = df.corr()

# 处理相关系数数据，保留阈值以上且非自相关的边
links = corr.stack().reset_index()
links.columns = ['var1', 'var2', 'value']
links_filtered = links.loc[(links['value'] > 0.7) & (links['var1'] != links['var2'])]

# 构建图，同时把相关系数作为边属性存入
G = nx.from_pandas_edgelist(links_filtered, 'var1', 'var2', edge_attr='value')

# 1. 设置布局，控制最小边长度（节点间距）
# k参数控制节点间距，值越大间距越大；seed保证布局可复现
pos = nx.spring_layout(G, k=0.5, seed=42, scale=2)

# 2. 准备边的颜色：根据相关系数做深浅梯度
edge_values = [G[u][v]['value'] for u, v in G.edges()]
norm = Normalize(vmin=min(edge_values), vmax=max(edge_values))
cmap = cm.coolwarm
edge_colors = [cmap(norm(val)) for val in edge_values]

# 绘制节点
nx.draw(G, pos, with_labels=True, node_color='orange', node_size=400, 
        edge_color=edge_colors, linewidths=1, font_size=15)

# 3. 为边标注相关系数（保留两位小数）
edge_labels = {(u, v): f"{G[u][v]['value']:.2f}" for u, v in G.edges()}
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=12)

plt.show()

需求实现细节 #

• 设置最小边长度：使用nx.spring_layout的k参数调整节点间的排斥力，k值越大，节点间距越大，避免真实数据中节点重叠无法阅读。同时设置seed参数保证每次运行布局一致，方便调试。
• 边标注相关系数：通过字典推导式提取每条边的value属性（即相关系数），格式化保留两位小数后，用nx.draw_networkx_edge_labels将标注放在对应边的中间位置。
• 边的深浅梯度着色：借助matplotlib的颜色映射工具，先将相关系数值归一化到0-1区间，再通过coolwarm色映射生成对应深浅的颜色，相关系数越高（越接近1）颜色越深，反之越浅。

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