""" engines/python/dynamics_lib.py ------------------------------- 纯 NumPy 计算引擎：无文件 I/O，所有数据以 NumPy 数组传入， 结果作为 NumPy 数组返回。 接口与 C/C++/Fortran DLL 的 run_dynamics() 完全一致， 算法与 compute.py 的 run_simulation() 保持一致。 用法（由 engine_dll.py 内部调用）： from engines.python.dynamics_lib import run_dynamics out_x, out_y, out_z, out_vx, out_vy, out_vz = run_dynamics(...) """ import numpy as np TWO_PI = 2.0 * np.pi # ── method_id 映射 ────────────────────────────────────────── # 0=euler 1=implicit_euler 2=midpoint 3=leapfrog # ── 保守加速度（弹簧键 + 均匀重力场，不含阻尼）────────────── def _accel_conservative(x, y, z, m, Gx, Gy, Gz, gravity_field, elastic_force, bond_pairs, bond_k, bond_r0): ax = np.full_like(x, Gx) if gravity_field else np.zeros_like(x) ay = np.full_like(y, Gy) if gravity_field else np.zeros_like(y) az = np.full_like(z, Gz) if gravity_field else np.zeros_like(z) if elastic_force and len(bond_pairs) > 0: i1 = bond_pairs[:, 0] i2 = bond_pairs[:, 1] dx = x[i2] - x[i1] dy = y[i2] - y[i1] dz = z[i2] - z[i1] dist = np.sqrt(dx*dx + dy*dy + dz*dz) valid = dist > 1e-12 fac = np.where(valid, bond_k * (dist - bond_r0) / dist, 0.0) fx = fac * dx fy = fac * dy fz_b = fac * dz np.add.at(ax, i1, fx / m[i1]); np.add.at(ax, i2, -fx / m[i2]) np.add.at(ay, i1, fy / m[i1]); np.add.at(ay, i2, -fy / m[i2]) np.add.at(az, i1, fz_b / m[i1]); np.add.at(az, i2, -fz_b / m[i2]) return ax, ay, az # ── 完整加速度（含阻尼）────────────────────────────────────── def _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0): ax, ay, az = _accel_conservative(x, y, z, m, Gx, Gy, Gz, gravity_field, elastic_force, bond_pairs, bond_k, bond_r0) if damping_force: ax -= Bx * vx / m ay -= By * vy / m az -= Bz * vz / m return ax, ay, az # ── 蛙跳法（半隐式阻尼，与 compute.py leapfrog_staggered_step 一致）─ def _leapfrog_step(x, y, z, vx, vy, vz, fixed, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt): ax, ay, az = _accel_conservative(x, y, z, m, Gx, Gy, Gz, gravity_field, elastic_force, bond_pairs, bond_k, bond_r0) has_damp = damping_force and (Bx != 0.0 or By != 0.0 or Bz != 0.0) if has_damp: alpha_x = Bx * dt / (2.0 * m) alpha_y = By * dt / (2.0 * m) alpha_z = Bz * dt / (2.0 * m) vx_new = (vx * (1.0 - alpha_x) + ax * dt) / (1.0 + alpha_x) vy_new = (vy * (1.0 - alpha_y) + ay * dt) / (1.0 + alpha_y) vz_new = (vz * (1.0 - alpha_z) + az * dt) / (1.0 + alpha_z) else: vx_new = vx + ax * dt vy_new = vy + ay * dt vz_new = vz + az * dt # 全固定原子保持不变 all_fixed = np.all(fixed, axis=1) vx_new = np.where(all_fixed, vx, vx_new) vy_new = np.where(all_fixed, vy, vy_new) vz_new = np.where(all_fixed, vz, vz_new) x_new = x + vx_new * dt y_new = y + vy_new * dt z_new = z + vz_new * dt return x_new, y_new, z_new, vx_new, vy_new, vz_new # ── 显式欧拉法 ─────────────────────────────────────────────── def _euler_step(x, y, z, vx, vy, vz, fixed, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt): ax, ay, az = _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0) all_fixed = np.all(fixed, axis=1) mask = ~all_fixed x_new = np.where(mask, x + vx * dt, x) y_new = np.where(mask, y + vy * dt, y) z_new = np.where(mask, z + vz * dt, z) vx_new = np.where(mask, vx + ax * dt, vx) vy_new = np.where(mask, vy + ay * dt, vy) vz_new = np.where(mask, vz + az * dt, vz) return x_new, y_new, z_new, vx_new, vy_new, vz_new # ── 隐式欧拉法（与 compute.py Implicit_Euler_Method 一致）────── def _implicit_euler_step(x, y, z, vx, vy, vz, fixed, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt): gamma_x = Bx / m gamma_y = By / m gamma_z = Bz / m vx_next = (vx + Gx * dt) / (1.0 + gamma_x * dt) vy_next = (vy + Gy * dt) / (1.0 + gamma_y * dt) vz_next = (vz + Gz * dt) / (1.0 + gamma_z * dt) ax, ay, az = _accel_full(x, y, z, vx_next, vy_next, vz_next, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0) all_fixed = np.all(fixed, axis=1) mask = ~all_fixed vx_new = np.where(mask, vx + ax * dt, vx) vy_new = np.where(mask, vy + ay * dt, vy) vz_new = np.where(mask, vz + az * dt, vz) x_new = np.where(mask, x + vx_new * dt, x) y_new = np.where(mask, y + vy_new * dt, y) z_new = np.where(mask, z + vz_new * dt, z) return x_new, y_new, z_new, vx_new, vy_new, vz_new # ── 中点法（与 compute.py Midpoint_Method 一致）──────────────── def _midpoint_step(x, y, z, vx, vy, vz, fixed, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt): ax, ay, az = _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0) all_fixed = np.all(fixed, axis=1) mask = ~all_fixed xm = np.where(mask, x + 0.5*vx*dt, x) ym = np.where(mask, y + 0.5*vy*dt, y) zm = np.where(mask, z + 0.5*vz*dt, z) vxm = np.where(mask, vx + 0.5*ax*dt, 0.0) vym = np.where(mask, vy + 0.5*ay*dt, 0.0) vzm = np.where(mask, vz + 0.5*az*dt, 0.0) x_new = np.where(mask, x + vxm * dt, x) y_new = np.where(mask, y + vym * dt, y) z_new = np.where(mask, z + vzm * dt, z) axm, aym, azm = _accel_full(xm, ym, zm, vxm, vym, vzm, m, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0) vx_new = np.where(mask, vx + axm * dt, vx) vy_new = np.where(mask, vy + aym * dt, vy) vz_new = np.where(mask, vz + azm * dt, vz) return x_new, y_new, z_new, vx_new, vy_new, vz_new # ── 边界：反弹 + 回绕 + 逐自由度固定约束 ─────────────────────── def _apply_bc(x, y, z, vx, vy, vz, fixed, pos_init, box_a): lo, hi = -box_a, box_a # 反弹（全固定原子跳过） all_fixed = np.all(fixed, axis=1) do_bc = ~all_fixed over_x = do_bc & (x > hi); under_x = do_bc & (x < lo) over_y = do_bc & (y > hi); under_y = do_bc & (y < lo) over_z = do_bc & (z > hi); under_z = do_bc & (z < lo) x = np.where(over_x, hi, np.where(under_x, lo, x)) y = np.where(over_y, hi, np.where(under_y, lo, y)) z = np.where(over_z, hi, np.where(under_z, lo, z)) vx = np.where(over_x | under_x, -np.abs(vx)*np.sign(np.where(over_x, 1, -1)), vx) vy = np.where(over_y | under_y, -np.abs(vy)*np.sign(np.where(over_y, 1, -1)), vy) vz = np.where(over_z | under_z, -np.abs(vz)*np.sign(np.where(over_z, 1, -1)), vz) # 反弹速度简化：越界则取反绝对值（与 C 版 _limit1 一致） vx = np.where(over_x, -np.abs(vx), np.where(under_x, np.abs(vx), vx)) vy = np.where(over_y, -np.abs(vy), np.where(under_y, np.abs(vy), vy)) vz = np.where(over_z, -np.abs(vz), np.where(under_z, np.abs(vz), vz)) # 回绕 x = np.where(x > hi, lo, np.where(x < lo, hi, x)) y = np.where(y > hi, lo, np.where(y < lo, hi, y)) z = np.where(z > hi, lo, np.where(z < lo, hi, z)) # 逐自由度固定约束 fx = fixed[:, 0].astype(bool) fy = fixed[:, 1].astype(bool) fz = fixed[:, 2].astype(bool) x = np.where(fx, pos_init[:, 0], x); vx = np.where(fx, 0.0, vx) y = np.where(fy, pos_init[:, 1], y); vy = np.where(fy, 0.0, vy) z = np.where(fz, pos_init[:, 2], z); vz = np.where(fz, 0.0, vz) return x, y, z, vx, vy, vz # ── 驱动力（与 compute.py apply_driving_force 逻辑一致）───────── def _apply_driving(x, y, z, vx, vy, vz, t, step, dt, drv_idx, drv_amp, drv_freq, drv_phi, drv_eq, drv_ncycles, drv_has_period, freeze): """freeze: (n_drivers, 3) mutable array for frozen positions.""" nd = len(drv_idx) for d in range(nd): idx = drv_idx[d] fx_ = drv_freq[d, 0]; fy_ = drv_freq[d, 1]; fz_ = drv_freq[d, 2] if drv_has_period[d]: mf = max(abs(fx_), abs(fy_), abs(fz_)) ps = int(drv_ncycles[d] / mf / dt) if mf > 1e-12 else 0 if step > ps: x[idx] = freeze[d, 0]; y[idx] = freeze[d, 1]; z[idx] = freeze[d, 2] vx[idx] = vy[idx] = vz[idx] = 0.0 continue px = drv_eq[d,0] + drv_amp[d,0]*np.cos(TWO_PI*fx_*t + drv_phi[d,0]) py = drv_eq[d,1] + drv_amp[d,1]*np.cos(TWO_PI*fy_*t + drv_phi[d,1]) pz = drv_eq[d,2] + drv_amp[d,2]*np.cos(TWO_PI*fz_*t + drv_phi[d,2]) if step == ps: freeze[d, 0] = px; freeze[d, 1] = py; freeze[d, 2] = pz x[idx] = drv_eq[d,0] + drv_amp[d,0]*np.cos(TWO_PI*fx_*t + drv_phi[d,0]) y[idx] = drv_eq[d,1] + drv_amp[d,1]*np.cos(TWO_PI*fy_*t + drv_phi[d,1]) z[idx] = drv_eq[d,2] + drv_amp[d,2]*np.cos(TWO_PI*fz_*t + drv_phi[d,2]) vx[idx] = -drv_amp[d,0]*TWO_PI*fx_*np.sin(TWO_PI*fx_*t + drv_phi[d,0]) vy[idx] = -drv_amp[d,1]*TWO_PI*fy_*np.sin(TWO_PI*fy_*t + drv_phi[d,1]) vz[idx] = -drv_amp[d,2]*TWO_PI*fz_*np.sin(TWO_PI*fz_*t + drv_phi[d,2]) def _do_step(x, y, z, vx, vy, vz, fixed, masses, method_id, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt, pos_init, box_a): if method_id == 0: x, y, z, vx, vy, vz = _euler_step( x, y, z, vx, vy, vz, fixed, masses, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt) elif method_id == 1: x, y, z, vx, vy, vz = _implicit_euler_step( x, y, z, vx, vy, vz, fixed, masses, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt) elif method_id == 2: x, y, z, vx, vy, vz = _midpoint_step( x, y, z, vx, vy, vz, fixed, masses, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt) else: x, y, z, vx, vy, vz = _leapfrog_step( x, y, z, vx, vy, vz, fixed, masses, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt) x, y, z, vx, vy, vz = _apply_bc(x, y, z, vx, vy, vz, fixed, pos_init, box_a) return x, y, z, vx, vy, vz # ══════════════════════════════════════════════════════════════ # 主函数：run_dynamics # 接口与 C/C++/Fortran DLL 的 run_dynamics() 对应， # 参数格式：numpy 数组（替代 ctypes 指针）。 # # method_id: 0=euler 1=implicit_euler 2=midpoint 3=leapfrog # drv_amp/freq/phi/eq: (n_drivers, 3) float64 # drv_ncycles: (n_drivers,) float64 0=不限 # drv_has_period: (n_drivers,) int # # 返回：(out_x, out_y, out_z, out_vx, out_vy, out_vz) # 各 shape=(n_frames, n_atoms) # ══════════════════════════════════════════════════════════════ def run_dynamics( n_atoms, pos_init, vel_init, masses, fixed, n_bonds, bond_pairs, bond_k, bond_r0, box_a, dt, NT, NSTEP, warmup_steps, method_id, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, gravity_strength, n_drivers, drv_idx, drv_amp, drv_freq, drv_phi, drv_eq, drv_ncycles, drv_has_period, n_frames, progress_cb=None, ): """运行动力学模拟，返回抽帧轨迹数组。 Args: pos_init: (n_atoms, 3) float64 vel_init: (n_atoms, 3) float64 masses: (n_atoms,) float64 fixed: (n_atoms, 3) int — 1=固定 bond_pairs: (n_bonds, 2) int — 0-based 局部索引 bond_k: (n_bonds,) float64 bond_r0: (n_bonds,) float64 drv_idx: (n_drivers,) int — 0-based drv_amp/freq/phi/eq: (n_drivers, 3) float64 drv_ncycles: (n_drivers,) float64 drv_has_period: (n_drivers,) int n_frames: 预分配的输出帧数 Returns: out_x, out_y, out_z, out_vx, out_vy, out_vz — 各 (n_frames, n_atoms) """ pos_init = np.asarray(pos_init, dtype=np.float64) vel_init = np.asarray(vel_init, dtype=np.float64) masses = np.asarray(masses, dtype=np.float64) fixed = np.asarray(fixed, dtype=np.int32) bond_pairs = np.asarray(bond_pairs, dtype=np.int64).reshape(-1, 2) if n_bonds else np.zeros((0,2), dtype=np.int64) bond_k = np.asarray(bond_k, dtype=np.float64) if n_bonds else np.zeros(0) bond_r0 = np.asarray(bond_r0, dtype=np.float64) if n_bonds else np.zeros(0) n = n_atoms x = pos_init[:, 0].copy() y = pos_init[:, 1].copy() z = pos_init[:, 2].copy() vx = vel_init[:, 0].copy() vy = vel_init[:, 1].copy() vz = vel_init[:, 2].copy() # 驱动力数据（保证正确形状） nd = n_drivers if nd > 0: drv_idx = np.asarray(drv_idx, dtype=np.int64) drv_amp = np.asarray(drv_amp, dtype=np.float64).reshape(nd, 3) drv_freq = np.asarray(drv_freq, dtype=np.float64).reshape(nd, 3) drv_phi = np.asarray(drv_phi, dtype=np.float64).reshape(nd, 3) drv_eq = np.asarray(drv_eq, dtype=np.float64).reshape(nd, 3) drv_nc = np.asarray(drv_ncycles, dtype=np.float64) drv_hp = np.asarray(drv_has_period, dtype=np.int32) freeze = np.zeros((nd, 3), dtype=np.float64) else: drv_idx = drv_amp = drv_freq = drv_phi = drv_eq = drv_nc = drv_hp = freeze = None def _drive(t_, step_): if nd > 0: _apply_driving(x, y, z, vx, vy, vz, t_, step_, dt, drv_idx, drv_amp, drv_freq, drv_phi, drv_eq, drv_nc, drv_hp, freeze) # ── 蛙跳法：初始化 v(-dt/2) ───────────────────────────── if method_id == 3: ax0, ay0, az0 = _accel_conservative(x, y, z, masses, Gx, Gy, Gz, gravity_field, elastic_force, bond_pairs, bond_k, bond_r0) all_fixed = np.all(fixed, axis=1) vx = np.where(all_fixed, vx, vx - 0.5 * ax0 * dt) vy = np.where(all_fixed, vy, vy - 0.5 * ay0 * dt) vz = np.where(all_fixed, vz, vz - 0.5 * az0 * dt) # ── 初始驱动 t=0 ───────────────────────────────────────── _drive(0.0, 0) # ── 预热 ───────────────────────────────────────────────── for s in range(warmup_steps): tw = (s + 1) * dt _drive(tw, s) x, y, z, vx, vy, vz = _do_step( x, y, z, vx, vy, vz, fixed, masses, method_id, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt, pos_init, box_a) # ── 记录循环 ───────────────────────────────────────────── record_steps = NT - warmup_steps prog_interval = max(1, record_steps // 100) out_x = np.zeros((n_frames, n), dtype=np.float64) out_y = np.zeros((n_frames, n), dtype=np.float64) out_z = np.zeros((n_frames, n), dtype=np.float64) out_vx = np.zeros((n_frames, n), dtype=np.float64) out_vy = np.zeros((n_frames, n), dtype=np.float64) out_vz = np.zeros((n_frames, n), dtype=np.float64) frame_idx = 0 for s in range(record_steps): if progress_cb is not None and s % prog_interval == 0 and s > 0: progress_cb(s, record_steps) t = (s + warmup_steps) * dt _drive(t, s) if s % NSTEP == 0 and frame_idx < n_frames: out_x[frame_idx] = x out_y[frame_idx] = y out_z[frame_idx] = z out_vx[frame_idx] = vx out_vy[frame_idx] = vy out_vz[frame_idx] = vz frame_idx += 1 x, y, z, vx, vy, vz = _do_step( x, y, z, vx, vy, vz, fixed, masses, method_id, Gx, Gy, Gz, Bx, By, Bz, gravity_field, elastic_force, damping_force, bond_pairs, bond_k, bond_r0, dt, pos_init, box_a) return out_x, out_y, out_z, out_vx, out_vy, out_vz