ESMFold2 生产级部署工程实践
从结构预测到分子动力学的三层管线

// config.json (ESMFold2)
{
  "architectures": ["ESMFold2Model"],
  "esmc_id": "/path/to/esmc-6b",    // 指向 ESMC 骨干目录
  "lm_d_model": 2560,                // ESMC 表征维度
  "lm_num_layers": 80,               // ESMC Transformer 层数
  "folding_trunk": {
    "n_layers": 48,                  // 折叠主干层数
    "n_heads": 8
  },
  "structure_head": {
    "diffusion_module": {
      "c_atom": 128,                 // 原子级表征维度
      "c_token": 768,                // token 级表征维度
      "token_num_blocks": 12,        // 扩散模块 token 块数
      "inference_num_steps": 14      // 推理扩散步数
    }
  }
}

# 在导入任何 esm 模块之前设置
import os
os.environ['ESMCFOLD_CCD_PATH'] = os.path.join(BASE_DIR, 'ccd.pkl')

# 417 MB 的 ccd.pkl 包含全部 PDB 化学组件
# 配体解析从网络请求变为内存查找，延迟从秒级降到微秒级

def ensure_cuda_memory(min_free_gb: float) -> None:
    if min_free_gb <= 0 or not torch.cuda.is_available():
        return
    free, total = torch.cuda.mem_get_info()
    free_gb = free / 1024**3
    if free_gb < min_free_gb:
        raise RuntimeError(
            f"CUDA free memory is {free_gb:.2f}/{total/1024**3:.2f} GiB, "
            f"below --min-free-gb={min_free_gb:.2f}. Lower parameters, "
            "choose another GPU, or wait for other jobs."
        )

model = ESMFold2Model.from_pretrained(BASE_DIR, torch_dtype=dtype, low_cpu_mem_usage=True)
model = model.to(device).eval()
model.set_chunk_size(chunk_size)  # chunk_size=16 时峰值显存显著降低

try:
    result = fold(model, spi, params, seed)
except torch.cuda.OutOfMemoryError:
    if args.no_oom_retry or device != "cuda":
        raise
    print("CUDA OOM detected. Retrying with safer low-memory parameters...")
    del model
    clear_cuda_cache()  # gc.collect() + empty_cache() + ipc_collect()

    retry_params = FoldParams(
        num_loops=0,                                    # 跳过循环精化
        num_sampling_steps=max(8, min(12, params.num_sampling_steps)),
        num_diffusion_samples=1,
        chunk_size=16,                                  # 最小分块
        dtype=torch.float32,                            # fp32（避免 bf16 的额外激活）
    )
    model = load_model(device, retry_params.dtype, retry_params.chunk_size)
    result = fold(model, spi, retry_params, args.seed)

def print_cuda_memory(prefix: str) -> None:
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    free, total = torch.cuda.mem_get_info()
    print(f"{prefix}: allocated={allocated:.2f} GiB, "
          f"reserved={reserved:.2f} GiB, "
          f"free={free/1024**3:.2f}/{total/1024**3:.2f} GiB")

# 调用点：Before model load → After model load → After prediction

os.environ.setdefault(
    "PYTORCH_CUDA_ALLOC_CONF",
    "expandable_segments:True,max_split_size_mb:128"
)

def default_cpu_threads(reserve_cores: int = 10) -> int:
    return max(1, (os.cpu_count() or 1) - reserve_cores)

# 在导入 torch/numpy 之前设置 BLAS 线程数
DEFAULT_CPU_THREADS = default_cpu_threads()
for env_name in (
    "OMP_NUM_THREADS",      # OpenMP
    "MKL_NUM_THREADS",      # Intel MKL
    "OPENBLAS_NUM_THREADS", # OpenBLAS
    "NUMEXPR_NUM_THREADS",  # NumExpr
    "VECLIB_MAXIMUM_THREADS", # Accelerate
):
    os.environ.setdefault(env_name, str(DEFAULT_CPU_THREADS))

def apply_cpu_limits(args):
    available_cpus = sorted(os.sched_getaffinity(0))
    requested_threads = args.cpu_threads or max(1, len(available_cpus) - args.reserve_cpu_cores)
    cpu_threads = max(1, min(requested_threads, len(available_cpus)))

    # PyTorch 线程池
    torch.set_num_threads(cpu_threads)
    torch.set_num_interop_threads(max(1, min(4, cpu_threads)))

    # OS 级亲和性绑定
    if not args.no_cpu_affinity and hasattr(os, "sched_setaffinity"):
        os.sched_setaffinity(0, set(available_cpus[:cpu_threads]))

