update

app.py

@@ -1,3 +1,6 @@
+import zipfile
+
+import yaml
 import gradio as gr
 import os
 import json
@@ -115,14 +118,20 @@ def run_query(protein, compound, disease, pathway, go, depth,
     must_show = sample_dict.copy()
     try:
         print('start sampling')
-        sub_g, new2orig, node_map_sub = subgraph_by_node(graph, sample_dict, node_map, depth=depth)
+        sub_g, new2orig, node_map_sub, statistics = subgraph_by_node(graph, sample_dict, node_map, depth=depth)
         id_map_sub = {k: {vv: kk for kk, vv in v.items()} for k, v in node_map_sub.items()}
+        # save statistics as json
+        with open(join(results_root, "statistics.yaml"), "w") as f:
+            yaml.dump(statistics, f)
+
+        statistics_text = "\n".join([f"{k}: {v}" for k, v in statistics.items()])
+        
         # 储存subgraph
         # save_subgraph_and_metadata(sub_g, id_map_sub, must_show)
         # 调用公共函数生成展示 HTML
         iframe_html, html_code = generate_iframe(sub_g, id_map_sub, must_show, display_limits)
 
-        return iframe_html, 'success', html_code, sub_g, id_map_sub, must_show
+        return iframe_html, 'success', statistics_text, sub_g, id_map_sub, must_show
     except Exception as e:
         return f"Error: {str(e)}", f"Error: {str(e)}", sample_dict, None, None, None
 
@@ -171,14 +180,32 @@ def get_text_content(file_path="static/gr_head.md"):
     with open(file_path, "r", encoding="utf-8") as f:
         return f.read()
 
+# def download_entity(sub_g, id_map_sub):
+#     try:
+#         report_subgraph(sub_g, id_map_sub, save_root=results_root)
+#         path = join(results_root, 'triples.txt')
+#         return gr.update(value=path, visible=True)
+#     except Exception as e:
+#         return gr.update(value=f"Error: {e}", visible=True)
+
 def download_entity(sub_g, id_map_sub):
     try:
+        # 生成统计数据和三元组文件
         report_subgraph(sub_g, id_map_sub, save_root=results_root)
-        path = os.path.join(results_root, 'triples.txt')
-        return gr.update(value=path, visible=True)
+        triples_path = join(results_root, 'triples.txt')
+        statistics_path = join(results_root, 'statistics.yaml')
+
+        # 创建一个压缩文件
+        zip_path = join(results_root, 'results.zip')
+        with zipfile.ZipFile(zip_path, 'w') as zipf:
+            zipf.write(triples_path, arcname='triples.txt')
+            zipf.write(statistics_path, arcname='statistics.yaml')
+
+        # 返回压缩文件路径
+        return gr.update(value=zip_path, visible=True)
     except Exception as e:
         return gr.update(value=f"Error: {e}", visible=True)
-
+    
 # 新增：加载静态文件的函数，如果存在则加载保存的子图数据
 def load_static_files():
     subgraph_file = "static/subgraph.dgl"

demo.ipynb

@@ -133,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -150,7 +150,7 @@
    "source": [
     "# 生成子图\n",
     "# 使用 subgraph_by_node 函数获取子图\n",
-    "sub_g, new2orig, node_map = subgraph_by_node(graph, sample_dict, node_map, depth=depth)\n",
+    "sub_g, new2orig, node_map, _ = subgraph_by_node(graph, sample_dict, node_map, depth=depth)\n",
     "# 打印子图的基本信息\n",
     "print(sub_g)\n",
     "# 获取子图的实体和三元组信息\n",

utils.py

@@ -104,6 +104,32 @@ def degree_search(graph, node_type, node_name, node_map):
             if deg > 0:
                 print(f"  {etype}: {deg}")
 
+
+def analyze_connections(graph, sample_dict, id_map):
+    connection_stats = {}
+    for node_type, node_ids in sample_dict.items():
+        for node_id in node_ids:
+            node_stats = {"connected_nodes": {}, "connected_edges": {}}
+            for etype in graph.canonical_etypes:
+                # 统计出度
+                if etype[0] == node_type:
+                    neighbors = graph.successors(node_id, etype=etype).tolist()
+                    node_stats["connected_nodes"].setdefault(etype[2], 0)
+                    node_stats["connected_nodes"][etype[2]] += len(neighbors)
+                    node_stats["connected_edges"].setdefault(etype, 0)
+                    node_stats["connected_edges"][etype] += len(neighbors)
+                # 统计入度
+                if etype[2] == node_type:
+                    neighbors = graph.predecessors(node_id, etype=etype).tolist()
+                    node_stats["connected_nodes"].setdefault(etype[0], 0)
+                    node_stats["connected_nodes"][etype[0]] += len(neighbors)
+                    node_stats["connected_edges"].setdefault(etype, 0)
+                    node_stats["connected_edges"][etype] += len(neighbors)
+            connection_stats[(node_type, id_map[node_type][node_id])] = node_stats
+    return connection_stats
+
+
+
 def subgraph_by_node(graph, sample_dict, node_map, depth=1):
     """
     Get a subgraph centered around a specific node.
@@ -115,6 +141,7 @@ def subgraph_by_node(graph, sample_dict, node_map, depth=1):
     Output:
         - full_g: The subgraph centered around the node.
     """
+    cur_id_map = {}
     # print(f"Getting subgraph from: {sample_dict}")
     for node_type, node_names in sample_dict.items():
         for node_name in node_names:
@@ -123,7 +150,9 @@ def subgraph_by_node(graph, sample_dict, node_map, depth=1):
                 return
         # convert node names to node IDs
         sample_dict[node_type] = [node_map[node_type][node_name] for node_name in node_names]
-    
+        cur_id_map[node_type] = {node_map[node_type][node_name]: node_name for node_name in node_names}
+
+    connection_stats = analyze_connections(graph, sample_dict, cur_id_map)
 
     out_g, _ = dgl.khop_out_subgraph(graph, sample_dict, k=depth,
                                         relabel_nodes=True, store_ids=True)
@@ -167,8 +196,7 @@ def subgraph_by_node(graph, sample_dict, node_map, depth=1):
     # for etype in full_g.canonical_etypes:
     #     print(f"{etype}: {full_g.num_edges(etype)}")
     
-    return full_g, new2orig, new_node_map
-
+    return full_g, new2orig, new_node_map, connection_stats
 
 def report_subgraph(graph, id_map, save_root='static'):
     entities = defaultdict(list)