294 lines
9.3 KiB
Python
294 lines
9.3 KiB
Python
import os #用于文件系统操作
|
||
import pandas as pd #用于数据处理和操作
|
||
import networkx as nx #用于创建和分析图结构
|
||
import plotly.graph_objects as go #plotly：用于创建交互式可视化 plotly.graph_objects：用于创建低级的plotly图形对象
|
||
from plotly.subplots import make_subplots #用于创建子图
|
||
import plotly.express as px #用于快速创建统计图表
|
||
|
||
def read_parquet_files(directory):
|
||
"""
|
||
读取指定目录下的所有Parquet文件并合并
|
||
功能：读取指定目录下的所有Parquet文件并合并成一个DataFrame
|
||
实现：使用os.listdir遍历目录，pd.read_parquet读取每个文件，然后用pd.concat合并
|
||
"""
|
||
dataframes = []
|
||
for filename in os.listdir(directory):
|
||
if filename.endswith('.parquet'):
|
||
file_path = os.path.join(directory, filename)
|
||
df = pd.read_parquet(file_path)
|
||
dataframes.append(df)
|
||
return pd.concat(dataframes, ignore_index=True) if dataframes else pd.DataFrame()
|
||
|
||
|
||
def clean_dataframe(df):
|
||
"""
|
||
清理DataFrame，移除无效的行
|
||
功能：清理DataFrame，移除无效的行
|
||
实现：删除source和target列中的空值，将这两列转换为字符串类型
|
||
"""
|
||
df = df.dropna(subset=['source', 'target'])
|
||
df['source'] = df['source'].astype(str)
|
||
df['target'] = df['target'].astype(str)
|
||
return df
|
||
|
||
|
||
def create_knowledge_graph(df):
|
||
"""
|
||
从DataFrame创建知识图谱
|
||
功能：从DataFrame创建知识图谱
|
||
实现：使用networkx创建有向图，遍历DataFrame的每一行，添加边和属性
|
||
"""
|
||
G = nx.DiGraph()
|
||
for _, row in df.iterrows():
|
||
source = row['source']
|
||
target = row['target']
|
||
attributes = {k: v for k, v in row.items() if k not in ['source', 'target']}
|
||
G.add_edge(source, target, **attributes)
|
||
return G
|
||
|
||
|
||
def create_node_link_trace(G, pos):
|
||
"""
|
||
功能：创建节点和边的3D轨迹
|
||
实现：使用networkx的布局信息创建Plotly的Scatter3d对象
|
||
"""
|
||
edge_x = []
|
||
edge_y = []
|
||
edge_z = []
|
||
for edge in G.edges():
|
||
x0, y0, z0 = pos[edge[0]]
|
||
x1, y1, z1 = pos[edge[1]]
|
||
edge_x.extend([x0, x1, None])
|
||
edge_y.extend([y0, y1, None])
|
||
edge_z.extend([z0, z1, None])
|
||
|
||
edge_trace = go.Scatter3d(
|
||
x=edge_x, y=edge_y, z=edge_z,
|
||
line=dict(width=0.5, color='#888'),
|
||
hoverinfo='none',
|
||
mode='lines')
|
||
|
||
node_x = [pos[node][0] for node in G.nodes()]
|
||
node_y = [pos[node][1] for node in G.nodes()]
|
||
node_z = [pos[node][2] for node in G.nodes()]
|
||
|
||
node_trace = go.Scatter3d(
|
||
x=node_x, y=node_y, z=node_z,
|
||
mode='markers',
|
||
hoverinfo='text',
|
||
marker=dict(
|
||
showscale=True,
|
||
colorscale='YlGnBu',
|
||
size=10,
|
||
colorbar=dict(
|
||
thickness=15,
|
||
title='Node Connections',
|
||
xanchor='left',
|
||
titleside='right'
|
||
)
|
||
)
|
||
)
|
||
|
||
node_adjacencies = []
|
||
node_text = []
|
||
for node, adjacencies in G.adjacency():
|
||
node_adjacencies.append(len(adjacencies))
|
||
node_text.append(f'Node: {node}<br># of connections: {len(adjacencies)}')
|
||
|
||
node_trace.marker.color = node_adjacencies
|
||
node_trace.text = node_text
|
||
|
||
return edge_trace, node_trace
|
||
|
||
|
||
def create_edge_label_trace(G, pos, edge_labels):
|
||
"""
|
||
功能：创建边标签的3D轨迹
|
||
实现：计算边的中点位置，创建Scatter3d对象显示标签
|
||
"""
|
||
return go.Scatter3d(
|
||
x=[pos[edge[0]][0] + (pos[edge[1]][0] - pos[edge[0]][0]) / 2 for edge in edge_labels],
|
||
y=[pos[edge[0]][1] + (pos[edge[1]][1] - pos[edge[0]][1]) / 2 for edge in edge_labels],
|
||
z=[pos[edge[0]][2] + (pos[edge[1]][2] - pos[edge[0]][2]) / 2 for edge in edge_labels],
|
||
mode='text',
|
||
text=list(edge_labels.values()),
|
||
textposition='middle center',
|
||
hoverinfo='none'
|
||
)
|
||
|
||
|
||
def create_degree_distribution(G):
|
||
"""
|
||
功能：创建节点度分布直方图
|
||
实现：使用plotly.express创建直方图
|
||
"""
|
||
degrees = [d for n, d in G.degree()]
|
||
fig = px.histogram(x=degrees, nbins=20, labels={'x': 'Degree', 'y': 'Count'})
|
||
fig.update_layout(
|
||
title_text='Node Degree Distribution',
|
||
margin=dict(l=0, r=0, t=30, b=0),
|
||
height=300
|
||
)
|
||
return fig
|
||
|
||
|
||
def create_centrality_plot(G):
|
||
"""
|
||
功能：创建节点中心性分布箱线图
|
||
实现：计算度中心性，使用plotly.express创建箱线图
|
||
"""
|
||
centrality = nx.degree_centrality(G)
|
||
centrality_values = list(centrality.values())
|
||
fig = px.box(y=centrality_values, labels={'y': 'Centrality'})
|
||
fig.update_layout(
|
||
title_text='Degree Centrality Distribution',
|
||
margin=dict(l=0, r=0, t=30, b=0),
|
||
height=300
|
||
)
|
||
return fig
|
||
|
||
|
||
def visualize_graph_plotly(G):
|
||
"""功能：使用Plotly创建全面优化布局的高级交互式知识图谱可视化
|
||
实现：
|
||
创建3D布局
|
||
生成节点和边的轨迹
|
||
创建子图，包括3D图、度分布图和中心性分布图
|
||
添加交互式按钮和滑块
|
||
优化整体布局
|
||
"""
|
||
if G.number_of_nodes() == 0:
|
||
print("Graph is empty. Nothing to visualize.")
|
||
return
|
||
|
||
pos = nx.spring_layout(G, dim=3) # 3D layout
|
||
edge_trace, node_trace = create_node_link_trace(G, pos)
|
||
|
||
edge_labels = nx.get_edge_attributes(G, 'relation')
|
||
edge_label_trace = create_edge_label_trace(G, pos, edge_labels)
|
||
|
||
degree_dist_fig = create_degree_distribution(G)
|
||
centrality_fig = create_centrality_plot(G)
|
||
|
||
fig = make_subplots(
|
||
rows=2, cols=2,
|
||
column_widths=[0.7, 0.3],
|
||
row_heights=[0.7, 0.3],
|
||
specs=[
|
||
[{"type": "scene", "rowspan": 2}, {"type": "xy"}],
|
||
[None, {"type": "xy"}]
|
||
],
|
||
subplot_titles=("3D Knowledge Graph Code by AI超元域频道", "Node Degree Distribution", "Degree Centrality Distribution")
|
||
)
|
||
|
||
fig.add_trace(edge_trace, row=1, col=1)
|
||
fig.add_trace(node_trace, row=1, col=1)
|
||
fig.add_trace(edge_label_trace, row=1, col=1)
|
||
|
||
fig.add_trace(degree_dist_fig.data[0], row=1, col=2)
|
||
fig.add_trace(centrality_fig.data[0], row=2, col=2)
|
||
|
||
# Update 3D layout
|
||
fig.update_layout(
|
||
scene=dict(
|
||
xaxis=dict(showticklabels=False, showgrid=False, zeroline=False),
|
||
yaxis=dict(showticklabels=False, showgrid=False, zeroline=False),
|
||
zaxis=dict(showticklabels=False, showgrid=False, zeroline=False),
|
||
aspectmode='cube'
|
||
),
|
||
scene_camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))
|
||
)
|
||
|
||
# Add buttons for different layouts
|
||
fig.update_layout(
|
||
updatemenus=[
|
||
dict(
|
||
type="buttons",
|
||
direction="left",
|
||
buttons=list([
|
||
dict(args=[{"visible": [True, True, True, True, True]}], label="Show All", method="update"),
|
||
dict(args=[{"visible": [True, True, False, True, True]}], label="Hide Edge Labels",
|
||
method="update"),
|
||
dict(args=[{"visible": [False, True, False, True, True]}], label="Nodes Only", method="update")
|
||
]),
|
||
pad={"r": 10, "t": 10},
|
||
showactive=True,
|
||
x=0.05,
|
||
xanchor="left",
|
||
y=1.1,
|
||
yanchor="top"
|
||
),
|
||
]
|
||
)
|
||
|
||
# Add slider for node size
|
||
fig.update_layout(
|
||
sliders=[dict(
|
||
active=0,
|
||
currentvalue={"prefix": "Node Size: "},
|
||
pad={"t": 50},
|
||
steps=[dict(method='update',
|
||
args=[{'marker.size': [i] * len(G.nodes)}],
|
||
label=str(i)) for i in range(5, 21, 5)]
|
||
)]
|
||
)
|
||
|
||
# 优化整体布局
|
||
# fig.update_layout(
|
||
# height=1198, # 增加整体高度
|
||
# width=2055, # 增加整体宽度
|
||
# title_text="Advanced Interactive Knowledge Graph",
|
||
# margin=dict(l=10, r=10, t=25, b=10),
|
||
# legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01)
|
||
# )
|
||
|
||
fig.show()
|
||
|
||
|
||
def main():
|
||
""" 功能：主函数，协调整个程序的执行流程
|
||
实现：
|
||
读取Parquet文件
|
||
清理数据
|
||
创建知识图谱
|
||
打印图的统计信息
|
||
调用可视化函数
|
||
"""
|
||
directory = '/Users/charlesqin/PycharmProjects/RAGCode/inputs/artifacts' # 替换为实际的目录路径
|
||
df = read_parquet_files(directory)
|
||
|
||
if df.empty:
|
||
print("No data found in the specified directory.")
|
||
return
|
||
|
||
print("Original DataFrame shape:", df.shape)
|
||
print("Original DataFrame columns:", df.columns.tolist())
|
||
print("Original DataFrame head:")
|
||
print(df.head())
|
||
|
||
df = clean_dataframe(df)
|
||
|
||
print("\nCleaned DataFrame shape:", df.shape)
|
||
print("Cleaned DataFrame head:")
|
||
print(df.head())
|
||
|
||
if df.empty:
|
||
print("No valid data remaining after cleaning.")
|
||
return
|
||
|
||
G = create_knowledge_graph(df)
|
||
|
||
print(f"\nGraph statistics:")
|
||
print(f"Nodes: {G.number_of_nodes()}")
|
||
print(f"Edges: {G.number_of_edges()}")
|
||
|
||
if G.number_of_nodes() > 0:
|
||
print(f"Connected components: {nx.number_connected_components(G.to_undirected())}")
|
||
visualize_graph_plotly(G)
|
||
else:
|
||
print("Graph is empty. Cannot visualize.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|