@dataclass(frozen=True)
class FoldParams:
    num_loops: int
    num_sampling_steps: int
    num_diffusion_samples: int
    chunk_size: int | None
    dtype: torch.dtype

PRESETS = {
    "low_mem": FoldParams(
        num_loops=1, num_sampling_steps=20,
        num_diffusion_samples=1, chunk_size=16,
        dtype=torch.float32,
    ),
    "balanced": FoldParams(
        num_loops=2, num_sampling_steps=35,
        num_diffusion_samples=1, chunk_size=48,
        dtype=torch.bfloat16,
    ),
    "quality": FoldParams(
        num_loops=3, num_sampling_steps=50,
        num_diffusion_samples=1, chunk_size=64,
        dtype=torch.float32,
    ),
}

# 使用 balanced 预设但提高采样步数
python predict.py --preset balanced --sampling-steps 50

# 使用 low_mem 但切换到 bf16 省显存
python predict.py --preset low_mem --dtype bf16

# 禁用分块（需要足够显存）
python predict.py --preset quality --chunk-size -1

sys.path.insert(0, os.path.join(BASE_DIR, "vendor/esm-c94ed8d"))
sys.path.insert(0, os.path.join(BASE_DIR, "vendor/transformers-main/src"))

from esm.models.esmfold2 import (
    ESMFold2InputBuilder, LigandInput, ProteinInput, StructurePredictionInput,
)
from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model

from Bio.PDB import PDBIO, MMCIFParser, Select

class LigandSelect(Select):
    """选择配体链（chain L）的原子。"""
    def accept_chain(self, chain):
        return chain.get_id() == "L"

def convert_cif_to_pdb(cif_path, pdb_path):
    parser = MMCIFParser(QUIET=True)  # QUIET 抑制格式警告
    structure = parser.get_structure("structure", cif_path)
    io = PDBIO()
    io.set_structure(structure)
    io.save(pdb_path)

def cif_to_pdb_line(cif_line):
    """将 CIF ATOM 行手动转换为 PDB 格式。"""
    parts = cif_line.split()
    atom_id = parts[-1]
    atom_name = parts[12]    # auth_atom
    res_name = parts[9]      # auth_res
    chain = parts[11]        # auth_chain
    res_seq = parts[6]       # seq_id
    x, y, z = float(parts[14]), float(parts[15]), float(parts[16])
    element = parts[1]
    # PDB 固定列宽格式
    return f"ATOM  {atom_id:>5} {atom_name:<4} {res_name:<3} {chain:1} {res_seq:>4}    {x:>8.3f}{y:>8.3f}{z:>8.3f}  1.00  0.00          {element:>2}\n"

from meeko import MoleculePreparation, PDBQTWriterLegacy
from rdkit import Chem
from rdkit.Chem import AllChem

def prepare_pdbqt_meeko(input_pdb, output_pdbqt, molecule_type):
    # 受体：提取 chain A 避免多片段问题
    if molecule_type == 'receptor':
        temp_pdb = str(input_pdb).replace('.pdb', '_chainA.pdb')
        extract_chain_from_pdb(str(input_pdb), temp_pdb, 'A')
        input_pdb = temp_pdb

    mol = Chem.MolFromPDBFile(str(input_pdb), removeHs=False)
    mol = Chem.AddHs(mol)  # 添加显式氢

    # 配体：生成 3D 坐标
    if molecule_type == 'ligand':
        AllChem.EmbedMolecule(mol, randomSeed=42)
        AllChem.MMFFOptimizeMolecule(mol)

    # Meeko 准备：自动分配 AutoDock 原子类型和可旋转键
    preparator = MoleculePreparation()
    mol_setups = preparator.prepare(mol)
    pdbqt_string, _, _ = PDBQTWriterLegacy.write_string(mol_setups[0])

    with open(output_pdbqt, 'w') as f:
        f.write(pdbqt_string)

