hotwa
34102cf459
models/broad_spectrum_predictor.py: ✅ 新增 StrainPrediction dataclass（单个菌株预测结果） ✅ 更新 BroadSpectrumResult 添加 strain_predictions 字段（pandas.DataFrame 类型） ✅ 添加 to_strain_predictions_list() 方法（类型安全转换） ✅ 新增 _prepare_strain_level_predictions() 方法 ✅ 修改 predict_batch() 方法支持 include_strain_predictions 参数 utils/mole_predictor.py: ✅ 添加 include_strain_predictions 参数到所有函数 ✅ 添加命令行参数 --include-strain-predictions ✅ 实现菌株级别数据与聚合结果的合并逻辑 ✅ 更新所有函数签名和文档字符串 2. 测试验证 ✅ 测试基本功能（仅聚合结果）: test_3.csv → 3 行输出 ✅ 测试菌株级别预测功能: test_3.csv → 120 行输出（3 × 40） ✅ 验证输出格式正确性 ✅ 验证每个分子都有完整的 40 个菌株预测 ✅ 验证革兰染色信息正确（18 个阴性菌 + 22 个阳性菌） 3. 文档更新 README.md: ✅ 更新命令行使用示例 ✅ 添加 Python API 使用示例（包含菌株预测） ✅ 添加详细的输出格式说明 ✅ 添加 40 种菌株列表概览 ✅ 添加数据使用场景示例（强化学习、筛选、可视化） Data/mole/README.md: ✅ 新增"菌株级别预测详情"章节 ✅ 完整的 40 种菌株列表（分革兰阴性/阳性） ✅ 数据访问方式示例（CSV 读取、Python API） ✅ 强化学习应用场景（状态表示、奖励函数设计） ✅ 数据可视化代码示例 ✅ 性能和存储建议
661 lines
23 KiB
Python
661 lines
23 KiB
Python
"""
|
||
并行广谱抗菌预测器模块
|
||
|
||
提供高性能的分子广谱抗菌活性预测功能，支持批量处理和多进程并行计算。
|
||
基于MolE分子表示和XGBoost模型进行预测。
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import pickle
|
||
import torch
|
||
import numpy as np
|
||
import pandas as pd
|
||
import multiprocessing as mp
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
from typing import List, Dict, Union, Optional, Tuple, Any
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from scipy.stats.mstats import gmean
|
||
from sklearn.preprocessing import OneHotEncoder
|
||
|
||
from .mole_representation import process_representation
|
||
|
||
|
||
@dataclass
|
||
class PredictionConfig:
|
||
"""预测配置参数"""
|
||
xgboost_model_path: str = None
|
||
mole_model_path: str = None
|
||
strain_categories_path: str = None
|
||
gram_info_path: str = None
|
||
app_threshold: float = 0.04374140128493309
|
||
min_nkill: int = 10
|
||
batch_size: int = 100
|
||
n_workers: Optional[int] = None
|
||
device: str = "auto"
|
||
|
||
def __post_init__(self):
|
||
"""设置默认路径"""
|
||
from pathlib import Path
|
||
|
||
# 获取当前文件所在目录
|
||
current_file = Path(__file__).resolve()
|
||
project_root = current_file.parent.parent # models -> 项目根
|
||
|
||
data_dir = project_root / "Data/mole/pretrained_model/model_ginconcat_btwin_100k_d8000_l0.0001"
|
||
|
||
# 设置所有路径
|
||
if self.mole_model_path is None:
|
||
self.mole_model_path = str(data_dir)
|
||
|
||
if self.xgboost_model_path is None:
|
||
self.xgboost_model_path = str(data_dir / "MolE-XGBoost-08.03.2025_10.17.pkl")
|
||
|
||
if self.strain_categories_path is None:
|
||
self.strain_categories_path = str(data_dir / "maier_screening_results.tsv.gz")
|
||
|
||
if self.gram_info_path is None:
|
||
self.gram_info_path = str(data_dir / "strain_info_SF2.xlsx")
|
||
|
||
@dataclass
|
||
class MoleculeInput:
|
||
"""分子输入数据结构"""
|
||
smiles: str
|
||
chem_id: Optional[str] = None
|
||
|
||
|
||
@dataclass
|
||
class StrainPrediction:
|
||
"""单个菌株的预测结果"""
|
||
pred_id: str # 格式: "chem_id:strain_name"
|
||
chem_id: str
|
||
strain_name: str
|
||
antimicrobial_predictive_probability: float # 预测概率
|
||
no_growth_probability: float # 不生长概率（1 - antimicrobial_predictive_probability）
|
||
growth_inhibition: int # 二值化结果 (0/1)
|
||
gram_stain: Optional[str] = None # 革兰染色类型
|
||
|
||
|
||
@dataclass
|
||
class BroadSpectrumResult:
|
||
"""广谱抗菌预测结果（聚合结果）"""
|
||
chem_id: str
|
||
apscore_total: float
|
||
apscore_gnegative: float
|
||
apscore_gpositive: float
|
||
ginhib_total: int
|
||
ginhib_gnegative: int
|
||
ginhib_gpositive: int
|
||
broad_spectrum: int
|
||
strain_predictions: Optional[pd.DataFrame] = None # 菌株级别预测（40行）
|
||
|
||
def to_dict(self) -> Dict[str, Union[str, float, int]]:
|
||
"""转换为字典格式（仅聚合字段）"""
|
||
return {
|
||
'chem_id': self.chem_id,
|
||
'apscore_total': self.apscore_total,
|
||
'apscore_gnegative': self.apscore_gnegative,
|
||
'apscore_gpositive': self.apscore_gpositive,
|
||
'ginhib_total': self.ginhib_total,
|
||
'ginhib_gnegative': self.ginhib_gnegative,
|
||
'ginhib_gpositive': self.ginhib_gpositive,
|
||
'broad_spectrum': self.broad_spectrum
|
||
}
|
||
|
||
def to_strain_predictions_list(self) -> List['StrainPrediction']:
|
||
"""
|
||
将 DataFrame 转换为 StrainPrediction 列表（用于类型安全场景）
|
||
|
||
Returns:
|
||
StrainPrediction 对象列表
|
||
"""
|
||
if self.strain_predictions is None or self.strain_predictions.empty:
|
||
return []
|
||
|
||
strain_list = []
|
||
for _, row in self.strain_predictions.iterrows():
|
||
strain_pred = StrainPrediction(
|
||
pred_id=row['pred_id'],
|
||
chem_id=row['chem_id'],
|
||
strain_name=row['strain_name'],
|
||
antimicrobial_predictive_probability=row['antimicrobial_predictive_probability'],
|
||
no_growth_probability=row['no_growth_probability'],
|
||
growth_inhibition=row['growth_inhibition'],
|
||
gram_stain=row.get('gram_stain', None)
|
||
)
|
||
strain_list.append(strain_pred)
|
||
|
||
return strain_list
|
||
|
||
|
||
class BroadSpectrumPredictor:
|
||
"""
|
||
广谱抗菌预测器
|
||
|
||
基于MolE分子表示和XGBoost模型预测分子的广谱抗菌活性。
|
||
支持单分子和批量预测，提供详细的抗菌潜力分析。
|
||
"""
|
||
|
||
def __init__(self, config: Optional[PredictionConfig] = None) -> None:
|
||
"""
|
||
初始化预测器
|
||
|
||
Args:
|
||
config: 预测配置参数，如果为None则使用默认配置
|
||
"""
|
||
self.config = config or PredictionConfig()
|
||
self.n_workers = self.config.n_workers or mp.cpu_count()
|
||
|
||
# 验证文件路径
|
||
self._validate_paths()
|
||
|
||
# 预加载共享数据
|
||
self._load_shared_data()
|
||
|
||
def _validate_paths(self) -> None:
|
||
"""验证必要文件路径是否存在"""
|
||
required_files = {
|
||
"mole_model": self.config.mole_model_path,
|
||
"xgboost_model": self.config.xgboost_model_path,
|
||
"strain_categories": self.config.strain_categories_path,
|
||
"gram_info": self.config.gram_info_path,
|
||
}
|
||
|
||
for name, file_path in required_files.items():
|
||
if file_path is None:
|
||
raise ValueError(f"{name} is None! Check __post_init__ configuration")
|
||
if not Path(file_path).exists():
|
||
raise FileNotFoundError(f"Required {name} not found: {file_path}")
|
||
|
||
def _load_shared_data(self) -> None:
|
||
"""加载共享数据（菌株信息、革兰染色信息等）"""
|
||
try:
|
||
# 加载菌株筛选数据
|
||
self.maier_screen: pd.DataFrame = pd.read_csv(
|
||
self.config.strain_categories_path, sep='\t', index_col=0
|
||
)
|
||
|
||
# 准备菌株独热编码
|
||
self.strain_ohe: pd.DataFrame = self._prep_ohe(self.maier_screen.columns)
|
||
|
||
# 加载革兰染色信息
|
||
self.maier_strains: pd.DataFrame = pd.read_excel(
|
||
self.config.gram_info_path,
|
||
skiprows=[0, 1, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54],
|
||
index_col="NT data base"
|
||
)
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to load shared data: {str(e)}")
|
||
|
||
def _prep_ohe(self, categories: pd.Index) -> pd.DataFrame:
|
||
"""
|
||
准备菌株的独热编码
|
||
|
||
Args:
|
||
categories: 菌株类别索引
|
||
|
||
Returns:
|
||
独热编码后的DataFrame
|
||
"""
|
||
try:
|
||
# 新版本 sklearn 使用 sparse_output
|
||
ohe = OneHotEncoder(sparse_output=False)
|
||
except TypeError:
|
||
# 旧版本 sklearn 使用 sparse
|
||
ohe = OneHotEncoder(sparse=False)
|
||
|
||
ohe.fit(pd.DataFrame(categories))
|
||
cat_ohe = pd.DataFrame(
|
||
ohe.transform(pd.DataFrame(categories)),
|
||
columns=categories,
|
||
index=categories
|
||
)
|
||
return cat_ohe
|
||
|
||
def _get_mole_representation(self, molecules: List[MoleculeInput]) -> pd.DataFrame:
|
||
"""
|
||
获取分子的MolE表示
|
||
|
||
Args:
|
||
molecules: 分子输入列表
|
||
|
||
Returns:
|
||
MolE特征表示DataFrame
|
||
"""
|
||
# 准备输入数据
|
||
df_data = []
|
||
for i, mol in enumerate(molecules):
|
||
chem_id = mol.chem_id or f"mol{i+1}"
|
||
df_data.append({"smiles": mol.smiles, "chem_id": chem_id})
|
||
|
||
df = pd.DataFrame(df_data)
|
||
|
||
# 确定设备
|
||
device = self.config.device
|
||
if device == "auto":
|
||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||
|
||
# 获取MolE表示
|
||
return process_representation(
|
||
dataset_path=df,
|
||
smile_column_str="smiles",
|
||
id_column_str="chem_id",
|
||
pretrained_dir=self.config.mole_model_path,
|
||
device=device
|
||
)
|
||
|
||
def _add_strains(self, chemfeats_df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
添加菌株信息到化学特征（笛卡尔积）
|
||
|
||
Args:
|
||
chemfeats_df: 化学特征DataFrame
|
||
|
||
Returns:
|
||
包含菌株信息的特征DataFrame
|
||
"""
|
||
# 准备化学特征
|
||
chemfe = chemfeats_df.reset_index().rename(columns={"index": "chem_id"})
|
||
chemfe["chem_id"] = chemfe["chem_id"].astype(str)
|
||
|
||
# 准备独热编码
|
||
sohe = self.strain_ohe.reset_index().rename(columns={"index": "strain_name"})
|
||
|
||
# 笛卡尔积合并
|
||
xpred = chemfe.merge(sohe, how="cross")
|
||
xpred["pred_id"] = xpred["chem_id"].str.cat(xpred["strain_name"], sep=":")
|
||
|
||
xpred = xpred.set_index("pred_id")
|
||
xpred = xpred.drop(columns=["chem_id", "strain_name"])
|
||
|
||
return xpred
|
||
|
||
def _gram_stain(self, label_df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
添加革兰染色信息
|
||
|
||
Args:
|
||
label_df: 包含菌株名称的DataFrame
|
||
|
||
Returns:
|
||
添加革兰染色信息后的DataFrame
|
||
"""
|
||
df_label = label_df.copy()
|
||
|
||
# 提取NT编号
|
||
df_label["nt_number"] = df_label["strain_name"].apply(
|
||
lambda x: re.search(r".*?\((NT\d+)\)", x).group(1) if re.search(r".*?\((NT\d+)\)", x) else None
|
||
)
|
||
|
||
# 创建革兰染色字典
|
||
gram_dict = self.maier_strains[["Gram stain"]].to_dict()["Gram stain"]
|
||
|
||
# 添加染色信息
|
||
df_label["gram_stain"] = df_label["nt_number"].apply(gram_dict.get)
|
||
|
||
return df_label
|
||
|
||
def _prepare_strain_level_predictions(self, score_df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
准备菌株级别的预测数据
|
||
|
||
Args:
|
||
score_df: 原始预测分数DataFrame，包含 pred_id, 0, 1, growth_inhibition 列
|
||
|
||
Returns:
|
||
格式化的菌株级别预测DataFrame
|
||
"""
|
||
# 创建副本
|
||
strain_df = score_df.copy()
|
||
|
||
# 分离化合物ID和菌株名
|
||
strain_df["chem_id"] = strain_df["pred_id"].str.split(":", expand=True)[0]
|
||
strain_df["strain_name"] = strain_df["pred_id"].str.split(":", expand=True)[1]
|
||
|
||
# 添加革兰染色信息
|
||
strain_df = self._gram_stain(strain_df)
|
||
|
||
# 重命名列为用户友好的名称
|
||
strain_df = strain_df.rename(columns={
|
||
"1": "antimicrobial_predictive_probability",
|
||
"0": "no_growth_probability"
|
||
})
|
||
|
||
# 选择并排序输出列
|
||
output_columns = [
|
||
"pred_id",
|
||
"chem_id",
|
||
"strain_name",
|
||
"antimicrobial_predictive_probability",
|
||
"no_growth_probability",
|
||
"growth_inhibition",
|
||
"gram_stain"
|
||
]
|
||
|
||
return strain_df[output_columns]
|
||
|
||
def _antimicrobial_potential(self, score_df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
计算抗菌潜力分数
|
||
|
||
Args:
|
||
score_df: 预测分数DataFrame
|
||
|
||
Returns:
|
||
聚合后的抗菌潜力DataFrame
|
||
"""
|
||
# 分离化合物ID和菌株名
|
||
score_df["chem_id"] = score_df["pred_id"].str.split(":", expand=True)[0]
|
||
score_df["strain_name"] = score_df["pred_id"].str.split(":", expand=True)[1]
|
||
|
||
# 添加革兰染色信息
|
||
pred_df = self._gram_stain(score_df)
|
||
|
||
# 计算抗菌潜力分数（几何平均数的对数）
|
||
apscore_total = pred_df.groupby("chem_id")["1"].apply(gmean).to_frame().rename(
|
||
columns={"1": "apscore_total"}
|
||
)
|
||
apscore_total["apscore_total"] = np.log(apscore_total["apscore_total"])
|
||
|
||
# 按革兰染色分组的抗菌分数
|
||
apscore_gram = pred_df.groupby(["chem_id", "gram_stain"])["1"].apply(gmean).unstack().rename(
|
||
columns={"negative": "apscore_gnegative", "positive": "apscore_gpositive"}
|
||
)
|
||
apscore_gram["apscore_gnegative"] = np.log(apscore_gram["apscore_gnegative"])
|
||
apscore_gram["apscore_gpositive"] = np.log(apscore_gram["apscore_gpositive"])
|
||
|
||
# 被抑制菌株数统计
|
||
inhibted_total = pred_df.groupby("chem_id")["growth_inhibition"].sum().to_frame().rename(
|
||
columns={"growth_inhibition": "ginhib_total"}
|
||
)
|
||
|
||
# 按革兰染色分组的被抑制菌株数
|
||
inhibted_gram = pred_df.groupby(["chem_id", "gram_stain"])["growth_inhibition"].sum().unstack().rename(
|
||
columns={"negative": "ginhib_gnegative", "positive": "ginhib_gpositive"}
|
||
)
|
||
|
||
# 合并所有结果
|
||
agg_pred = apscore_total.join(apscore_gram).join(inhibted_total).join(inhibted_gram)
|
||
|
||
# 填充NaN值
|
||
agg_pred = agg_pred.fillna(0)
|
||
|
||
return agg_pred
|
||
|
||
|
||
def _predict_batch_worker(batch_data: Tuple[pd.DataFrame, int],
|
||
model_path: str,
|
||
app_threshold: float) -> Tuple[int, pd.DataFrame]:
|
||
"""
|
||
批次预测工作函数（用于多进程）
|
||
|
||
Args:
|
||
batch_data: (特征数据, 批次ID)
|
||
model_path: XGBoost模型路径
|
||
app_threshold: 抑制阈值
|
||
|
||
Returns:
|
||
(批次ID, 预测结果DataFrame)
|
||
"""
|
||
import warnings
|
||
# 忽略所有XGBoost版本相关的警告
|
||
warnings.filterwarnings("ignore", category=UserWarning, module="xgboost")
|
||
X_input, batch_id = batch_data
|
||
|
||
# 加载模型
|
||
with open(model_path, "rb") as file:
|
||
model = pickle.load(file)
|
||
# 修复特征名称兼容性问题
|
||
# 原因：模型使用旧版 XGBoost 保存时，特征列为元组格式（如 "('bacteria_name',)"）
|
||
# 新版 XGBoost 严格检查特征名称匹配，导致预测失败。
|
||
# 解决：清除 XGBoost 内部的特征名称验证，直接使用输入特征进行预测
|
||
# 注意：此操作不改变模型权重和预测逻辑，只禁用格式检查，预测结果保持一致
|
||
if hasattr(model, 'get_booster'):
|
||
model.get_booster().feature_names = None
|
||
|
||
# 进行预测
|
||
y_pred = model.predict_proba(X_input)
|
||
pred_df = pd.DataFrame(y_pred, columns=["0", "1"], index=X_input.index)
|
||
|
||
# 二值化预测结果
|
||
pred_df["growth_inhibition"] = pred_df["1"].apply(
|
||
lambda x: 1 if x >= app_threshold else 0
|
||
)
|
||
|
||
return batch_id, pred_df
|
||
|
||
|
||
class ParallelBroadSpectrumPredictor(BroadSpectrumPredictor):
|
||
"""
|
||
并行广谱抗菌预测器
|
||
|
||
继承自BroadSpectrumPredictor，添加了多进程并行处理能力，
|
||
适用于大规模分子批量预测。
|
||
"""
|
||
|
||
def predict_single(self, molecule: MoleculeInput) -> BroadSpectrumResult:
|
||
"""
|
||
预测单个分子的广谱抗菌活性
|
||
|
||
Args:
|
||
molecule: 分子输入数据
|
||
|
||
Returns:
|
||
广谱抗菌预测结果
|
||
"""
|
||
results = self.predict_batch([molecule])
|
||
return results[0]
|
||
|
||
def predict_batch(self, molecules: List[MoleculeInput],
|
||
include_strain_predictions: bool = False) -> List[BroadSpectrumResult]:
|
||
"""
|
||
批量预测分子的广谱抗菌活性
|
||
|
||
Args:
|
||
molecules: 分子输入列表
|
||
include_strain_predictions: 是否在结果中包含菌株级别预测数据
|
||
|
||
Returns:
|
||
广谱抗菌预测结果列表
|
||
"""
|
||
if not molecules:
|
||
return []
|
||
|
||
# 获取MolE表示
|
||
print(f"Processing {len(molecules)} molecules...")
|
||
mole_representation = self._get_mole_representation(molecules)
|
||
|
||
# 添加菌株信息
|
||
print("Preparing strain-level features...")
|
||
X_input = self._add_strains(mole_representation)
|
||
|
||
# 分批处理
|
||
print(f"Starting parallel prediction with {self.n_workers} workers...")
|
||
batches = []
|
||
for i in range(0, len(X_input), self.config.batch_size):
|
||
batch = X_input.iloc[i:i+self.config.batch_size]
|
||
batches.append((batch, i // self.config.batch_size))
|
||
|
||
# 并行预测
|
||
results = {}
|
||
with ProcessPoolExecutor(max_workers=self.n_workers) as executor:
|
||
futures = {
|
||
executor.submit(_predict_batch_worker, (batch_data, batch_id),
|
||
self.config.xgboost_model_path,
|
||
self.config.app_threshold): batch_id
|
||
for batch_data, batch_id in batches
|
||
}
|
||
|
||
for future in as_completed(futures):
|
||
batch_id, pred_df = future.result()
|
||
results[batch_id] = pred_df
|
||
print(f"Batch {batch_id} completed")
|
||
|
||
# 合并结果
|
||
print("Merging prediction results...")
|
||
all_pred_df = pd.concat([results[i] for i in sorted(results.keys())])
|
||
all_pred_df = all_pred_df.reset_index()
|
||
|
||
# 准备菌株级别数据（如果需要）
|
||
strain_level_data = None
|
||
if include_strain_predictions:
|
||
print("Preparing strain-level predictions...")
|
||
strain_level_data = self._prepare_strain_level_predictions(all_pred_df)
|
||
|
||
# 计算抗菌潜力
|
||
print("Calculating antimicrobial potential scores...")
|
||
agg_df = self._antimicrobial_potential(all_pred_df)
|
||
|
||
# 判断广谱抗菌
|
||
agg_df["broad_spectrum"] = agg_df["ginhib_total"].apply(
|
||
lambda x: 1 if x >= self.config.min_nkill else 0
|
||
)
|
||
|
||
# 转换为结果对象
|
||
results_list = []
|
||
for _, row in agg_df.iterrows():
|
||
# 获取该分子的菌株级别预测
|
||
mol_strain_preds = None
|
||
if strain_level_data is not None:
|
||
mol_strain_preds = strain_level_data[
|
||
strain_level_data['chem_id'] == row.name
|
||
].reset_index(drop=True)
|
||
|
||
result = BroadSpectrumResult(
|
||
chem_id=row.name,
|
||
apscore_total=row["apscore_total"],
|
||
apscore_gnegative=row["apscore_gnegative"],
|
||
apscore_gpositive=row["apscore_gpositive"],
|
||
ginhib_total=int(row["ginhib_total"]),
|
||
ginhib_gnegative=int(row["ginhib_gnegative"]),
|
||
ginhib_gpositive=int(row["ginhib_gpositive"]),
|
||
broad_spectrum=int(row["broad_spectrum"]),
|
||
strain_predictions=mol_strain_preds
|
||
)
|
||
results_list.append(result)
|
||
|
||
return results_list
|
||
|
||
def predict_from_smiles(self,
|
||
smiles_list: List[str],
|
||
chem_ids: Optional[List[str]] = None) -> List[BroadSpectrumResult]:
|
||
"""
|
||
从SMILES字符串列表预测广谱抗菌活性
|
||
|
||
Args:
|
||
smiles_list: SMILES字符串列表
|
||
chem_ids: 化合物ID列表，如果为None则自动生成
|
||
|
||
Returns:
|
||
广谱抗菌预测结果列表
|
||
"""
|
||
if chem_ids is None:
|
||
chem_ids = [f"mol{i+1}" for i in range(len(smiles_list))]
|
||
|
||
if len(smiles_list) != len(chem_ids):
|
||
raise ValueError("smiles_list and chem_ids must have the same length")
|
||
|
||
molecules = [
|
||
MoleculeInput(smiles=smiles, chem_id=chem_id)
|
||
for smiles, chem_id in zip(smiles_list, chem_ids)
|
||
]
|
||
|
||
return self.predict_batch(molecules)
|
||
|
||
def predict_from_file(self,
|
||
file_path: str,
|
||
smiles_column: str = "smiles",
|
||
id_column: str = "chem_id") -> List[BroadSpectrumResult]:
|
||
"""
|
||
从文件预测广谱抗菌活性
|
||
|
||
Args:
|
||
file_path: 输入文件路径（支持CSV/TSV）
|
||
smiles_column: SMILES列名
|
||
id_column: 化合物ID列名
|
||
|
||
Returns:
|
||
广谱抗菌预测结果列表
|
||
"""
|
||
# 读取文件
|
||
if file_path.endswith('.tsv'):
|
||
df = pd.read_csv(file_path, sep='\t')
|
||
else:
|
||
df = pd.read_csv(file_path)
|
||
|
||
# 验证列存在（大小写不敏感）
|
||
columns_lower = {col.lower(): col for col in df.columns}
|
||
|
||
smiles_col_actual = columns_lower.get(smiles_column.lower())
|
||
if smiles_col_actual is None:
|
||
raise ValueError(f"Column '{smiles_column}' not found in file. Available columns: {list(df.columns)}")
|
||
|
||
# 处理ID列
|
||
id_col_actual = columns_lower.get(id_column.lower())
|
||
if id_col_actual is None:
|
||
df[id_column] = [f"mol{i+1}" for i in range(len(df))]
|
||
id_col_actual = id_column
|
||
|
||
# 创建分子输入
|
||
molecules = [
|
||
MoleculeInput(smiles=row[smiles_col_actual], chem_id=str(row[id_col_actual]))
|
||
for _, row in df.iterrows()
|
||
]
|
||
|
||
return self.predict_batch(molecules)
|
||
|
||
|
||
def create_predictor(config: Optional[PredictionConfig] = None) -> ParallelBroadSpectrumPredictor:
|
||
"""
|
||
创建并行广谱抗菌预测器实例
|
||
|
||
Args:
|
||
config: 预测配置参数
|
||
|
||
Returns:
|
||
预测器实例
|
||
"""
|
||
return ParallelBroadSpectrumPredictor(config)
|
||
|
||
|
||
# 便捷函数
|
||
def predict_smiles(smiles_list: List[str],
|
||
chem_ids: Optional[List[str]] = None,
|
||
config: Optional[PredictionConfig] = None) -> List[BroadSpectrumResult]:
|
||
"""
|
||
便捷函数：直接从SMILES列表预测广谱抗菌活性
|
||
|
||
Args:
|
||
smiles_list: SMILES字符串列表
|
||
chem_ids: 化合物ID列表
|
||
config: 预测配置
|
||
|
||
Returns:
|
||
预测结果列表
|
||
"""
|
||
predictor = create_predictor(config)
|
||
return predictor.predict_from_smiles(smiles_list, chem_ids)
|
||
|
||
|
||
def predict_file(file_path: str,
|
||
smiles_column: str = "smiles",
|
||
id_column: str = "chem_id",
|
||
config: Optional[PredictionConfig] = None) -> List[BroadSpectrumResult]:
|
||
"""
|
||
便捷函数：从文件预测广谱抗菌活性
|
||
|
||
Args:
|
||
file_path: 输入文件路径
|
||
smiles_column: SMILES列名
|
||
id_column: ID列名
|
||
config: 预测配置
|
||
|
||
Returns:
|
||
预测结果列表
|
||
"""
|
||
predictor = create_predictor(config)
|
||
return predictor.predict_from_file(file_path, smiles_column, id_column)
|
||
|