v2431/FormulaShowcase
0
1# data_process.py2# -*- coding: utf-8 -*-3"""4数据读取与统一数据结构定义模块5 6本文件主要提供：71. KGFileConfig 小配置类：管理数据库文件路径82. UniBioMapKG 类：9 - 从 CSV 读取主图和天然产物图10 - 提供 PPI 子图 (networkx.Graph)11 - 提供任务输入 {type: [ids]} 映射到蛋白集合的功能12 - 提供天然来源 natural_source -> 蛋白靶标 映射13 14说明：15- 所有大规模网络分析都基于主图 (links_df)，天然来源关系来自 natural_sources_df。16"""17 18from __future__ import annotations19 20from dataclasses import dataclass21import json22from pathlib import Path23from typing import Dict, List, Optional, Set, Iterable24 25import pandas as pd26import networkx as nx27import math28from collections import defaultdict29from utils import custom_csv_loader, custom_chem_txt_loader, custom_ns_txt_loader30 31@dataclass32class KGFileConfig:33 """34 用来管理路径的小配置类35 36 Attributes37 ----------38 db_dir : str39 数据库目录，如 'database'40 links_filename : str41 主图三元组文件名，如 'unibiomap.links.csv'42 natural_sources_filename : str43 天然产物三元组文件名，如 'unibiomap.natural_sources.csv'44 """45 46 db_dir: str = "database"47 links_filename: str = "unibiomap.links.csv"48 natural_sources_filename: str = "unibiomap.natural_sources.csv"49 50 @property51 def links_path(self) -> Path:52 return Path(self.db_dir) / self.links_filename53 54 @property55 def natural_sources_path(self) -> Path:56 return Path(self.db_dir) / self.natural_sources_filename57 58 def desc_path(self, entity_type: str) -> Path:59 filename = f"{entity_type}_desc.json"60 return Path(self.db_dir) / filename61 62 63class UniBioMapKG:64 """65 用于封装 UniBioMap 知识图谱的核心数据结构。66 67 核心职责：68 ----------69 1. 从 CSV 读取主图 (links) 和天然产物图 (natural_sources)70 2. 提供 PPI 子图 (仅 protein-protein 边) 的 networkx.Graph 视图71 3. 提供任务输入 {type: [ids]} 映射到蛋白集合的工具方法72 4. 提供天然来源 natural_source -> 蛋白靶标集合 的映射方法73 74 注意：75 ----------76 - 节点 ID 统一使用原始实体编号（不拼 type 前缀）。77 - PPI 图中所有节点都是 protein 类型。78 """79 80 def __init__(81 self,82 file_config: Optional[KGFileConfig] = None,83 *,84 min_conf: Optional[float] = None,85 species_filter: Optional[List[str]] = None, # 使用 taxonomy ID 过滤物种86 ) -> None:87 """88 Parameters89 ----------90 file_config : KGFileConfig, optional91 文件路径配置，如果为 None 则使用默认 'database/*'92 min_conf : float, optional93 对于主图中带 conf 的边，如果指定则只保留 conf >= min_conf 的边。94 若为 None 则不按 conf 过滤。95 """96 self.config = file_config or KGFileConfig()97 self.min_conf = min_conf98 99 self.species_filter = species_filter100 self._protein_set = self._load_species_proteins()101 102 # 主图三元组：htype, ttype, h, r, t, src, conf103 self.links_df: pd.DataFrame = self._load_links()104 105 # 天然产物图三元组：htype, ttype, h, r, t, src106 self.natural_sources_df: Optional[pd.DataFrame] = self._load_natural_sources()107 108 # PPI 子图缓存（networkx.Graph）109 self._ppi_graph: Optional[nx.Graph] = None110 111 # 复用缓存：compound -> set(protein)112 self._compound_to_proteins: Optional[Dict[str, Set[str]]] = None113 # 复用缓存：compound -> rel -> protein -> max_conf（用于按关系/置信度进行精细裁剪）114 self._compound_to_protein_edges: Optional[Dict[str, Dict[str, Dict[str, float]]]] = None115 # 复用缓存：protein -> 关联到多少个 compound（全局频次，用于 IDF/泛化度惩罚）116 self._protein_df_compound: Optional[Dict[str, int]] = None117 # 复用缓存：用于 IDF 计算的 compound 总数（基于 compound-protein 边）118 self._num_compounds_in_edges: Optional[int] = None119 # 复用缓存：natural_source -> set(compound)120 self._natural_source_to_compounds: Optional[Dict[str, Set[str]]] = None121 # 复用缓存：natural_source -> set(protein)122 self._natural_source_to_proteins: Optional[Dict[str, Set[str]]] = None123 124 # PPI 的 GCC 子图缓存（仅包含最大连通分量，用于 NetMedPy 等网络医学分析）125 self._ppi_gcc_graph: Optional[nx.Graph] = None126 127 # prefilter 日志：key 为候选/来源ID，val 为多次运行记录列表128 self.prefilter_log: Dict[str, List[Dict]] = {}129 130 # ------------------------------------------------------------------131 # 数据加载部分132 # ------------------------------------------------------------------133 def _load_species_proteins(self) -> set:134 if self.species_filter is None:135 return None136 with open(self.config.desc_path("protein"), "r", encoding="utf-8") as f:137 prot_desc = json.load(f)138 protein_set = set()139 for k, v in prot_desc.items():140 tax_id = v.get("specie")141 if tax_id in self.species_filter:142 protein_set.add(k)143 return protein_set144 145 def _load_links(self) -> pd.DataFrame:146 """147 读取主图三元组 CSV，并做基本类型转换。148 如果设置了物种过滤，则只保留符合条件的边。149 150 Returns151 -------152 pd.DataFrame153 包含列 ['htype', 'ttype', 'h', 'r', 't', 'src', 'conf']154 """155 path = self.config.links_path156 if not path.exists():157 raise FileNotFoundError(f"主图文件不存在：{path}")158 159 df = pd.read_csv(160 path,161 dtype={162 "htype": "string",163 "ttype": "string",164 "h": "string",165 "r": "string",166 "t": "string",167 "src": "string",168 },169 )170 171 if "conf" in df.columns:172 df["conf"] = pd.to_numeric(df["conf"], errors="coerce")173 else:174 df["conf"] = pd.NA175 176 # 按置信度过滤177 if self.min_conf is not None:178 df = df[df["conf"].fillna(-1.0) >= self.min_conf].copy()179 180 df["htype"] = df["htype"].str.lower()181 df["ttype"] = df["ttype"].str.lower()182 183 # 按物种过滤蛋白质184 if self._protein_set is not None:185 # 判断头节点是否为蛋白质186 h_is_protein = df["htype"] == "protein"187 # 判断尾节点是否为蛋白质188 t_is_protein = df["ttype"] == "protein"189 190 # 头节点是蛋白质且在集合中191 h_valid = (~h_is_protein) | df["h"].isin(self._protein_set)192 # 尾节点是蛋白质且在集合中193 t_valid = (~t_is_protein) | df["t"].isin(self._protein_set)194 195 # 两个条件都满足才保留196 df = df[h_valid & t_valid].copy()197 198 return df199 200 201 def _load_natural_sources(self) -> Optional[pd.DataFrame]:202 """203 读取天然产物三元组 CSV（如果存在）。204 205 Returns206 -------207 pd.DataFrame or None208 包含列 ['htype', 'ttype', 'h', 'r', 't', 'src']209 """210 path = self.config.natural_sources_path211 if not path.exists():212 return None213 214 df = pd.read_csv(215 path,216 dtype={217 "htype": "string",218 "ttype": "string",219 "h": "string",220 "r": "string",221 "t": "string",222 "src": "string",223 },224 )225 df["htype"] = df["htype"].str.lower()226 df["ttype"] = df["ttype"].str.lower()227 return df228 229 # ------------------------------------------------------------------230 # PPI 子图部分231 # ------------------------------------------------------------------232 def build_ppi_graph(233 self,234 relation_filter: Optional[Iterable[str]] = None,235 ) -> nx.Graph:236 """237 从主图三元组中提取 protein-protein 边，构建 PPI graph。238 239 Parameters240 ----------241 relation_filter : Iterable[str], optional242 如果提供，只保留 r 在该集合中的边。243 若为 None 则不过滤 r。244 245 Returns246 -------247 networkx.Graph248 无向 PPI 图，节点为 protein 实体编号（字符串）249 """250 df = self.links_df251 252 mask = (df["htype"] == "protein") & (df["ttype"] == "protein")253 254 if relation_filter is not None:255 rel_set = set(relation_filter)256 mask &= df["r"].isin(rel_set)257 258 ppi_df = df.loc[mask, ["h", "t"]]259 260 G = nx.Graph()261 nodes = pd.unique(pd.concat([ppi_df["h"], ppi_df["t"]], ignore_index=True))262 G.add_nodes_from(nodes)263 G.add_edges_from(ppi_df.itertuples(index=False, name=None))264 return G265 266 @property267 def ppi_graph(self) -> nx.Graph:268 """269 懒加载 PPI graph，只构建一次，避免重复开销。270 """271 if self._ppi_graph is None:272 self._ppi_graph = self.build_ppi_graph()273 return self._ppi_graph274 275 # ------------------------------------------------------------------276 # PPI Giant Connected Component 巨型连通分量部分277 # ------------------------------------------------------------------278 def build_ppi_gcc(279 self,280 relation_filter: Optional[Iterable[str]] = None,281 ) -> nx.Graph:282 """283 从主图三元组中提取 protein-protein 边，并进一步只保留284 最大连通分量（GCC），构建 PPI GCC 图。285 286 说明：287 ----------288 - NetMedPy 的 all_pair_distances 要求图是连通的，289 因此我们在这里强制裁剪到最大连通分量。290 - 对于网络医学分析（LCC、proximity、separation 等），291 都建议使用这个 GCC 图。292 293 Parameters294 ----------295 relation_filter : Iterable[str], optional296 如果提供，只保留 r 在该集合中的边。297 若为 None 则不过滤 r。298 299 Returns300 -------301 networkx.Graph302 无向 PPI GCC 图，节点为 protein 实体编号（字符串）303 """304 # 先构建完整的 protein-protein PPI 图305 full_ppi = self.ppi_graph306 307 # 如果图为空，直接返回空图308 if full_ppi.number_of_nodes() == 0:309 return full_ppi310 311 # 如果本来就是连通图，就直接返回312 if nx.is_connected(full_ppi):313 return full_ppi314 315 # 否则只保留最大连通分量（GCC）316 gcc_nodes = max(nx.connected_components(full_ppi), key=len)317 gcc_graph = full_ppi.subgraph(gcc_nodes).copy()318 return gcc_graph319 320 @property321 def ppi_gcc(self) -> nx.Graph:322 """323 懒加载 PPI 的 GCC 子图（仅包含最大连通分量）。324 325 用途：326 ----------327 - 所有需要网络距离 / NetMedPy 的分析（LCC、proximity 等），328 建议都使用 ppi_gcc，而不是 ppi_graph。329 """330 if self._ppi_gcc_graph is None:331 self._ppi_gcc_graph = self.build_ppi_gcc()332 return self._ppi_gcc_graph333 334 def report_gcc_coverage(335 self,336 seed_proteins: Iterable[str],337 module_nodes: Optional[Iterable[str]] = None,338 ns_targets: Optional[Dict[str, Set[str]]] = None,339 *,340 verbose: bool = True,341 ) -> Dict[str, float]:342 """343 一键统计并输出：344 1. PPI GCC 的大小（节点数 / 边数）345 2. 疲劳种子蛋白在 GCC 中的覆盖率346 3. 疲劳模块（例如 LCC）在 GCC 中的覆盖率（如果提供）347 4. 所有天然来源靶点在 GCC 中的覆盖率（如果有天然来源映射）348 349 Parameters350 ----------351 seed_proteins : Iterable[str]352 例如由 get_seed_proteins(task_input) 得到的种子蛋白集合353 module_nodes : Iterable[str], optional354 例如 LCCResult.lcc_nodes（疲劳模块 LCC 的节点集合）355 ns_targets : dict, optional356 key: natural_source id357 value: set of protein id358 若为 None，则自动调用 get_natural_source_targets()359 verbose : bool, default True360 若为 True，则在 stdout 打印可读的统计信息361 362 Returns363 -------364 dict365 一个包含各类覆盖率指标的字典，字段包括：366 - gcc_num_nodes, gcc_num_edges367 - seed_total, seed_in_gcc, seed_coverage368 - module_total, module_in_gcc, module_coverage369 - ns_targets_total, ns_targets_in_gcc, ns_targets_coverage370 """371 # 1) 取出 GCC 图372 gcc_graph = self.ppi_gcc373 gcc_nodes = set(gcc_graph.nodes())374 375 gcc_num_nodes = gcc_graph.number_of_nodes()376 gcc_num_edges = gcc_graph.number_of_edges()377 378 # 2) 种子蛋白覆盖率379 seed_set = set(seed_proteins)380 seed_total = len(seed_set)381 seed_in_gcc = len(seed_set & gcc_nodes)382 seed_coverage = (seed_in_gcc / seed_total) if seed_total > 0 else 0.0383 384 # 3) 模块节点覆盖率（如果给了 module_nodes）385 if module_nodes is not None:386 module_set = set(module_nodes)387 module_total = len(module_set)388 module_in_gcc = len(module_set & gcc_nodes)389 module_coverage = (module_in_gcc / module_total) if module_total > 0 else 0.0390 else:391 module_set = set()392 module_total = 0393 module_in_gcc = 0394 module_coverage = 0.0395 396 # 4) 天然来源靶点覆盖率397 if ns_targets is None:398 ns_targets = self.get_natural_source_targets()399 400 # 收集所有天然来源的蛋白靶点并去重401 all_ns_proteins: Set[str] = set()402 for prot_set in ns_targets.values():403 all_ns_proteins.update(prot_set)404 405 ns_targets_total = len(all_ns_proteins)406 ns_targets_in_gcc = len(all_ns_proteins & gcc_nodes)407 ns_targets_coverage = (408 ns_targets_in_gcc / ns_targets_total if ns_targets_total > 0 else 0.0409 )410 411 # 5) 汇总结果字典412 stats = {413 "gcc_num_nodes": gcc_num_nodes,414 "gcc_num_edges": gcc_num_edges,415 "seed_total": seed_total,416 "seed_in_gcc": seed_in_gcc,417 "seed_coverage": seed_coverage,418 "module_total": module_total,419 "module_in_gcc": module_in_gcc,420 "module_coverage": module_coverage,421 "ns_targets_total": ns_targets_total,422 "ns_targets_in_gcc": ns_targets_in_gcc,423 "ns_targets_coverage": ns_targets_coverage,424 }425 426 # 6) 可选打印427 if verbose:428 print("========== GCC 覆盖率报告 ==========")429 print(f"GCC 大小: 节点数 = {gcc_num_nodes}, 边数 = {gcc_num_edges}")430 print(431 f"种子蛋白: 总数 = {seed_total}, "432 f"在 GCC 中 = {seed_in_gcc} ({seed_coverage:.2%})"433 )434 if module_total > 0:435 print(436 f"模块蛋白: 总数 = {module_total}, "437 f"在 GCC 中 = {module_in_gcc} ({module_coverage:.2%})"438 )439 else:440 print("模块蛋白: 未提供（module_nodes=None）")441 442 print(443 f"天然来源靶点: 去重总数 = {ns_targets_total}, "444 f"在 GCC 中 = {ns_targets_in_gcc} ({ns_targets_coverage:.2%})"445 )446 print("====================================")447 448 return stats449 450 # ------------------------------------------------------------------451 # 任务输入映射部分：{type: [ids]} -> seed proteins452 # ------------------------------------------------------------------453 def get_seed_proteins(454 self,455 task_input: Dict[str, List[str]],456 ) -> Set[str]:457 """458 根据任务输入 {type: [ids]}，得到对应的蛋白集合。459 460 当前策略：461 1. "protein" key 下的 id 直接加入集合462 2. 对于其他类型（phenotype / disease / compound / go / pathway）：463 - (node_type -> protein)464 - (protein -> node_type)465 检索一跳邻居为 protein 的边466 467 Parameters468 ----------469 task_input : dict470 例如：471 {472 "protein": ["P12345"],473 "phenotype": ["HP:0012378"],474 "disease": [],475 "compound": [],476 "go": [],477 "pathway": []478 }479 480 Returns481 -------482 set of str483 映射到的蛋白实体编号集合484 """485 df = self.links_df486 487 normalized_input: Dict[str, List[str]] = {488 (k.lower() if isinstance(k, str) else k): v for k, v in task_input.items()489 }490 491 seed_proteins: Set[str] = set(normalized_input.get("protein", []))492 493 for node_type, ids in normalized_input.items():494 if node_type == "protein":495 continue496 if not ids:497 continue498 499 id_list = list(set(ids))500 501 mask1 = (502 (df["htype"] == node_type)503 & (df["h"].isin(id_list))504 & (df["ttype"] == "protein")505 )506 proteins1 = df.loc[mask1, "t"].tolist()507 508 mask2 = (509 (df["ttype"] == node_type)510 & (df["t"].isin(id_list))511 & (df["htype"] == "protein")512 )513 proteins2 = df.loc[mask2, "h"].tolist()514 515 seed_proteins.update(proteins1)516 seed_proteins.update(proteins2)517 518 return seed_proteins519 520 # ------------------------------------------------------------------521 # 新增部分：天然来源 -> compound / protein 映射522 # ------------------------------------------------------------------523 def _build_compound_to_proteins(self) -> None:524 """525 内部方法：构建 compound -> set(protein) 的缓存字典。526 只在第一次需要时调用一次。527 """528 if self._compound_to_proteins is not None:529 return530 531 df = self.links_df532 comp2prot: Dict[str, Set[str]] = {}533 534 # 情况1：compound 在头、protein 在尾535 mask1 = (df["htype"] == "compound") & (df["ttype"] == "protein")536 for comp_id, sub in df.loc[mask1, ["h", "t"]].groupby("h"):537 comp2prot.setdefault(comp_id, set()).update(sub["t"].tolist())538 539 # 情况2：protein 在头、compound 在尾540 mask2 = (df["ttype"] == "compound") & (df["htype"] == "protein")541 for comp_id, sub in df.loc[mask2, ["t", "h"]].groupby("t"):542 comp2prot.setdefault(comp_id, set()).update(sub["h"].tolist())543 544 self._compound_to_proteins = comp2prot545 def _build_compound_to_protein_edges(self) -> None:546 """547 内部方法：构建 compound -> (rel -> (protein -> max_conf)) 的缓存字典。548 549 说明550 ----551 - 用于天然来源 target 裁剪（按 r 分层、按 conf 阈值等）。552 - 默认忽略 r=HAS_METABOLITE（按你们当前需求不考虑）。553 - 如果同一 (compound, rel, protein) 出现多次，取最大 conf。554 """555 if self._compound_to_protein_edges is not None:556 return557 558 df = self.links_df559 required_cols = {"htype", "ttype", "h", "t", "r", "conf"}560 if not required_cols.issubset(set(df.columns)):561 # 数据列不完整，直接返回空缓存，避免报错562 self._compound_to_protein_edges = {}563 self._num_compounds_in_edges = 0564 return565 566 # 情况1：compound -> protein567 mask1 = (df["htype"] == "compound") & (df["ttype"] == "protein")568 df1 = df.loc[mask1, ["h", "t", "r", "conf"]].rename(569 columns={"h": "comp", "t": "prot"}570 )571 572 # 情况2：protein -> compound（统一转成 compound -> protein）573 mask2 = (df["htype"] == "protein") & (df["ttype"] == "compound")574 df2 = df.loc[mask2, ["t", "h", "r", "conf"]].rename(575 columns={"t": "comp", "h": "prot"}576 )577 578 cp = pd.concat([df1, df2], ignore_index=True)579 if cp.empty:580 self._compound_to_protein_edges = {}581 self._num_compounds_in_edges = 0582 return583 584 # 归一化585 cp["r"] = cp["r"].astype(str).str.upper()586 cp["conf"] = pd.to_numeric(cp["conf"], errors="coerce").fillna(0.0)587 588 # 忽略 HAS_METABOLITE589 cp = cp[cp["r"] != "HAS_METABOLITE"]590 591 # 同一 (comp, r, prot) 取最大 conf592 cp = cp.groupby(["comp", "r", "prot"], as_index=False)["conf"].max()593 594 edges: Dict[str, Dict[str, Dict[str, float]]] = {}595 for comp_id, sub in cp.groupby("comp"):596 comp_key = str(comp_id)597 rel_map: Dict[str, Dict[str, float]] = {}598 for rel, subr in sub.groupby("r"):599 rel_map[str(rel)] = {600 str(pid): float(c)601 for pid, c in zip(subr["prot"].tolist(), subr["conf"].tolist())602 }603 edges[comp_key] = rel_map604 605 self._compound_to_protein_edges = edges606 self._num_compounds_in_edges = len(edges)607 608 def _build_protein_df_compound(self) -> None:609 """610 内部方法：构建 protein 的“全局泛化度”统计：该 protein 出现在多少个 compound 的 compound-protein 边上。611 612 用途613 ----614 - 在不使用任务模块信息的前提下，对“到处都出现的通用蛋白”进行降权（IDF）。615 """616 if self._protein_df_compound is not None and self._num_compounds_in_edges is not None:617 return618 619 self._build_compound_to_protein_edges()620 edges = self._compound_to_protein_edges or {}621 622 df_compound: Dict[str, int] = defaultdict(int)623 for comp_id, rel_map in edges.items():624 # 该 compound 下所有 protein（跨关系去重）625 prots: Set[str] = set()626 for prot_dict in rel_map.values():627 prots.update(prot_dict.keys())628 for p in prots:629 df_compound[p] += 1630 631 self._protein_df_compound = dict(df_compound)632 if self._num_compounds_in_edges is None:633 self._num_compounds_in_edges = len(edges)634 635 def natural_source_target_prefilter(636 self,637 source_id: str,638 *,639 comp_set: Optional[Set[str]] = None,640 max_size: int = -1,641 interaction_conf: float = 0.95,642 comp_buffer_ratio: float = 1.2,643 is_single_chem: Optional[bool] = None,644 verbose: bool = False,645 ) -> Set[str]:646 """647 对单个来源（natural_source / 自定义实体 / 单分子化合物）的 targets 进行“自身信息”裁剪（不使用任务 module 信息）。648 649 说明650 ----651 - 该函数既可用于 natural_source（内部从 natural_sources_df 取 compound 集合），652 也可用于自定义实体/化合物（通过 comp_set 显式传入 compound 集合）。653 - 当 comp_set 只有 1 个 compound 时，默认进入单分子模式（is_single_chem=True），654 排序时更侧重 conf / 特异性 / 结构连通性，而不依赖 support（因为 support≈1）。655 - 会将每次裁剪过程的关键统计写入 self.prefilter_log[source_id]（append，不覆盖）。656 657 裁剪策略（不泄露任务信息）658 --------------------------659 Part1：按关系类型 r 分层（优先保留更像“靶标”的关系）660 Tier0: DRUG_TARGET661 Tier1: + COMPOUND_INHIBITION + COMPOUND_ACTIVATION662 Tier2: + HAS_SUBSTRATE663 Tier3: COMPOUND_INTERACTION（仅作为填充池，且 conf >= interaction_conf）664 HAS_METABOLITE: 永远忽略665 666 Part2：在候选范围内做 top-k（或填充），仅使用：667 - tier（证据层级）668 - support：该 protein 被该来源的多少 compound 支持（多化合物时有效）669 - max_conf：最大 conf670 - idf：protein 的全局“泛化度”惩罚（出现于多少 compound）671 - PPI GCC 结构：components 聚焦（优先保留大连通分量）672 - hub 惩罚：对高 degree 节点轻度降权673 674 Returns675 -------676 set677 裁剪后的 protein 集合（尽量不超过 max_size；允许略小于 max_size）678 """679 run_log: Dict = {680 "source_id": str(source_id),681 "max_size": int(max_size),682 "interaction_conf": float(interaction_conf),683 "comp_buffer_ratio": float(comp_buffer_ratio),684 "is_single_chem": None,685 "comp_set_size": None,686 "tiers": {},687 "case": None,688 "stage": None,689 "ppi": {},690 "components": {},691 "final_size": 0,692 }693 694 if max_size <= 0:695 run_log["final_size"] = 0696 self.prefilter_log.setdefault(str(source_id), []).append(run_log)697 return set()698 699 # 1) 确定 comp_set700 if comp_set is None:701 self._build_natural_source_to_compounds()702 ns2comp = self._natural_source_to_compounds or {}703 comp_set = ns2comp.get(str(source_id), set())704 else:705 comp_set = set(comp_set)706 707 run_log["comp_set_size"] = len(comp_set)708 709 if not comp_set:710 run_log["final_size"] = 0711 self.prefilter_log.setdefault(str(source_id), []).append(run_log)712 return set()713 714 # 2) 单分子模式判定715 if is_single_chem is None:716 is_single_chem = (len(comp_set) == 1)717 run_log["is_single_chem"] = bool(is_single_chem)718 719 # 3) 需要 compound-protein 边（带 r/conf）与全局 df720 self._build_compound_to_protein_edges()721 self._build_protein_df_compound()722 723 edges = self._compound_to_protein_edges or {}724 protein_df_compound = self._protein_df_compound or {}725 N_comp = int(self._num_compounds_in_edges or 0)726 727 # 4) PPI GCC（仅自身网络信息）728 G = self.ppi_gcc729 ppi_nodes = set(G.nodes())730 run_log["ppi"]["gcc_nodes"] = int(G.number_of_nodes())731 run_log["ppi"]["gcc_edges"] = int(G.number_of_edges())732 733 # -------------------------734 # 内部工具：关系收集 + 统计摘要735 # -------------------------736 def _conf_summary(conf_map: Dict[str, float]) -> Dict[str, float]:737 if not conf_map:738 return {"max": 0.0, "mean": 0.0}739 vals = list(conf_map.values())740 return {"max": float(max(vals)), "mean": float(sum(vals) / max(len(vals), 1))}741 742 def _collect_by_rels(rels: Set[str], conf_min: Optional[float] = None):743 prot_max_conf: Dict[str, float] = {}744 prot_support_sets: Dict[str, Set[str]] = defaultdict(set)745 746 for comp_id in comp_set:747 rel_map = edges.get(str(comp_id), {})748 for rel in rels:749 prot_map = rel_map.get(rel, {})750 if not prot_map:751 continue752 for prot_id, conf in prot_map.items():753 cf = float(conf)754 if conf_min is not None and cf < float(conf_min):755 continue756 p = str(prot_id)757 c = str(comp_id)758 prot_support_sets[p].add(c)759 if cf > float(prot_max_conf.get(p, 0.0)):760 prot_max_conf[p] = cf761 762 prots = set(prot_support_sets.keys())763 return prots, prot_max_conf, prot_support_sets764 765 # -------------------------766 # Part1：关系分层767 # -------------------------768 REL_TIER0 = {"DRUG_TARGET"}769 REL_TIER1 = {"COMPOUND_INHIBITION", "COMPOUND_ACTIVATION"}770 REL_TIER2 = {"HAS_SUBSTRATE"}771 REL_INTERACTION = {"COMPOUND_INTERACTION"}772 773 # protein -> 属性（tier/support/conf）774 tier_rank: Dict[str, int] = {}775 max_conf_map: Dict[str, float] = {}776 support_sets_all: Dict[str, Set[str]] = defaultdict(set)777 778 def _merge(prots: Set[str], tier: int, conf_map: Dict[str, float], support_sets: Dict[str, Set[str]]):779 for p in prots:780 if p not in tier_rank:781 tier_rank[p] = tier782 else:783 tier_rank[p] = min(int(tier_rank[p]), int(tier))784 if p in conf_map:785 max_conf_map[p] = max(float(max_conf_map.get(p, 0.0)), float(conf_map[p]))786 if p in support_sets:787 support_sets_all[p].update(support_sets[p])788 789 # Tier0790 t0, t0_conf, t0_supp = _collect_by_rels(REL_TIER0)791 _merge(t0, 0, t0_conf, t0_supp)792 candidates = set(t0)793 run_log["tiers"]["tier0"] = {"count": int(len(t0)), "conf": _conf_summary(t0_conf)}794 795 stage = "tier0"796 797 # Tier1798 if len(candidates) < max_size:799 t1, t1_conf, t1_supp = _collect_by_rels(REL_TIER1)800 _merge(t1, 1, t1_conf, t1_supp)801 candidates |= t1802 run_log["tiers"]["tier1"] = {"count": int(len(t1)), "conf": _conf_summary(t1_conf)}803 stage = "tier0+tier1"804 else:805 run_log["tiers"]["tier1"] = {"count": 0, "conf": {"max": 0.0, "mean": 0.0}}806 807 # Tier2808 if len(candidates) < max_size:809 t2, t2_conf, t2_supp = _collect_by_rels(REL_TIER2)810 _merge(t2, 2, t2_conf, t2_supp)811 candidates |= t2812 run_log["tiers"]["tier2"] = {"count": int(len(t2)), "conf": _conf_summary(t2_conf)}813 stage = "tier0+tier1+tier2"814 else:815 run_log["tiers"]["tier2"] = {"count": 0, "conf": {"max": 0.0, "mean": 0.0}}816 817 core = set(candidates)818 run_log["stage"] = stage819 run_log["tiers"]["core_tier0_2"] = {"count": int(len(core))}820 821 # interaction 作为填充池822 interaction_pool: Set[str] = set()823 if len(core) < max_size:824 ti, ti_conf, ti_supp = _collect_by_rels(REL_INTERACTION, conf_min=interaction_conf)825 interaction_pool = set(ti)826 run_log["tiers"]["tier3_interaction_pool"] = {"count": int(len(interaction_pool)), "conf": _conf_summary(ti_conf)}827 else:828 run_log["tiers"]["tier3_interaction_pool"] = {"count": 0, "conf": {"max": 0.0, "mean": 0.0}}829 830 # -------------------------831 # Part2：top-k / 填充（不使用 module）832 # -------------------------833 def _idf(p: str) -> float:834 if N_comp <= 0:835 return 0.0836 dfc = int(protein_df_compound.get(p, 0))837 return float(math.log((N_comp + 1.0) / (dfc + 1.0)))838 839 def _select_top(pool: Set[str], k: int) -> Set[str]:840 if k <= 0 or not pool:841 return set()842 843 # 1) PPI GCC 映射过滤844 pool_ppi = [p for p in pool if p in ppi_nodes]845 if not pool_ppi:846 return set()847 848 # 2) components 聚焦：优先保留更大的连通分量（减少散点噪声）849 subG = G.subgraph(pool_ppi)850 comps = list(nx.connected_components(subG))851 comps_sorted = sorted(comps, key=len, reverse=True)852 853 comp_size_map: Dict[str, int] = {}854 for comp in comps_sorted:855 cs = len(comp)856 for node in comp:857 comp_size_map[str(node)] = cs858 859 buffer_cap = int(max(k, 1) * float(comp_buffer_ratio))860 kept_components = 0861 if len(pool_ppi) > buffer_cap:862 kept: Set[str] = set()863 for comp in comps_sorted:864 if len(kept) >= buffer_cap:865 break866 kept.update(comp)867 kept_components += 1868 pool_ppi = list(kept)869 else:870 kept_components = len(comps_sorted)871 872 run_log["components"] = {873 "num_components": int(len(comps_sorted)),874 "kept_components": int(kept_components),875 "buffer_cap": int(buffer_cap),876 "pool_ppi_before_sort": int(len(pool_ppi)),877 }878 879 # 3) 排序：根据是否单分子切换模式（单分子时 support 几乎无信息）880 def sort_key_multi(p: str):881 t = int(tier_rank.get(p, 9))882 support = len(support_sets_all.get(p, set()))883 confv = float(max_conf_map.get(p, 0.0))884 idfv = _idf(p)885 deg = int(G.degree(p))886 csz = int(comp_size_map.get(p, 1))887 return (t, -support, -confv, -idfv, deg, -csz, p)888 889 def sort_key_single(p: str):890 t = int(tier_rank.get(p, 9))891 confv = float(max_conf_map.get(p, 0.0))892 idfv = _idf(p)893 deg = int(G.degree(p))894 csz = int(comp_size_map.get(p, 1))895 return (t, -confv, -idfv, deg, -csz, p)896 897 # 记录排序策略898 if is_single_chem:899 run_log["sorting_strategy"] = {900 "mode": "single_compound",901 "factors": ["tier", "max_conf", "idf", "hub_penalty", "component_size"],902 "sort_order": [903 {"factor": "tier", "direction": "asc", "priority": 1, "description": "Tier level (0=highest)"},904 {"factor": "max_conf", "direction": "desc", "priority": 2, "description": "Maximum confidence score"},905 {"factor": "idf", "direction": "desc", "priority": 3, "description": "IDF specificity (lower = more specific)"},906 {"factor": "hub_penalty", "direction": "asc", "priority": 4, "description": "PPI degree (lower = less hub)"},907 {"factor": "component_size", "direction": "desc", "priority": 5, "description": "Connected component size"},908 ],909 "top_k_selected": min(k, len(pool_ppi)),910 }911 else:912 run_log["sorting_strategy"] = {913 "mode": "multi_compound",914 "factors": ["tier", "support", "max_conf", "idf", "hub_penalty", "component_size"],915 "sort_order": [916 {"factor": "tier", "direction": "asc", "priority": 1, "description": "Tier level (0=highest)"},917 {"factor": "support", "direction": "desc", "priority": 2, "description": "Number of compounds supporting this target"},918 {"factor": "max_conf", "direction": "desc", "priority": 3, "description": "Maximum confidence score"},919 {"factor": "idf", "direction": "desc", "priority": 4, "description": "IDF specificity (lower = more generic)"},920 {"factor": "hub_penalty", "direction": "asc", "priority": 5, "description": "PPI degree (lower = less hub)"},921 {"factor": "component_size", "direction": "desc", "priority": 6, "description": "Connected component size"},922 ],923 "top_k_selected": min(k, len(pool_ppi)),924 }925 926 pool_ppi_sorted = sorted(pool_ppi, key=sort_key_single if is_single_chem else sort_key_multi)927 return set(pool_ppi_sorted[: min(k, len(pool_ppi_sorted))])928 929 # 若 core+interaction <= max_size：直接返回（但做 PPI 映射过滤 + 记录）930 if len(core | interaction_pool) <= max_size:931 out = (core | interaction_pool) & ppi_nodes932 run_log["case"] = "direct_return"933 run_log["ppi"]["final_in_gcc"] = int(len(out))934 run_log["final_size"] = int(len(out))935 self.prefilter_log.setdefault(str(source_id), []).append(run_log)936 if verbose:937 print(f"[prefilter] id={source_id} direct_return core={len(core)} pool={len(interaction_pool)} -> {len(out)}")938 return out939 940 # Case A：core 已经 >= max_size 或没有 interaction_pool（对 core top-k）941 if len(core) >= max_size or not interaction_pool:942 selected = _select_top(core, max_size)943 run_log["case"] = "A_core_topk"944 run_log["ppi"]["core_raw"] = int(len(core))945 run_log["ppi"]["core_selected_in_gcc"] = int(len(selected))946 run_log["final_size"] = int(len(selected))947 self.prefilter_log.setdefault(str(source_id), []).append(run_log)948 if verbose:949 print(f"[prefilter] id={source_id} case=A core_raw={len(core)} -> {len(selected)}")950 return selected951 952 # Case B：core 不够，interaction 填充953 core_ppi = set([p for p in core if p in ppi_nodes])954 # merge interaction 信息（仅用于排序）955 ti, ti_conf, ti_supp = _collect_by_rels(REL_INTERACTION, conf_min=interaction_conf)956 _merge(ti, 3, ti_conf, ti_supp)957 958 k_fill = max_size - len(core_ppi)959 fill_selected = _select_top(interaction_pool, k_fill)960 out = core_ppi | fill_selected961 962 run_log["case"] = "B_core_plus_fill"963 run_log["ppi"]["core_ppi"] = int(len(core_ppi))964 run_log["ppi"]["fill_k"] = int(k_fill)965 run_log["ppi"]["fill_selected"] = int(len(fill_selected))966 run_log["final_size"] = int(len(out))967 self.prefilter_log.setdefault(str(source_id), []).append(run_log)968 969 if verbose:970 print(f"[prefilter] id={source_id} case=B core_ppi={len(core_ppi)} fill_k={k_fill} pool={len(interaction_pool)} -> {len(out)}")971 972 return out973 974 975 def _build_natural_source_to_compounds(self) -> None:976 """977 内部方法：构建 natural_source -> set(compound) 的缓存字典。978 基于 natural_sources_df，其中 htype=compound, ttype=natural_source。979 """980 if self._natural_source_to_compounds is not None:981 return982 if self.natural_sources_df is None:983 # 没有天然来源文件，保持空字典984 self._natural_source_to_compounds = {}985 return986 987 ns_df = self.natural_sources_df988 # 只保留 htype=compound, ttype=natural_source 的行（按规范应该全是）989 mask = (ns_df["htype"] == "compound") & (ns_df["ttype"] == "natural_source")990 ns_df = ns_df.loc[mask, ["h", "t"]]991 992 ns2comp: Dict[str, Set[str]] = {}993 for ns_id, sub in ns_df.groupby("t"):994 ns2comp[ns_id] = set(sub["h"].tolist())995 996 self._natural_source_to_compounds = ns2comp997 998 def _build_natural_source_to_proteins(self) -> None:999 """1000 内部方法：组合 compound->protein 和 natural_source->compound，1001 得到 natural_source->protein 的缓存字典。1002 """1003 if self._natural_source_to_proteins is not None:1004 return1005 1006 self._build_compound_to_proteins()1007 self._build_natural_source_to_compounds()1008 1009 comp2prot = self._compound_to_proteins or {}1010 ns2comp = self._natural_source_to_compounds or {}1011 1012 ns2prot: Dict[str, Set[str]] = {}1013 for ns_id, comp_set in ns2comp.items():1014 proteins: Set[str] = set()1015 for comp_id in comp_set:1016 if comp_id in comp2prot:1017 proteins.update(comp2prot[comp_id])1018 if proteins:1019 ns2prot[ns_id] = proteins1020 1021 self._natural_source_to_proteins = ns2prot1022 1023 def get_natural_source_targets(1024 self,1025 natural_source_ids: Optional[Iterable[str]] = None,1026 *,1027 prefilter: bool = True,1028 max_size: int = -1,1029 interaction_conf: float = 0.95,1030 verbose: bool = False,1031 ) -> Dict[str, Set[str]]:1032 """1033 对外接口：获取天然来源（草药/食物）对应的蛋白靶标集合。1034 1035 Parameters1036 ----------1037 natural_source_ids : Iterable[str], optional1038 如果为 None，返回所有 natural_source 的映射；1039 否则仅返回传入 id 子集的映射（存在于图中的）。1040 prefilter : bool, default True1041 是否对超大 target 集进行裁剪（不使用任务 module 信息）。1042 仅当 raw_targets 数量 > max_size 时触发裁剪。1043 max_size : int, default 30001044 触发裁剪的上限，同时也是裁剪后的目标规模上限（实际可能略小于 max_size）。1045 interaction_conf : float, default 0.951046 COMPOUND_INTERACTION 作为填充池时的置信度阈值。1047 verbose : bool, default False1048 是否打印裁剪日志（English prints）。1049 1050 Returns1051 -------1052 dict1053 key: natural_source id1054 value: set of protein id（可能已裁剪）1055 """1056 self._build_natural_source_to_proteins()1057 ns2prot = self._natural_source_to_proteins or {}1058 1059 # 先做 ID 子集过滤，减少不必要的计算1060 if natural_source_ids is None:1061 selected = {k: set(v) for k, v in ns2prot.items()}1062 else:1063 id_set = set(natural_source_ids)1064 selected = {1065 ns_id: set(prot_set)1066 for ns_id, prot_set in ns2prot.items()1067 if ns_id in id_set and prot_set1068 }1069 1070 if not prefilter or max_size < 0:1071 return selected1072 1073 # 只对超大集合触发 prefilter（大多数 natural_source 不受影响）1074 out: Dict[str, Set[str]] = {}1075 for ns_id, prot_set in selected.items():1076 if max_size > 0 and len(prot_set) > int(max_size):1077 out[ns_id] = self.natural_source_target_prefilter(1078 ns_id,1079 max_size=int(max_size),1080 interaction_conf=float(interaction_conf),1081 verbose=bool(verbose),1082 )1083 else:1084 out[ns_id] = set(prot_set)1085 return out1086 1087 def get_targets_from_compounds(1088 self,1089 source_id: str,1090 compound_set: Iterable[str],1091 *,1092 prefilter: bool = True,1093 max_size: int = -1,1094 interaction_conf: float = 0.95,1095 verbose: bool = False,1096 ) -> Set[str]:1097 """给定 compound 集合，返回其对应的 protein 集合（可选 prefilter）。"""1098 self._build_compound_to_proteins()1099 comp2prot = self._compound_to_proteins or {}1100 1101 proteins: Set[str] = set()1102 comp_set = set([str(c) for c in compound_set if c is not None])1103 for cid in comp_set:1104 prots = comp2prot.get(cid)1105 if prots:1106 proteins.update(prots)1107 1108 if not prefilter or max_size <= 0:1109 return proteins1110 1111 if len(proteins) > int(max_size):1112 # 使用同一套 prefilter（不泄露 module），单分子模式由 comp_set 大小自动决定1113 return self.natural_source_target_prefilter(1114 str(source_id),1115 comp_set=comp_set,1116 max_size=int(max_size),1117 interaction_conf=float(interaction_conf),1118 is_single_chem=(len(comp_set) == 1),1119 verbose=bool(verbose),1120 )1121 1122 return proteins1123 1124 def save_prefilter_log(self, out_path: str) -> None:1125 """保存 self.prefilter_log 为 JSON（外部可调用）。"""1126 path = Path(out_path)1127 path.parent.mkdir(parents=True, exist_ok=True)1128 with open(path, "w", encoding="utf-8") as f:1129 json.dump(self.prefilter_log, f, ensure_ascii=False, indent=2)1130 1131 def clear_prefilter_log(self) -> None:1132 """清空 prefilter 日志（外部可调用）。"""1133 self.prefilter_log = {}1134 1135 1136 1137 def get_compound_targets(1138 self,1139 compound_ids: Optional[Iterable[str]] = None,1140 *,1141 prefilter: bool = False,1142 max_size: int = -1,1143 interaction_conf: float = 0.95,1144 verbose: bool = False,1145 ) -> Dict[str, Set[str]]:1146 """1147 获取化合物对应的蛋白靶标集合。1148 1149 兼容说明1150 --------1151 - 默认 prefilter=False，保持旧行为不变。1152 - 当 prefilter=True 且某个化合物 targets 规模过大（>max_size）时，1153 将复用 natural_source_target_prefilter（单分子模式）进行裁剪。1154 1155 Parameters1156 ----------1157 compound_ids : Iterable[str], optional1158 化合物 UCI 列表。如果为 None，返回所有化合物的映射。1159 prefilter : bool, default False1160 是否启用裁剪（仅当 len(targets)>max_size 时触发）。1161 max_size : int, default -11162 裁剪上限（prefilter=True 时生效）。1163 interaction_conf : float, default 0.951164 interaction 作为填充池时的置信度阈值。1165 verbose : bool, default False1166 是否打印裁剪日志（English prints）。1167 1168 Returns1169 -------1170 dict1171 key: compound UCI1172 value: set of protein id（可能已裁剪）1173 """1174 self._build_compound_to_proteins()1175 comp2prot = self._compound_to_proteins or {}1176 1177 if compound_ids is None:1178 # 返回全部（可选 prefilter）1179 out: Dict[str, Set[str]] = {}1180 for cid, prots in comp2prot.items():1181 if not prefilter or max_size <= 0 or len(prots) <= int(max_size):1182 out[str(cid)] = set(prots)1183 else:1184 out[str(cid)] = self.get_targets_from_compounds(1185 source_id=str(cid),1186 compound_set=[str(cid)],1187 prefilter=True,1188 max_size=int(max_size),1189 interaction_conf=float(interaction_conf),1190 verbose=bool(verbose),1191 )1192 return out1193 1194 id_set = set([str(c) for c in compound_ids])1195 out: Dict[str, Set[str]] = {}1196 for cid in id_set:1197 prots = comp2prot.get(cid, set())1198 if not prots:1199 continue1200 if not prefilter or max_size <= 0 or len(prots) <= int(max_size):