def calculate_search_box(ligand_coords, padding=10.0):
    """从配体坐标计算搜索盒中心和尺寸。"""
    min_coords = ligand_coords.min(axis=0)
    max_coords = ligand_coords.max(axis=0)

    center = (min_coords + max_coords) / 2
    size = (max_coords - min_coords) + 2 * padding  # 10 Å 填充

    return center, size

receptor = /path/to/receptor.pdbqt
ligand = /path/to/ligand.pdbqt

center_x = -7.900
center_y = 26.137
center_z = 29.241

size_x = 26.480
size_y = 24.311
size_z = 21.468

exhaustiveness = 8    # 搜索彻底度
num_modes = 9         # 输出构象数
energy_range = 3.0    # 能量窗口 (kcal/mol)

def analyze_interactions(protein_atoms, ligand_atoms):
    """基于距离的蛋白质-配体相互作用分析。"""
    interactions = {'hbonds': [], 'hydrophobic': [], 'van_der_waals': []}

    hbond_cutoff = 3.5    # 氢键截止距离 (Å)
    hydrophobic_cutoff = 4.5
    vdw_cutoff = 5.0

    for prot_atom in protein_atoms:
        for lig_atom in ligand_atoms:
            dist = calculate_distance(prot_atom, lig_atom)

            # 氢键：供体-受体对且距离 < 3.5 Å
            if dist < hbond_cutoff:
                if (is_hbond_donor(prot_atom) and is_hbond_acceptor(lig_atom)) or \
                   (is_hbond_donor(lig_atom) and is_hbond_acceptor(prot_atom)):
                    interactions['hbonds'].append({...})

            # 疏水接触：双方均为疏水残基且距离 < 4.5 Å
            if dist < hydrophobic_cutoff and is_hydrophobic(prot_atom) and is_hydrophobic(lig_atom):
                interactions['hydrophobic'].append({...})

            # 范德华接触：距离 < 5.0 Å
            if dist < vdw_cutoff:
                interactions['van_der_waals'].append({...})

    return interactions

def calculate_binding_score(interactions):
    """距离加权的结合亲和力评分。"""
    hbond_weight = -1.5      # kcal/mol per H-bond
    hydrophobic_weight = -0.5
    vdw_weight = -0.1

    score = 0.0
    for hbond in interactions['hbonds']:
        score += hbond_weight * (3.5 / hbond['distance'])  # 距离越短贡献越大
    for contact in interactions['hydrophobic']:
        score += hydrophobic_weight * (4.5 / contact['distance'])
    for contact in interactions['van_der_waals']:
        if contact['distance'] > 0:
            score += vdw_weight * (5.0 / contact['distance'])

    return score  # 负值表示有利结合

{
  "binding_score_kcal_mol": -49.11,
  "num_hbonds": 9,
  "num_hydrophobic": 0,
  "num_vdw_contacts": 220,
  "protein_atoms": 7230,
  "ligand_atoms": 10
}

from openff.toolkit.topology import Molecule
from openmmforcefields.generators import SMIRNOFFTemplateGenerator

# 从 RDKit 构建 OpenFF 分子
ligand_rdkit = Chem.MolFromPDBFile(str(ligand_pdb), removeHs=False, sanitize=True)
ligand_rdkit = Chem.AddHs(ligand_rdkit, addCoords=True)
ligand_off = Molecule.from_rdkit(ligand_rdkit, allow_undefined_stereo=True)

# 注册到力场
forcefield = ForceField("amber14-all.xml", "amber14/tip3p.xml")
ligand_generator = SMIRNOFFTemplateGenerator(
    molecules=[ligand_off],
    forcefield="openff-2.2.1.offxml"
)
forcefield.registerTemplateGenerator(ligand_generator.generator)

# 1. 拆分蛋白和配体
split_protein_and_ligand(input_pdb, protein_pdb, ligand_pdb, ligand_resname)

# 2. PDBFixer 修复蛋白结构
fixer = PDBFixer(filename=str(protein_pdb))
fixer.findMissingResidues()
fixer.missingResidues = {}        # 不添加缺失残基
fixer.findMissingAtoms()
fixer.addMissingAtoms()           # 添加缺失原子
fixer.addMissingHydrogens(pH=7.0) # 添加氢原子

# 3. 合并蛋白和配体
modeller = Modeller(fixer.topology, fixer.positions)
modeller.add(ligand_topology, ligand_positions)

