蛋白质折叠模拟跟踪 10 万个原子的运动轨迹——每个原子都受到周围原子的范德华力、静电力和键结力。牛顿第二定律每个时间步1fs 10^-15 秒计算所有原子对的作用力 → 更新速度 → 更新位置。100,000 个原子 → 50 亿对相互作用/步 → 1μs 的模拟需要 10 亿步 → 传统 CPU 集群跑一个月。但非键结作用的截断半径通常只有 10-12Å1.2nm在 10 万个原子中每个原子只与最近的 200-500 个邻居有实际相互作用——稀疏性是把模拟搬上 NPU 的关键。mat-chem-sim-pred 仓库用 Verlet 邻接列表空间哈希 截断半径把 50 亿对缩减到 2000 万对再用 NPU 的 Cube 单元并行计算所有作用力。Verlet 邻接列表——空间哈希桶# mat-chem-sim-pred/md/verlet_list.py## Verlet 邻接列表: 为每个原子维护截断半径内的邻居索引# 核心: 三维空间哈希 → 每个原子只检查相邻 27 个桶## 空间哈希桶: 边长 cut_off12Å# 网格: (x // cut_off, y // cut_off, z // cut_off)# 邻域: 26 个相邻桶 自身桶 27 桶importtorchimporttorch_npuclassVerletNeighborList: 基于空间哈希的 Verlet 邻接列表 NPU 加速: 1. 哈希桶分配: scatter argsort O(N log N)在 NPU 上用并行排序 2. 邻域搜索: 对每个桶查 27 个邻域桶 → tensordot 批量计算距离 3. 截断: 距离 cut_off 的原子对直接 mask def__init__(self,cut_off12.0,skin2.0): cut_off: 截断半径 (Å)非键结作用力在此距离外可忽略 skin: 缓冲区 (Å)更新邻接列表的频率控制 邻接列表包含 cut_off skin 范围内的原子 更新频率 skin / (max_velocity * dt) self.cut_offcut_off# 12Åself.skinskin# 2Å → 邻接列表实际半径 14Åself.cell_sizecut_offskin# 哈希桶边长 14Ådefbuild(self,positions,box_size): 构建邻接列表 positions: [N, 3] 原子坐标 (Å) box_size: [3] 模拟盒子大小周期边界 returns: neighbors: [N, max_neighbors] 邻居索引-1 表示空位 n_neighbors: [N] 每原子的邻居数 Npositions.shape[0]devicepositions.device# Step 1: 三维空间哈希每个原子的哈希桶cell_idx(positions/self.cell_size).long()# [N, 3]cell_hash(cell_idx[:,0]*73856093cell_idx[:,1]*19349663cell_idx[:,2]*83492791)%(2**31-1)# [N]# 按哈希值排序相同哈希 同一桶sorted_hash,sort_idxcell_hash.sort()sorted_positionspositions[sort_idx]# [N, 3]# 桶边界: 找到每个哈希值的起始位置hash_diffsorted_hash[1:]-sorted_hash[:-1]# [N-1]bucket_startstorch.cat([torch.zeros(1,dtypetorch.long,devicedevice),(hash_diff!0).nonzero(as_tupleTrue)[0]1,torch.tensor([N],dtypetorch.long,devicedevice)])# [n_buckets 1]num_bucketslen(bucket_starts)-1# Step 2: 为每个桶构建 27 个邻域偏移# 偏移: [-1,0,1]³ 27 种grid_offsetstorch.tensor([[dx,dy,dz]fordxin[-1,0,1]fordyin[-1,0,1]fordzin[-1,0,1]],dtypetorch.long,devicedevice)# [27, 3]# Step 3: 收集所有候选邻居对# 最大候选数估算: N * atoms_per_bucket * 27max_atoms_per_bucket100# 典型截断半径球内 200-500每桶 ~50-150max_candidatesN*max_atoms_per_bucket*27pair_itorch.zeros(max_candidates,dtypetorch.long,devicedevice)pair_jtorch.zeros(max_candidates,dtypetorch.long,devicedevice)n_pairs0forbidinrange(num_buckets):b_startbucket_starts[bid].item()b_endbucket_starts[b_start1].item()ifb_start1len(bucket_starts)elseN atoms_in_bucketb_end-b_startifatoms_in_bucket0:continueatom_indicessort_idx[b_start:b_end]# 桶内原子原始索引# 当前桶的中心坐标 (cx, cy, cz)first_atom_idxsort_idx[b_start]cxcell_idx[first_atom_idx]# 检查 27 个邻域桶foroffsetingrid_offsets:neighbor_cellcxoffset# 周期边界neighbor_cellself._wrap_cell(neighbor_cell,box_size/self.cell_size)# 查找邻域桶的哈希neighbor_hash(int(neighbor_cell[0].item())*73856093int(neighbor_cell[1].item())*19349663int(neighbor_cell[2].item())*83492791)%(2**31-1)# 二分查找邻域桶neighbor_bidself._find_bucket(sorted_hash,bucket_starts,neighbor_hash)ifneighbor_bid-1:continuenb_startbucket_starts[neighbor_bid].item()nb_endbucket_starts[neighbor_bid1].item()ifneighbor_bid1num_bucketselseN# 对当前桶中所有原子 × 邻域桶中所有原子 → 候选对neighbor_atomssort_idx[nb_start:nb_end]n_curratoms_in_bucket n_neighnb_end-nb_start iiatom_indices.unsqueeze(1).expand(-1,n_neigh).reshape(-1)jjneighbor_atoms.unsqueeze(0).expand(n_curr,-1).reshape(-1)n_newlen(ii)ifn_pairsn_newmax_candidates:breakpair_i[n_pairs:n_pairsn_new]ii pair_j[n_pairs:n_pairsn_new]jj n_pairsn_new# Step 4: 距离筛选截断 cut_off skinpair_ipair_i[:n_pairs]pair_jpair_j[:n_pairs]# 排除自身配对self_maskpair_i!pair_j pair_ipair_i[self_mask]pair_jpair_j[self_mask]# 计算所有候选对的距离pos_ipositions[pair_i]# [n_valid, 3]pos_jpositions[pair_j]# 周期边界下的最小镜像距离deltapos_i-pos_j deltadelta-box_size*torch.round(delta/box_size)distancesdelta.norm(dim1)# [n_valid]# 截断: 距离 cut_off skinwithin_rangedistances(self.cut_offself.skin)pair_ipair_i[within_range]pair_jpair_j[within_range]returnpair_i,pair_j# [n_pairs], [n_pairs]def_wrap_cell(self,cell_idx,n_cells):周期边界包裹returncell_idx%n_cells.long()def_find_bucket(self,sorted_hashes,bucket_starts,target_hash):二分查找哈希值对应的桶索引# 简化版: 线性搜索foriinrange(len(bucket_starts)-1):startbucket_starts[i].item()ifsorted_hashes[start].item()target_hash:returnireturn-1Lennard-Jones 势——并行非键结作用力# mat-chem-sim-pred/md/force_field.py## 力场计算: 对邻接列表中的每对原子计算作用力# Lennard-Jones 12-6: U(r) 4ε[(σ/r)^12 - (σ/r)^6]# 力: F -∇U(r)classLennardJonesForce: LJ 12-6 势的并行力计算 NPU: 所有原子对同时计算距离 → 力 → 归约到每个原子 def__init__(self,epsilon0.238,sigma3.4,cut_off12.0): epsilon: 势阱深度 (kcal/mol)Ar 原子的 LJ 参数 sigma: 零点距离 (Å)U(σ) 0 self.epsilonepsilon# kcal/molself.sigmasigma# Åself.cut_offcut_off# Å# 截断修正: LJ 在 cut_off 处应为零Weeks-Chandler-Andersen 移位self.U_cut4*epsilon*((sigma/cut_off)**12-(sigma/cut_off)**6)defcompute_forces(self,positions,pair_i,pair_j,box_size): 计算所有原子对之间的 LJ 力 力矢量: F_ij dU/dr * (r_ij / |r_ij|) 其中: dU/dr 24ε/r * [2(σ/r)^12 - (σ/r)^6] F_shift 作用力-反作用力: F_i F_ij, F_j - F_ij Npositions.shape[0]n_pairslen(pair_i)# 原子对的位置pipositions[pair_i]# [n_pairs, 3]pjpositions[pair_j]# [n_pairs, 3]# 最小镜像距离矢量deltapi-pj# [n_pairs, 3]deltadelta-box_size.unsqueeze(0)*torch.round(delta/box_size.unsqueeze(0))# 距离rdelta.norm(dim1)# [n_pairs]# LJ 力的大小: F(r) -dU/dr# dU/dr 24ε * (1/r) * [2(σ/r)^12 - (σ/r)^6]sigma_rself.sigma/r# [n_pairs]sigma_r_6sigma_r**6sigma_r_12sigma_r_6**2# dU/dr -24ε/r * (2*sigma_r_12 - sigma_r_6)# 负号来自 F -dU/drforce_magnitude24.0*self.epsilon/r*(2.0*sigma_r_12-sigma_r_6)# [n_pairs]# 截断附近平滑衰减shift function# 避免截断处力不连续导致能量漂移r_cself.cut_off shift_maskr(r_c-1.0)# 截断前 1Å 平滑ifshift_mask.any():# 三次样条平滑: F(r) F_original * S(r/r_c)x(r_c-r[shift_mask])shift_factor-20.0*x**770.0*x**6-84.0*x**535.0*x**4force_magnitude[shift_mask]*shift_factor# 方向单位矢量directiondelta/r.unsqueeze(1)# [n_pairs, 3]# 力矢量F_ijforce_magnitude.unsqueeze(1)*direction# [n_pairs, 3]# 归约: F_i F_ij, F_j - F_ijforcestorch.zeros(N,3,dtypetorch.float32,devicepositions.device)# scatter_add: 将 F_ij 加到 atoms i, 将 -F_ij 加到 atoms jforces.index_add_(0,pair_i,F_ij)forces.index_add_(0,pair_j,-F_ij)returnforcesdefcompute_energy(self,positions,pair_i,pair_j,box_size):计算系统的 LJ 势能pipositions[pair_i]pjpositions[pair_j]deltapi-pj deltadelta-box_size.unsqueeze(0)*torch.round(delta/box_size.unsqueeze(0))rdelta.norm(dim1)sigma_rself.sigma/r sigma_r_6sigma_r**6sigma_r_12sigma_r_6**2U_pair4.0*self.epsilon*(sigma_r_12-sigma_r_6)returnU_pair.sum()Velocity-Verlet 积分——位置和速度的交替推进# mat-chem-sim-pred/md/integrator.py## Velocity-Verlet: 时间步 dt 1fs (10^-15 s)# 1. v(t dt/2) v(t) a(t) * dt/2# 2. x(t dt) x(t) v(t dt/2) * dt# 3. a(t dt) F(x(tdt)) / m# 4. v(t dt) v(t dt/2) a(t dt) * dt/2classVelocityVerletIntegrator: Velocity-Verlet 积分器NPT 系综恒温恒压 NPU: 所有原子的位置、速度、力同时更新 def__init__(self,dt1.0,n_steps1000000):self.dtdt# fs (1 fs 10^-15 s)self.n_stepsn_stepsdefstep(self,positions,velocities,forces,masses,force_engine,pair_i,pair_j,box_size): 一步 Velocity-Verlet positions: [N, 3] (Å) velocities: [N, 3] (Å/fs) forces: [N, 3] (kcal/(mol·Å) 原子单位力) masses: [N] (amu 原子质量单位) dtself.dt devicepositions.device# 质量 → 加速度的转换因子# F m * a → a F / m# 校准: 1 kcal/(mol·Å) / 1 amu 4184 m/s² 4.184e-4 Å/fs²accel_factor4.184e-4# Å/fs² per kcal/(mol·Å·amu)# Step 1: 半步速度# v_half v a * dt/2accelerationsforces/masses.unsqueeze(1)*accel_factor# [N, 3]v_halfvelocitiesaccelerations*(dt/2.0)# Step 2: 更新位置positions_newpositionsv_half*dt# 周期边界: 原子不能漫游出模拟盒子positions_newpositions_new%box_size.unsqueeze(0)# Step 3: 计算新位置的力这是计算最密集的部分# 注: 实际模拟中会检测是否需要重建邻接列表forces_newforce_engine.compute_forces(positions_new,pair_i,pair_j,box_size)# Step 4: 更新速度为全步accelerations_newforces_new/masses.unsqueeze(1)*accel_factor velocities_newv_halfaccelerations_new*(dt/2.0)returnpositions_new,velocities_new,forces_newdefrun(self,positions,velocities,masses,force_engine,neighbor_list,box_size,rebuild_interval20): 批量运行模拟 rebuild_interval: 每 20 步重建邻接列表 重建间隔 skin / (max_velocity * dt) ≈ 2Å / (0.1Å/fs * 1fs) 20 步 # 初始邻接列表pair_i,pair_jneighbor_list.build(positions,box_size)traj[]# 轨迹存储energies[]forstepinrange(self.n_steps):# 定期重建邻接列表ifstep%rebuild_interval0andstep0:pair_i,pair_jneighbor_list.build(positions,box_size)positions,velocities,forcesself.step(positions,velocities,forces,masses,force_engine,pair_i,pair_j,box_size)# 每隔一定步数记录ifstep%10000:E_kin0.5*(masses*velocities.pow(2).sum(dim1)).sum()E_potforce_engine.compute_energy(positions,pair_i,pair_j,box_size)energies.append((E_kin.item(),E_pot.item()))# 温度: T 2*E_kin / (3*N*k_B)Nlen(positions)k_B0.001987# kcal/(mol·K)T2.0*E_kin/(3.0*N*k_B)traj.append(positions.clone())ifstep%100000:print(fStep{step:8d}: T{T:.1f}K fE_kin{E_kin:.2f}E_pot{E_pot:.2f}fE_tot{E_kinE_pot:.2f}kcal/mol)returntraj,energies踩坑邻接列表重建太频繁——截断半径紧挨着 LJ 势的截断刚超 12Å 就丢失相互作用# ❌ skin0邻接列表半径 cut_off 12Å# 原子移动超出 12Å → 邻接列表中丢失 → LJ 力计算遗漏 → 能量不守恒# 但每步重建邻接列表 → 占模拟时间 40%# ✅ skin2Å邻接列表半径 14Å# 缓冲区内原子即使 12-14Å 也不参与 LJforce 截断在 12Å# 但邻接列表中保留了它们 → 20 步内不会丢失# 重建间隔 20 → 邻接列表开销 5%踩坑NPT 系综的恒压器Berendsen barostat在大时间步下盒尺寸振荡# ❌ Berendsen barostat 直接缩放置信盒尺寸# box_new box * (1 - κ * dt/τ_p * (P_target - P_instant))# dt2fs两倍标准步长→ κ*dt/τ_p 过大 → 盒尺寸振荡 ±15%# ✅ Parrinello-Rahman barostat: 用扩展拉格朗日量把盒尺寸作为动力变量classParrinelloRahmanBarostat: P-R 恒压器: 盒尺寸是动力变量有惯性不会振荡 def__init__(self,P_target1.0,tau_p1.0,compressibility4.5e-5): P_target: 目标压力 (atm) tau_p: 恒压器时间常数 (ps) compressibility: 水的压缩系数 (atm^-1) self.P_targetP_target self.tau_ptau_p self.betacompressibility# 盒矩阵 h [a_x, b_x, c_x; a_y, b_y, c_y; a_z, b_z, c_z]self.hNoneself.h_velNone# dh/dt, 盒矩阵的速率defrescale(self,positions,forces,stress_tensor,dt): P-R 恒压器一步 stress_tensor: [3, 3] 瞬时应力 (atm) ifself.hisNone:self.h10.0*torch.eye(3,devicepositions.device)hself.h Vtorch.det(h)# 盒体积# 内能贡献: P_inst (2*E_kin/3 sum(r·F))/3V# stress_tensor 是从力场计算的 Virial 应力# 恒压器力: (P_inst - P_target) → 驱动 h 变化P_diff(stress_tensor.trace()/3.0-self.P_target)# 标量W3.0*300*100*self.tau_p**2# 盒矩阵的质量# 盒矩阵的加速度h_acc(V*P_diff/W)*self.beta# Verlet 更新盒矩阵h_halfself.h_velh_acc*dt/2.0self.hhh_half*dt# 原子坐标缩放: 单位坐标 s_i h^(-1) * r_i 不变storch.linalg.solve(h,positions.T).T# [N, 3] 单位坐标positions_news self.h.T# 新绝对坐标returnpositions_new踩坑index_add_ 在 FP16 下精度累积错误——10 万次迭代后原子飘移 0.3Å# ❌ forces forces.half() → index_add_ → FP16 累加# 每步 2000 万对原子 × 1fs → 100 万步 → FP16 精度不足# 原子位置误差 0.3Å → 相当于化学键长误差 20%# ✅ FP32 force buffer: 力的归约用 FP32位置和速度用 FP16classMixedPrecisionMD:混合精度 MD: FP16 坐标 FP32 力积累def__init__(self,N):self.NN# FP16: 坐标和速度节省 50% HBM2× 吞吐self.positionstorch.zeros(N,3,dtypetorch.float16,devicenpu)self.velocitiestorch.zeros(N,3,dtypetorch.float16,devicenpu)# FP32: 力积累缓冲区保证累加精度self.force_buffertorch.zeros(N,3,dtypetorch.float32,devicenpu)defaccumulate_force(self,F_ij,pair_i,pair_j):FP32 力归约每步清零 index_addself.force_buffer.zero_()self.force_buffer.index_add_(0,pair_i,F_ij.float())self.force_buffer.index_add_(0,pair_j,-F_ij.float())# 转 FP16 给速度更新returnself.force_buffer.half()mat-chem-sim-pred 的分子动力学方案空间哈希桶 27 邻域搜索把 50 亿对原子缩减到 2000 万对25×Lennard-Jones 12-6 势对所有候选对并行计算力矢量Cube 单元做 r 的乘幂 direction 归一化Velocity-Verlet 积分交替更新位置和速度半步 v → 全步 x → 新力 → 全步 v。10 万原子 × 1μs 模拟从集群一月压到 NPU 数小时。踩坑skin0 导致频繁重建邻接列表40% 开销→skin2Å 缓冲5%、Berendsen barostat 振荡→Parrinello-Rahman 盒矩阵动力变量、FP16 index_add_ 原子飘移 0.3Å→FP32 力缓冲区归约。