""" 配方优化模拟程序 通过迭代式 Grid Search 寻找最优 LNP 配方，最大化目标器官的 Biodistribution。 使用方法: python -m app.optimize --smiles "CC(C)..." --organ liver """ import itertools import json from pathlib import Path from typing import Dict, List, Optional, Tuple from dataclasses import dataclass, field import numpy as np import pandas as pd import torch from torch.utils.data import DataLoader from loguru import logger from tqdm import tqdm import typer from lnp_ml.config import MODELS_DIR from lnp_ml.dataset import ( LNPDataset, LNPDatasetConfig, collate_fn, SMILES_COL, COMP_COLS, HELP_COLS, TARGET_BIODIST, get_phys_cols, get_exp_cols, ) from lnp_ml.modeling.predict import load_model app = typer.Typer() # ============ 参数配置 ============ # 可用的目标器官 AVAILABLE_ORGANS = ["lymph_nodes", "heart", "liver", "spleen", "lung", "kidney", "muscle"] @dataclass class CompRanges: """组分参数范围配置（mol 比例为百分数 0-100）""" # 阳离子脂质/mRNA 重量比 weight_ratio_min: float = 5.0 weight_ratio_max: float = 30.0 # 阳离子脂质 mol 比例 (%) cationic_mol_min: float = 5.0 cationic_mol_max: float = 80.0 # 磷脂 mol 比例 (%) phospholipid_mol_min: float = 0.0 phospholipid_mol_max: float = 80.0 # 胆固醇 mol 比例 (%) cholesterol_mol_min: float = 0.0 cholesterol_mol_max: float = 80.0 # PEG 脂质 mol 比例 (%) peg_mol_min: float = 0.0 peg_mol_max: float = 5.0 def to_dict(self) -> Dict: """转换为字典""" return { "weight_ratio": (self.weight_ratio_min, self.weight_ratio_max), "cationic_mol": (self.cationic_mol_min, self.cationic_mol_max), "phospholipid_mol": (self.phospholipid_mol_min, self.phospholipid_mol_max), "cholesterol_mol": (self.cholesterol_mol_min, self.cholesterol_mol_max), "peg_mol": (self.peg_mol_min, self.peg_mol_max), } def get_validation_error(self) -> Optional[str]: """ 验证范围是否合理，返回错误信息（如果有）。 Returns: 错误信息字符串，如果验证通过则返回 None """ # 检查各范围是否有效（最小值不能大于最大值） if self.weight_ratio_min > self.weight_ratio_max: return f"阳离子脂质/mRNA重量比：最小值({self.weight_ratio_min})不能大于最大值({self.weight_ratio_max})" if self.cationic_mol_min > self.cationic_mol_max: return f"阳离子脂质mol比例：最小值({self.cationic_mol_min})不能大于最大值({self.cationic_mol_max})" if self.phospholipid_mol_min > self.phospholipid_mol_max: return f"磷脂mol比例：最小值({self.phospholipid_mol_min})不能大于最大值({self.phospholipid_mol_max})" if self.cholesterol_mol_min > self.cholesterol_mol_max: return f"胆固醇mol比例：最小值({self.cholesterol_mol_min})不能大于最大值({self.cholesterol_mol_max})" if self.peg_mol_min > self.peg_mol_max: return f"PEG脂质mol比例：最小值({self.peg_mol_min})不能大于最大值({self.peg_mol_max})" # 检查 mol ratio 是否可能加起来为 1 min_sum = self.cationic_mol_min + self.phospholipid_mol_min + self.cholesterol_mol_min + self.peg_mol_min max_sum = self.cationic_mol_max + self.phospholipid_mol_max + self.cholesterol_mol_max + self.peg_mol_max if min_sum > 100.0: return f"mol比例最小值之和({min_sum:.1f}%)超过100%，无法生成有效配方" if max_sum < 100.0: return f"mol比例最大值之和({max_sum:.1f}%)不足100%，无法生成有效配方" return None def validate(self) -> bool: """验证范围是否合理（至少存在一个可行解）""" return self.get_validation_error() is None # 默认组分范围 DEFAULT_COMP_RANGES = CompRanges() # PEG 最小 step size (百分数) MIN_STEP_SIZE = 1 # 迭代策略：mol ratio 的 step_size (百分数) MOL_STEP_SIZES = [10, 2, 1] # 迭代策略：weight ratio 的 step_size（与 mol ratio 解耦） WR_STEP_SIZES = [5, 2, 1] # Helper lipid 选项（不包含 DOTAP） HELPER_LIPID_OPTIONS = ["DOPE", "DSPC"] # Route of administration 选项 ROUTE_OPTIONS = ["intravenous", "intramuscular"] # delivery 统计量（由 preprocess_internal.py 生成） # 包含: mean/std（z-score 逆变换）、qd_min/qd_max（评分归一化） _DELIVERY_STATS_PATH = Path(__file__).resolve().parent / "delivery_zscore_stats.json" if _DELIVERY_STATS_PATH.exists(): with open(_DELIVERY_STATS_PATH) as _f: DELIVERY_ZSCORE_STATS: Dict[str, Dict[str, float]] = json.load(_f) logger.info(f"Loaded delivery stats from {_DELIVERY_STATS_PATH}") else: DELIVERY_ZSCORE_STATS = {} logger.warning(f"delivery_zscore_stats.json not found at {_DELIVERY_STATS_PATH}, " "run 'make preprocess' to generate it") # quantified_delivery 归一化常量（从统计量中提取 qd_min/qd_max，用于评分归一化到 [0,1]） DELIVERY_NORM: Dict[str, Dict[str, float]] = {} for _route, _stats in DELIVERY_ZSCORE_STATS.items(): if "qd_min" in _stats and "qd_max" in _stats: DELIVERY_NORM[_route] = {"min": _stats["qd_min"], "max": _stats["qd_max"]} if not DELIVERY_NORM: logger.warning("DELIVERY_NORM is empty — scoring normalization for delivery will be disabled") @dataclass class ScoringWeights: """ 评分权重配置。 总分 = biodist_score + delivery_score + size_score + ee_score + pdi_score + toxic_score 各项计算方式参见 SCORE.md。 默认值：仅按目标器官 biodistribution 排序（向后兼容）。 """ # 回归任务权重 biodist_weight: float = 1.0 # score = biodist_value * weight delivery_weight: float = 0.0 # score = normalize(delivery, route) * weight size_weight: float = 0.0 # score = (1 if 80<=size<=150 else 0) * weight # 分类任务：per-class 权重（预测为该类时，得分 = 对应权重） ee_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0]) # EE class 0, 1, 2 pdi_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0, 0.0]) # PDI class 0, 1, 2, 3 toxic_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0]) # Toxic class 0, 1 # 默认评分权重（仅按 biodist 排序） DEFAULT_SCORING_WEIGHTS = ScoringWeights() def compute_formulation_score( f: 'Formulation', organ: str, weights: ScoringWeights, ) -> float: """ 计算单个 Formulation 的综合评分。 Args: f: 配方对象 organ: 目标器官 weights: 评分权重 Returns: 综合评分 """ score = 0.0 # 1. biodistribution score += f.get_biodist(organ) * weights.biodist_weight # 2. quantified_delivery（按给药途径归一化到 [0, 1]） if f.quantified_delivery is not None and weights.delivery_weight != 0: norm = DELIVERY_NORM.get(f.route, DELIVERY_NORM["intravenous"]) d_range = norm["max"] - norm["min"] if d_range > 0: delivery_normalized = (f.quantified_delivery - norm["min"]) / d_range delivery_normalized = max(0.0, min(1.0, delivery_normalized)) else: delivery_normalized = 0.0 score += delivery_normalized * weights.delivery_weight # 3. size（60-150nm 为理想范围） if f.size is not None and weights.size_weight != 0: if 60 <= f.size <= 150: score += 1.0 * weights.size_weight # 4. EE 分类 if f.ee_class is not None and 0 <= f.ee_class < len(weights.ee_class_weights): score += weights.ee_class_weights[f.ee_class] # 5. PDI 分类 if f.pdi_class is not None and 0 <= f.pdi_class < len(weights.pdi_class_weights): score += weights.pdi_class_weights[f.pdi_class] # 6. 毒性分类 if f.toxic_class is not None and 0 <= f.toxic_class < len(weights.toxic_class_weights): score += weights.toxic_class_weights[f.toxic_class] return score def compute_df_score( df: pd.DataFrame, organ: str, weights: ScoringWeights, ) -> pd.Series: """ 为 DataFrame 中的所有行计算综合评分。 Args: df: 包含预测结果的 DataFrame organ: 目标器官 weights: 评分权重 Returns: 评分 Series """ score = pd.Series(0.0, index=df.index) # 1. biodistribution pred_col = f"pred_Biodistribution_{organ}" if pred_col in df.columns: score += df[pred_col] * weights.biodist_weight # 2. quantified_delivery（按给药途径归一化） if weights.delivery_weight != 0 and "pred_delivery" in df.columns: for route_name, norm in DELIVERY_NORM.items(): mask = df["_route"] == route_name if mask.any(): d_range = norm["max"] - norm["min"] if d_range > 0: delivery_normalized = (df.loc[mask, "pred_delivery"] - norm["min"]) / d_range delivery_normalized = delivery_normalized.clip(0.0, 1.0) score.loc[mask] += delivery_normalized * weights.delivery_weight # 3. size if weights.size_weight != 0 and "pred_size" in df.columns: size_ok = (df["pred_size"] >= 60) & (df["pred_size"] <= 150) score += size_ok.astype(float) * weights.size_weight # 4. EE 分类 if "pred_ee_class" in df.columns: for cls, w in enumerate(weights.ee_class_weights): if w != 0: score += (df["pred_ee_class"] == cls).astype(float) * w # 5. PDI 分类 if "pred_pdi_class" in df.columns: for cls, w in enumerate(weights.pdi_class_weights): if w != 0: score += (df["pred_pdi_class"] == cls).astype(float) * w # 6. 毒性分类 if "pred_toxic_class" in df.columns: for cls, w in enumerate(weights.toxic_class_weights): if w != 0: score += (df["pred_toxic_class"] == cls).astype(float) * w return score @dataclass class Formulation: """配方数据结构""" # comp token cationic_lipid_to_mrna_ratio: float cationic_lipid_mol_ratio: float phospholipid_mol_ratio: float cholesterol_mol_ratio: float peg_lipid_mol_ratio: float # 离散选项 helper_lipid: str = "DOPE" route: str = "intravenous" # 预测结果（填充后） biodist_predictions: Dict[str, float] = field(default_factory=dict) # 额外预测值 quantified_delivery: Optional[float] = None unnormalized_delivery: Optional[float] = None # 反推的原始递送值（z-score 逆变换） size: Optional[float] = None pdi_class: Optional[int] = None # PDI 分类 (0-3) ee_class: Optional[int] = None # EE 分类 (0-2) toxic_class: Optional[int] = None # 毒性分类 (0: 无毒, 1: 有毒) def to_dict(self) -> Dict: """转换为字典""" return { "Cationic_Lipid_to_mRNA_weight_ratio": self.cationic_lipid_to_mrna_ratio, "Cationic_Lipid_Mol_Ratio": self.cationic_lipid_mol_ratio, "Phospholipid_Mol_Ratio": self.phospholipid_mol_ratio, "Cholesterol_Mol_Ratio": self.cholesterol_mol_ratio, "PEG_Lipid_Mol_Ratio": self.peg_lipid_mol_ratio, "helper_lipid": self.helper_lipid, "route": self.route, } def get_biodist(self, organ: str) -> float: """获取指定器官的 biodistribution 预测值""" col = f"Biodistribution_{organ}" return self.biodist_predictions.get(col, 0.0) def unique_key(self) -> tuple: """生成唯一标识键，用于去重""" return ( round(self.cationic_lipid_to_mrna_ratio, 4), round(self.cationic_lipid_mol_ratio, 4), round(self.phospholipid_mol_ratio, 4), round(self.cholesterol_mol_ratio, 4), round(self.peg_lipid_mol_ratio, 4), self.helper_lipid, self.route, ) def generate_grid_values( center: float, step_size: float, min_val: float, max_val: float, radius: int = 2, ) -> List[float]: """ 围绕中心点生成网格值。 Args: center: 中心值 step_size: 步长 min_val: 最小值 max_val: 最大值 radius: 扩展半径（生成 2*radius+1 个点） Returns: 网格值列表 """ values = [] for i in range(-radius, radius + 1): val = center + i * step_size if min_val <= val <= max_val: values.append(round(val, 4)) return sorted(set(values)) def generate_initial_grid( mol_step: float, wr_step: float, comp_ranges: CompRanges = None, ) -> List[Tuple[float, float, float, float, float]]: """ 生成初始搜索网格（满足 mol ratio 和为 100% 的约束）。 Args: mol_step: mol ratio 搜索步长 (百分数) wr_step: weight ratio 搜索步长 comp_ranges: 组分范围配置（默认使用 DEFAULT_COMP_RANGES） Returns: List of (weight_ratio, cationic_mol, phospholipid_mol, cholesterol_mol, peg_mol) """ if comp_ranges is None: comp_ranges = DEFAULT_COMP_RANGES grid = [] weight_ratios = np.arange( comp_ranges.weight_ratio_min, comp_ranges.weight_ratio_max + 0.001, wr_step ) peg_values = np.arange( comp_ranges.peg_mol_min, comp_ranges.peg_mol_max + 0.001, MIN_STEP_SIZE ) for weight_ratio in weight_ratios: for peg in peg_values: remaining = 100.0 - peg cationic_max = min(comp_ranges.cationic_mol_max, remaining) + 0.001 for cationic_mol in np.arange(comp_ranges.cationic_mol_min, cationic_max, mol_step): phospholipid_max = min(comp_ranges.phospholipid_mol_max, remaining - cationic_mol) + 0.001 for phospholipid_mol in np.arange(comp_ranges.phospholipid_mol_min, phospholipid_max, mol_step): cholesterol_mol = remaining - cationic_mol - phospholipid_mol if (comp_ranges.cholesterol_mol_min <= cholesterol_mol <= comp_ranges.cholesterol_mol_max): grid.append(( round(weight_ratio, 4), round(cationic_mol, 4), round(phospholipid_mol, 4), round(cholesterol_mol, 4), round(peg, 4), )) return grid def generate_refined_grid( seeds: List[Formulation], mol_step: float, wr_step: float, radius: int = 2, comp_ranges: CompRanges = None, ) -> List[Tuple[float, float, float, float, float]]: """ 围绕种子点生成精细化网格。 Args: seeds: 种子配方列表 mol_step: mol ratio 步长 (百分数) wr_step: weight ratio 步长 radius: 扩展半径 comp_ranges: 组分范围配置（默认使用 DEFAULT_COMP_RANGES） Returns: 新的网格点列表 """ if comp_ranges is None: comp_ranges = DEFAULT_COMP_RANGES grid_set = set() for seed in seeds: weight_ratios = generate_grid_values( seed.cationic_lipid_to_mrna_ratio, wr_step, comp_ranges.weight_ratio_min, comp_ranges.weight_ratio_max, radius ) peg_values = generate_grid_values( seed.peg_lipid_mol_ratio, MIN_STEP_SIZE, comp_ranges.peg_mol_min, comp_ranges.peg_mol_max, radius ) cationic_mols = generate_grid_values( seed.cationic_lipid_mol_ratio, mol_step, comp_ranges.cationic_mol_min, comp_ranges.cationic_mol_max, radius ) phospholipid_mols = generate_grid_values( seed.phospholipid_mol_ratio, mol_step, comp_ranges.phospholipid_mol_min, comp_ranges.phospholipid_mol_max, radius ) for weight_ratio in weight_ratios: for peg in peg_values: remaining = 100.0 - peg for cationic_mol in cationic_mols: for phospholipid_mol in phospholipid_mols: cholesterol_mol = remaining - cationic_mol - phospholipid_mol if (comp_ranges.cationic_mol_min <= cationic_mol <= comp_ranges.cationic_mol_max and comp_ranges.phospholipid_mol_min <= phospholipid_mol <= comp_ranges.phospholipid_mol_max and comp_ranges.cholesterol_mol_min <= cholesterol_mol <= comp_ranges.cholesterol_mol_max and comp_ranges.peg_mol_min <= peg <= comp_ranges.peg_mol_max): grid_set.add(( round(weight_ratio, 4), round(cationic_mol, 4), round(phospholipid_mol, 4), round(cholesterol_mol, 4), round(peg, 4), )) return list(grid_set) def create_dataframe_from_formulations( smiles: str, grid: List[Tuple[float, float, float, float, float]], helper_lipids: List[str], routes: List[str], ) -> pd.DataFrame: """ 从配方网格创建 DataFrame。 使用固定的 phys token（Pure, Microfluidic, mRNA, FFL）和 exp token（Mouse, body, luminescence）。 """ rows = [] for comp_values in grid: for helper in helper_lipids: for route in routes: row = { SMILES_COL: smiles, # comp token "Cationic_Lipid_to_mRNA_weight_ratio": comp_values[0], "Cationic_Lipid_Mol_Ratio": comp_values[1], "Phospholipid_Mol_Ratio": comp_values[2], "Cholesterol_Mol_Ratio": comp_values[3], "PEG_Lipid_Mol_Ratio": comp_values[4], # phys token (固定值) "Purity_Pure": 1.0, "Purity_Crude": 0.0, "Mix_type_Microfluidic": 1.0, "Mix_type_Pipetting": 0.0, "Cargo_type_mRNA": 1.0, "Cargo_type_pDNA": 0.0, "Cargo_type_siRNA": 0.0, "Target_or_delivered_gene_FFL": 1.0, "Target_or_delivered_gene_Peptide_barcode": 0.0, "Target_or_delivered_gene_hEPO": 0.0, "Target_or_delivered_gene_FVII": 0.0, "Target_or_delivered_gene_GFP": 0.0, # help token "Helper_lipid_ID_DOPE": 1.0 if helper == "DOPE" else 0.0, "Helper_lipid_ID_DOTAP": 1.0 if helper == "DOTAP" else 0.0, "Helper_lipid_ID_DSPC": 1.0 if helper == "DSPC" else 0.0, "Helper_lipid_ID_MDOA": 0.0, # 不使用 # exp token (固定值) "Model_type_Mouse": 1.0, "Delivery_target_body": 1.0, f"Route_of_administration_{route}": 1.0, "Batch_or_individual_or_barcoded_Individual": 1.0, "Value_name_luminescence": 1.0, # 存储配方元信息 "_helper_lipid": helper, "_route": route, } # 其他 exp token 默认为 0 for col in get_exp_cols(): if col not in row: row[col] = 0.0 rows.append(row) return pd.DataFrame(rows) def predict_all( model: torch.nn.Module, df: pd.DataFrame, device: torch.device, batch_size: int = 256, ) -> pd.DataFrame: """ 使用模型预测所有输出（biodistribution、size、delivery、PDI、EE）。 Returns: 添加了预测列的 DataFrame """ dataset = LNPDataset(df) dataloader = DataLoader( dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn ) all_biodist_preds = [] all_size_preds = [] all_delivery_preds = [] all_pdi_preds = [] all_ee_preds = [] all_toxic_preds = [] with torch.no_grad(): for batch in dataloader: smiles = batch["smiles"] tabular = {k: v.to(device) for k, v in batch["tabular"].items()} outputs = model(smiles, tabular) # biodist 输出是 softmax 后的概率分布 [B, 7] biodist_pred = outputs["biodist"].cpu().numpy() all_biodist_preds.append(biodist_pred) # size 和 delivery 是回归值 all_size_preds.append(outputs["size"].squeeze(-1).cpu().numpy()) all_delivery_preds.append(outputs["delivery"].squeeze(-1).cpu().numpy()) # PDI、EE 和 toxic 是分类，取 argmax all_pdi_preds.append(outputs["pdi"].argmax(dim=-1).cpu().numpy()) all_ee_preds.append(outputs["ee"].argmax(dim=-1).cpu().numpy()) all_toxic_preds.append(outputs["toxic"].argmax(dim=-1).cpu().numpy()) biodist_preds = np.concatenate(all_biodist_preds, axis=0) size_preds = np.concatenate(all_size_preds, axis=0) delivery_preds = np.concatenate(all_delivery_preds, axis=0) pdi_preds = np.concatenate(all_pdi_preds, axis=0) ee_preds = np.concatenate(all_ee_preds, axis=0) toxic_preds = np.concatenate(all_toxic_preds, axis=0) # 添加到 DataFrame for i, col in enumerate(TARGET_BIODIST): df[f"pred_{col}"] = biodist_preds[:, i] # size 模型输出为 log(size)，转换回真实粒径 (nm) df["pred_size"] = np.exp(size_preds) df["pred_delivery"] = delivery_preds df["pred_pdi_class"] = pdi_preds df["pred_ee_class"] = ee_preds df["pred_toxic_class"] = toxic_preds # 反推 unnormalized_delivery: value = z-score * std + mean df["pred_unnorm_delivery"] = np.nan if DELIVERY_ZSCORE_STATS: for route_name, stats in DELIVERY_ZSCORE_STATS.items(): mask = df["_route"] == route_name if mask.any(): df.loc[mask, "pred_unnorm_delivery"] = ( delivery_preds[mask.values] * stats["std"] + stats["mean"] ) return df # 保持向后兼容 def predict_biodist( model: torch.nn.Module, df: pd.DataFrame, device: torch.device, batch_size: int = 256, ) -> pd.DataFrame: """向后兼容的别名""" return predict_all(model, df, device, batch_size) def select_top_k( df: pd.DataFrame, organ: str, k: int = 20, scoring_weights: Optional[ScoringWeights] = None, ) -> List[Formulation]: """ 选择 top-k 配方。 Args: df: 包含预测结果的 DataFrame organ: 目标器官 k: 选择数量 scoring_weights: 评分权重（默认仅按 biodist 排序） Returns: Top-k 配方列表 """ if scoring_weights is None: scoring_weights = DEFAULT_SCORING_WEIGHTS # 计算综合评分并排序 df = df.copy() df["_composite_score"] = compute_df_score(df, organ, scoring_weights) df_sorted = df.sort_values("_composite_score", ascending=False) # 创建配方对象 formulations = [] seen = set() for _, row in df_sorted.iterrows(): key = ( row["Cationic_Lipid_to_mRNA_weight_ratio"], row["Cationic_Lipid_Mol_Ratio"], row["Phospholipid_Mol_Ratio"], row["Cholesterol_Mol_Ratio"], row["PEG_Lipid_Mol_Ratio"], row["_helper_lipid"], row["_route"], ) if key not in seen: seen.add(key) unnorm_val = row.get("pred_unnorm_delivery") unnorm_delivery = float(unnorm_val) if pd.notna(unnorm_val) else None formulation = Formulation( cationic_lipid_to_mrna_ratio=row["Cationic_Lipid_to_mRNA_weight_ratio"], cationic_lipid_mol_ratio=row["Cationic_Lipid_Mol_Ratio"], phospholipid_mol_ratio=row["Phospholipid_Mol_Ratio"], cholesterol_mol_ratio=row["Cholesterol_Mol_Ratio"], peg_lipid_mol_ratio=row["PEG_Lipid_Mol_Ratio"], helper_lipid=row["_helper_lipid"], route=row["_route"], biodist_predictions={ col: row[f"pred_{col}"] for col in TARGET_BIODIST }, # 额外预测值 quantified_delivery=row.get("pred_delivery"), unnormalized_delivery=unnorm_delivery, size=row.get("pred_size"), pdi_class=int(row.get("pred_pdi_class")) if row.get("pred_pdi_class") is not None else None, ee_class=int(row.get("pred_ee_class")) if row.get("pred_ee_class") is not None else None, toxic_class=int(row.get("pred_toxic_class")) if row.get("pred_toxic_class") is not None else None, ) formulations.append(formulation) if len(formulations) >= k: break return formulations def generate_single_seed_grid( seed: Formulation, mol_step: float, wr_step: float, radius: int = 2, comp_ranges: CompRanges = None, ) -> List[Tuple[float, float, float, float, float]]: """ 为单个种子点生成邻域网格。 Args: seed: 种子配方 mol_step: mol ratio 步长 (百分数) wr_step: weight ratio 步长 radius: 扩展半径 comp_ranges: 组分范围配置（默认使用 DEFAULT_COMP_RANGES） Returns: 网格点列表 """ if comp_ranges is None: comp_ranges = DEFAULT_COMP_RANGES grid_set = set() weight_ratios = generate_grid_values( seed.cationic_lipid_to_mrna_ratio, wr_step, comp_ranges.weight_ratio_min, comp_ranges.weight_ratio_max, radius ) peg_values = generate_grid_values( seed.peg_lipid_mol_ratio, MIN_STEP_SIZE, comp_ranges.peg_mol_min, comp_ranges.peg_mol_max, radius ) cationic_mols = generate_grid_values( seed.cationic_lipid_mol_ratio, mol_step, comp_ranges.cationic_mol_min, comp_ranges.cationic_mol_max, radius ) phospholipid_mols = generate_grid_values( seed.phospholipid_mol_ratio, mol_step, comp_ranges.phospholipid_mol_min, comp_ranges.phospholipid_mol_max, radius ) for weight_ratio in weight_ratios: for peg in peg_values: remaining = 100.0 - peg for cationic_mol in cationic_mols: for phospholipid_mol in phospholipid_mols: cholesterol_mol = remaining - cationic_mol - phospholipid_mol if (comp_ranges.cationic_mol_min <= cationic_mol <= comp_ranges.cationic_mol_max and comp_ranges.phospholipid_mol_min <= phospholipid_mol <= comp_ranges.phospholipid_mol_max and comp_ranges.cholesterol_mol_min <= cholesterol_mol <= comp_ranges.cholesterol_mol_max and comp_ranges.peg_mol_min <= peg <= comp_ranges.peg_mol_max): grid_set.add(( round(weight_ratio, 4), round(cationic_mol, 4), round(phospholipid_mol, 4), round(cholesterol_mol, 4), round(peg, 4), )) return list(grid_set) def optimize( smiles: str, organ: str, model: torch.nn.Module, device: torch.device, top_k: int = 20, num_seeds: Optional[int] = None, top_per_seed: int = 1, step_sizes: Optional[List[float]] = None, wr_step_sizes: Optional[List[float]] = None, comp_ranges: Optional[CompRanges] = None, routes: Optional[List[str]] = None, scoring_weights: Optional[ScoringWeights] = None, batch_size: int = 256, ) -> List[Formulation]: """ 执行配方优化（层级搜索策略）。 采用层级搜索策略： 1. 第一次迭代：全局稀疏搜索，选择 top num_seeds 个分散的种子点 2. 后续迭代：对每个种子点分别在其邻域内搜索，各自保留 top_per_seed 个局部最优 3. 这样可以保持搜索的多样性，避免结果集中在单一区域 Args: smiles: SMILES 字符串 organ: 目标器官 model: 训练好的模型 device: 计算设备 top_k: 最终返回的最优配方数 num_seeds: 第一次迭代后保留的种子点数量（默认为 top_k * 5） top_per_seed: 每个种子点的邻域搜索后保留的局部最优点数量 step_sizes: mol ratio 每轮迭代的步长列表 (百分数，默认 [10, 2, 1]) wr_step_sizes: weight ratio 每轮迭代的步长列表 (默认 [5, 2, 1]) comp_ranges: 组分范围配置（默认使用 DEFAULT_COMP_RANGES） routes: 给药途径列表（默认使用 ROUTE_OPTIONS） scoring_weights: 评分权重配置（默认仅按 biodist 排序） batch_size: 预测批次大小 Returns: 最终 top-k 配方列表 """ if num_seeds is None: num_seeds = top_k * 5 if step_sizes is None: step_sizes = MOL_STEP_SIZES if wr_step_sizes is None: wr_step_sizes = WR_STEP_SIZES # 两组步长长度必须一致 if len(wr_step_sizes) != len(step_sizes): raise ValueError( f"step_sizes ({len(step_sizes)}) 和 wr_step_sizes ({len(wr_step_sizes)}) 长度不一致" ) if comp_ranges is None: comp_ranges = DEFAULT_COMP_RANGES if routes is None: routes = ROUTE_OPTIONS if scoring_weights is None: scoring_weights = DEFAULT_SCORING_WEIGHTS def _score(f: Formulation) -> float: return compute_formulation_score(f, organ, scoring_weights) logger.info(f"Starting optimization for organ: {organ}") logger.info(f"SMILES: {smiles}") logger.info(f"Strategy: num_seeds={num_seeds}, top_per_seed={top_per_seed}, top_k={top_k}") logger.info(f"Mol step sizes: {step_sizes}, WR step sizes: {wr_step_sizes}") logger.info(f"Routes: {routes}") logger.info(f"Scoring weights: biodist={scoring_weights.biodist_weight}, delivery={scoring_weights.delivery_weight}, size={scoring_weights.size_weight}") logger.info(f"Comp ranges: {comp_ranges.to_dict()}") seeds = None for iteration, (mol_step, wr_step) in enumerate(zip(step_sizes, wr_step_sizes)): logger.info(f"\n{'='*60}") logger.info(f"Iteration {iteration + 1}/{len(step_sizes)}, mol_step={mol_step}, wr_step={wr_step}") logger.info(f"{'='*60}") if seeds is None: # ==================== 第一次迭代：全局稀疏搜索 ==================== logger.info("Generating initial grid (global sparse search)...") grid = generate_initial_grid(mol_step, wr_step, comp_ranges) logger.info(f"Grid size: {len(grid)} comp combinations") # 扩展到所有 helper lipid 和 route 组合 total_combinations = len(grid) * len(HELPER_LIPID_OPTIONS) * len(routes) logger.info(f"Total combinations: {total_combinations}") # 创建 DataFrame df = create_dataframe_from_formulations( smiles, grid, HELPER_LIPID_OPTIONS, routes ) # 预测 logger.info("Running predictions...") df = predict_biodist(model, df, device, batch_size) # 选择 top num_seeds 个种子点 seeds = select_top_k(df, organ, num_seeds, scoring_weights) logger.info(f"Selected {len(seeds)} seeds for next iteration") else: # ==================== 后续迭代：层级局部搜索 ==================== # 对每个种子点分别搜索，各自保留局部最优 logger.info(f"Hierarchical local search around {len(seeds)} seeds...") all_local_best = [] for seed_idx, seed in enumerate(seeds): local_grid = generate_single_seed_grid(seed, mol_step, wr_step, radius=2, comp_ranges=comp_ranges) if len(local_grid) == 0: # 如果没有新的网格点，保留原种子 all_local_best.append(seed) continue # 创建 DataFrame df = create_dataframe_from_formulations( smiles, local_grid, [seed.helper_lipid], [seed.route] ) # 预测 df = predict_biodist(model, df, device, batch_size) # 选择该种子邻域内的 top top_per_seed 个局部最优 local_top = select_top_k(df, organ, top_per_seed, scoring_weights) all_local_best.extend(local_top) if seed_idx == 0 or (seed_idx + 1) % 5 == 0: logger.info(f" Seed {seed_idx + 1}/{len(seeds)}: local grid size={len(local_grid)}, " f"local best score={_score(local_top[0]):.4f}") # 更新种子为所有局部最优点（去重） seen_keys = set() unique_local_best = [] # 先按综合评分排序，确保保留最优的 all_local_best_sorted = sorted(all_local_best, key=_score, reverse=True) for f in all_local_best_sorted: key = f.unique_key() if key not in seen_keys: seen_keys.add(key) unique_local_best.append(f) seeds = unique_local_best logger.info(f"Collected {len(seeds)} unique local best formulations (from {len(all_local_best)} candidates)") # 显示当前最优 best = max(seeds, key=_score) logger.info(f"Current best score: {_score(best):.4f} (biodist_{organ}={best.get_biodist(organ):.4f})") logger.info(f"Best formulation: {best.to_dict()}") # 最终去重、按综合评分排序并返回 top_k seeds_sorted = sorted(seeds, key=_score, reverse=True) # 去重：保留每个唯一配方中得分最高的（已排序，所以第一个出现的就是最高的） seen_keys = set() unique_results = [] for f in seeds_sorted: key = f.unique_key() if key not in seen_keys: seen_keys.add(key) unique_results.append(f) logger.info(f"Final results: {len(unique_results)} unique formulations (from {len(seeds)} candidates)") return unique_results[:top_k] def format_results(formulations: List[Formulation], organ: str) -> pd.DataFrame: """格式化结果为 DataFrame""" rows = [] for i, f in enumerate(formulations): row = { "rank": i + 1, f"Biodistribution_{organ}": f.get_biodist(organ), **f.to_dict(), } # 添加其他器官的预测 for col in TARGET_BIODIST: if col != f"Biodistribution_{organ}": row[col] = f.biodist_predictions.get(col, 0.0) rows.append(row) return pd.DataFrame(rows) @app.command() def main( smiles: str = typer.Option(..., "--smiles", "-s", help="Cationic lipid SMILES string"), organ: str = typer.Option(..., "--organ", "-o", help=f"Target organ: {AVAILABLE_ORGANS}"), model_path: Path = typer.Option( MODELS_DIR / "final" / "model.pt", "--model", "-m", help="Path to trained model checkpoint" ), output_path: Optional[Path] = typer.Option( None, "--output", "-O", help="Output CSV path (optional)" ), top_k: int = typer.Option(20, "--top-k", "-k", help="Number of top formulations to return"), num_seeds: Optional[int] = typer.Option(None, "--num-seeds", "-n", help="Number of seed points from first iteration (default: top_k * 5)"), top_per_seed: int = typer.Option(1, "--top-per-seed", "-t", help="Number of local best to keep per seed"), step_sizes: Optional[str] = typer.Option(None, "--step-sizes", "-S", help="Mol ratio step sizes, comma-separated (e.g., '10,2,1')"), wr_step_sizes: Optional[str] = typer.Option(None, "--wr-step-sizes", help="Weight ratio step sizes, comma-separated (e.g., '5,2,1')"), batch_size: int = typer.Option(256, "--batch-size", "-b", help="Prediction batch size"), device: str = typer.Option("cuda" if torch.cuda.is_available() else "cpu", "--device", "-d", help="Device"), ): """ 配方优化程序：寻找最大化目标器官 Biodistribution 的最优 LNP 配方。 采用层级搜索策略： 1. 第一次迭代：全局稀疏搜索，选择 top num_seeds 个分散的种子点 2. 后续迭代：对每个种子点分别在其邻域内搜索，各自保留 top_per_seed 个局部最优 3. 这样可以保持搜索的多样性，避免结果集中在单一区域 示例: python -m app.optimize --smiles "CC(C)..." --organ liver python -m app.optimize -s "CC(C)..." -o spleen -k 10 -n 30 -t 2 python -m app.optimize -s "CC(C)..." -o liver -S "10,2,1" --wr-step-sizes "5,2,1" """ # 验证器官 if organ not in AVAILABLE_ORGANS: logger.error(f"Invalid organ: {organ}. Available: {AVAILABLE_ORGANS}") raise typer.Exit(1) # 解析步长 parsed_step_sizes = None if step_sizes: try: parsed_step_sizes = [float(s.strip()) for s in step_sizes.split(",")] except ValueError: logger.error(f"Invalid step sizes format: {step_sizes}") raise typer.Exit(1) parsed_wr_step_sizes = None if wr_step_sizes: try: parsed_wr_step_sizes = [float(s.strip()) for s in wr_step_sizes.split(",")] except ValueError: logger.error(f"Invalid wr step sizes format: {wr_step_sizes}") raise typer.Exit(1) # 加载模型 logger.info(f"Loading model from {model_path}") device = torch.device(device) model = load_model(model_path, device) results = optimize( smiles=smiles, organ=organ, model=model, device=device, top_k=top_k, num_seeds=num_seeds, top_per_seed=top_per_seed, step_sizes=parsed_step_sizes, wr_step_sizes=parsed_wr_step_sizes, batch_size=batch_size, ) # 格式化并显示结果 df_results = format_results(results, organ) logger.info(f"\n{'='*60}") logger.info(f"TOP {top_k} FORMULATIONS FOR {organ.upper()}") logger.info(f"{'='*60}") print(df_results.to_string(index=False)) # 保存结果 if output_path: df_results.to_csv(output_path, index=False) logger.success(f"Results saved to {output_path}") return df_results if __name__ == "__main__": app()