# 4. 溶剂化
modeller.addSolvent(
    forcefield,
    model="tip3p",
    padding=1.0 * unit.nanometer,       # 1 nm 水盒子边界
    ionicStrength=0.15 * unit.molar,    # 生理盐浓度
    positiveIon="Na+",
    negativeIon="Cl-",
)

platform = Platform.getPlatformByName("CUDA")
properties = {
    "DeviceIndex": gpu_id,
    "Precision": "mixed",            # mixed > single（稳定性）
    "UseCpuPme": "false",            # PME 在 GPU 上执行
    "DeterministicForces": "true",   # 确定性力计算（可复现）
}

def add_ligand_restraints(system, topology, positions, k, ligand_resname="LIG"):
    """对配体重原子施加谐和约束。
    k: 力常数 (kJ/mol/nm²)
    返回被约束的原子数。
    """
    if k <= 0:
        return 0

    force = CustomExternalForce("0.5*k*((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
    force.addGlobalParameter("k", k * unit.kilojoule_per_mole / unit.nanometer**2)
    force.addPerParticleParameter("x0")
    force.addPerParticleParameter("y0")
    force.addPerParticleParameter("z0")

    n = 0
    for idx, atom in enumerate(topology.atoms()):
        # 只约束配体重原子，不约束氢原子
        if atom.residue.name == ligand_resname and atom.element.symbol != "H":
            x, y, z = positions[idx].value_in_unit(unit.nanometer)
            force.addParticle(idx, [x, y, z])
            n += 1

    if n > 0:
        system.addForce(force)
    return n

# 阶段一：轻约束生产（让配体不逃逸，但允许局部松弛）
python run_openmm_md_production.py \
  --production-ns 50 \
  --ligand-restraint-k 50 \
  --ligand-restraint-duration-ns 10 \
  --out-dir .../md/production_run_v2

# 阶段二：释放约束继续生产（验证约束下松弛的姿势是否稳定）
python run_openmm_md_production.py \
  --resume \
  --ligand-restraint-k 0 \
  --production-ns 100 \
  --out-dir .../md/production_run_v2

# 检查恢复条件
has_resume_files = (
    args.resume
    and checkpoint_file.exists()
    and eq_state_file.exists()
    and (out_dir / "equilibrated.pdb").exists()
)

if has_resume_files:
    # 从平衡后状态重建模型（跳过整个平衡流程）
    forcefield, modeller, _, _ = prepare_resume_model(args, out_dir)
    platform, properties = make_platform(args)

    # 构建生产系统（包括重新添加约束）
    system_prod = create_base_system(forcefield, modeller.topology)
    system_prod.addForce(MonteCarloBarostat(pressure, temperature, 25))

    if args.ligand_restraint_k > 0:
        eq_state = XmlSerializer.deserialize(eq_state_file.read_text())
        add_ligand_restraints(system_prod, modeller.topology,
                              eq_state.getPositions(),
                              args.ligand_restraint_k)

    integrator_prod = LangevinMiddleIntegrator(temperature, friction, timestep)
    sim_prod = Simulation(modeller.topology, system_prod, integrator_prod, platform, properties)

    # 从检查点恢复（关键步骤）
    try:
        sim_prod.loadCheckpoint(str(checkpoint_file))
        resumed_from_checkpoint = True
    except Exception as exc:
        print(f"Checkpoint load failed ({exc}), falling back to equilibrated state.")
        eq_state = XmlSerializer.deserialize(eq_state_file.read_text())
        sim_prod.context.setPositions(eq_state.getPositions())
        sim_prod.context.setVelocities(eq_state.getVelocities())

payload = {
    "input_pdb": str(args.input_pdb),
    "temperature_k": args.temperature_k,
    "pressure_atm": args.pressure_atm,
    "timestep_fs": args.timestep_fs,
    "production_ns": args.production_ns,
    "production_steps_total": production_steps_total,
    "production_steps_completed": int(sim_prod.currentStep),
    "ligand_restraint_k": args.ligand_restraint_k,
    "ligand_restraint_duration_ns": args.ligand_restraint_duration_ns,
    "n_lig_restrained": n_lig_restrained,
    "resumed_from_checkpoint": resumed_from_checkpoint,
    "elapsed_seconds": elapsed,
}
save_metadata(out_dir / "final_state.json", payload)