197 lines
5.9 KiB
Python
197 lines
5.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
知识图谱可视化模块
|
|
使用 pyvis 生成交互式网络图，支持嵌入 Streamlit
|
|
"""
|
|
import logging
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import networkx as nx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# 节点类型对应的颜色和形状
|
|
NODE_STYLES = {
|
|
"科创企业": {"color": "#4CAF50", "shape": "dot", "size": 30},
|
|
"供应商": {"color": "#FF9800", "shape": "diamond", "size": 20},
|
|
"客户": {"color": "#9C27B0", "shape": "triangle", "size": 20},
|
|
"核心人员": {"color": "#2196F3", "shape": "star", "size": 15},
|
|
"投资方": {"color": "#607D8B", "shape": "square", "size": 20},
|
|
"企业": {"color": "#4CAF50", "shape": "dot", "size": 25},
|
|
}
|
|
|
|
# 受制裁节点的样式覆盖
|
|
SANCTIONED_STYLE = {"color": "#F44336", "size": 35}
|
|
# 受限边的样式
|
|
RESTRICTED_EDGE_STYLE = {"color": "#F44336", "dashes": True, "width": 3}
|
|
|
|
|
|
def generate_interactive_graph(
|
|
G: nx.DiGraph,
|
|
highlight_company: Optional[str] = None,
|
|
output_path: Optional[str] = None,
|
|
height: str = "600px",
|
|
width: str = "100%",
|
|
) -> str:
|
|
"""
|
|
生成交互式知识图谱 HTML
|
|
"""
|
|
try:
|
|
from pyvis.network import Network
|
|
except ImportError:
|
|
logger.error("pyvis 未安装，请运行: pip install pyvis")
|
|
return _generate_fallback_html(G)
|
|
|
|
net = Network(
|
|
height=height,
|
|
width=width,
|
|
directed=True,
|
|
notebook=False,
|
|
bgcolor="#1a1a2e",
|
|
font_color="white",
|
|
)
|
|
|
|
# 物理引擎配置
|
|
net.set_options("""
|
|
{
|
|
"physics": {
|
|
"forceAtlas2Based": {
|
|
"gravitationalConstant": -50,
|
|
"centralGravity": 0.01,
|
|
"springLength": 150,
|
|
"springConstant": 0.08
|
|
},
|
|
"solver": "forceAtlas2Based",
|
|
"stabilization": {"iterations": 100}
|
|
},
|
|
"interaction": {
|
|
"hover": true,
|
|
"tooltipDelay": 200,
|
|
"navigationButtons": true
|
|
}
|
|
}
|
|
""")
|
|
|
|
# 添加节点
|
|
for node, data in G.nodes(data=True):
|
|
node_type = data.get("node_type", "企业")
|
|
style = NODE_STYLES.get(node_type, NODE_STYLES["企业"]).copy()
|
|
|
|
# 受制裁节点特殊样式
|
|
if data.get("is_sanctioned", False):
|
|
style.update(SANCTIONED_STYLE)
|
|
|
|
# 高亮选中企业
|
|
if highlight_company and node == highlight_company:
|
|
style["color"] = "#FFD700" # 金色
|
|
style["size"] = 45
|
|
style["borderWidth"] = 3
|
|
|
|
# 构建标签和悬浮提示
|
|
label = node.split("@")[0] if "@" in node else node
|
|
title_parts = [f"<b>{label}</b>", f"类型: {node_type}"]
|
|
if data.get("is_sanctioned"):
|
|
title_parts.append("⚠️ <b>已被制裁</b>")
|
|
if data.get("stock_code"):
|
|
title_parts.append(f"代码: {data['stock_code']}")
|
|
if data.get("industry"):
|
|
title_parts.append(f"行业: {data['industry']}")
|
|
if data.get("title"):
|
|
title_parts.append(f"职位: {data['title']}")
|
|
if data.get("status"):
|
|
title_parts.append(f"状态: {data['status']}")
|
|
|
|
net.add_node(
|
|
node,
|
|
label=label,
|
|
title="<br>".join(title_parts),
|
|
color=style["color"],
|
|
shape=style["shape"],
|
|
size=style["size"],
|
|
)
|
|
|
|
# 添加边
|
|
for u, v, data in G.edges(data=True):
|
|
edge_style = {
|
|
"color": "#666666",
|
|
"width": 1,
|
|
"arrows": "to",
|
|
}
|
|
|
|
# 受限边特殊样式
|
|
if data.get("status") == "受限":
|
|
edge_style.update(RESTRICTED_EDGE_STYLE)
|
|
elif data.get("critical", False):
|
|
edge_style["width"] = 2
|
|
edge_style["color"] = "#FFC107" # 关键边用黄色
|
|
|
|
relation = data.get("relation", "")
|
|
net.add_edge(
|
|
u, v,
|
|
title=relation,
|
|
label=relation if len(relation) <= 8 else "",
|
|
**edge_style,
|
|
)
|
|
|
|
# 输出 HTML
|
|
if output_path is None:
|
|
output_path = str(Path(tempfile.gettempdir()) / "kg_visualization.html")
|
|
|
|
net.save_graph(output_path)
|
|
|
|
# 读取 HTML 内容
|
|
with open(output_path, "r", encoding="utf-8") as f:
|
|
return f.read()
|
|
|
|
|
|
def _generate_fallback_html(G: nx.DiGraph) -> str:
|
|
"""当 pyvis 不可用时的备用 HTML 可视化"""
|
|
nodes_info = []
|
|
for node, data in G.nodes(data=True):
|
|
label = node.split("@")[0] if "@" in node else node
|
|
is_sanctioned = data.get("is_sanctioned", False)
|
|
nodes_info.append(f"<li style='color: {'red' if is_sanctioned else 'green'}'>{label} ({data.get('node_type', '未知')})</li>")
|
|
|
|
return f"""
|
|
<html><body style='background: #1a1a2e; color: white; padding: 20px;'>
|
|
<h2>📊 知识图谱节点列表（pyvis 未安装，使用简化视图）</h2>
|
|
<p>节点数: {G.number_of_nodes()} | 边数: {G.number_of_edges()}</p>
|
|
<ul>{''.join(nodes_info[:50])}</ul>
|
|
<p style='color: #888;'>安装 pyvis 以获得交互式可视化: pip install pyvis</p>
|
|
</body></html>
|
|
"""
|
|
|
|
|
|
def get_subgraph_for_company(G: nx.DiGraph, company: str, depth: int = 2) -> nx.DiGraph:
|
|
"""
|
|
提取以指定企业为中心的子图（上下游 N 层）
|
|
"""
|
|
if company not in G:
|
|
return nx.DiGraph()
|
|
|
|
# 收集相关节点
|
|
related_nodes = {company}
|
|
|
|
# 上游（前驱）
|
|
current_layer = {company}
|
|
for _ in range(depth):
|
|
next_layer = set()
|
|
for node in current_layer:
|
|
next_layer.update(G.predecessors(node))
|
|
related_nodes.update(next_layer)
|
|
current_layer = next_layer
|
|
|
|
# 下游（后继）
|
|
current_layer = {company}
|
|
for _ in range(depth):
|
|
next_layer = set()
|
|
for node in current_layer:
|
|
next_layer.update(G.successors(node))
|
|
related_nodes.update(next_layer)
|
|
current_layer = next_layer
|
|
|
|
return G.subgraph(related_nodes).copy()
|