ylniu/dynamics
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: 更新 examples/Readme.md 并新增 Readme.html

Browse Source 
- 覆盖全部 10 个案例（原 Readme 只到 case06）
- 新增案例选择指南表格
- Readme.html 为深色主题独立 HTML 页面
  （含卡片布局、标签分类、代码高亮、响应式设计）
- 各案例详情对齐最新配置参数
This commit is contained in:
Ying-Li Niu
2026-06-17 15:33:49 +08:00
parent ea99f09f9b
commit 0e636e275d
59 changed files with 7302 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engines/python/__init__.py
View File
engines/python/dynamics_lib.py
+404
View File
@@ -0,0 +1,404 @@
"""
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
engines/python/main.py
+283
View File
@@ -0,0 +1,283 @@
"""
engines/python/main.py
-----------------------
独立 Python 计算引擎。
与 main.c / main.cpp / main.f90 结构一致：
输入: <input_dir>/coord.txt, connection.txt, bond.txt, [driver.txt]
<param_json> (同 engines/c/param.json 格式)
输出: <output_dir>/display.txt (+ display.npz)
<output_dir>/trajectory.txt (若 save_trajectory=1)
用法:
python main.py <input_dir> <output_dir> <param_json>
内部调用 dynamics_lib.run_dynamics()，算法与 compute.py 完全一致。
"""
import json
import os
import sys
import time
import numpy as np
# 将父目录（engines/python 的上级 engines）加入 sys.path，
# 以便在独立运行时也能找到 dynamics_lib
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
from dynamics_lib import run_dynamics
# 为读取 coord/bond/display，复用 compute.py 中的 I/O 函数
_COMPUTE = os.path.join(_HERE, "..", "..")
sys.path.insert(0, _COMPUTE)
import compute as _c
_METHOD_ID = {
"explicit_euler": 0,
"euler": 0,
"implicit_euler": 1,
"midpoint": 2,
"leapfrog": 3,
}
def _load_params(param_path):
"""读取 param.json（与 C 引擎格式相同）."""
with open(param_path, "r", encoding="utf-8") as f:
p = json.load(f)
return p
def main():
if len(sys.argv) < 4:
print("用法: python main.py <input_dir> <output_dir> <param_json>")
sys.exit(1)
input_dir = sys.argv[1]
output_dir = sys.argv[2]
param_path = sys.argv[3]
os.makedirs(output_dir, exist_ok=True)
# ── 读取参数 ─────────────────────────────────────────────
p = _load_params(param_path)
box_a = float(p.get("box_a", 10.0))
NT = int(p.get("NT", 10000))
dt = float(p.get("DT", 0.001))
NSTEP = int(p.get("NSTEP", 100))
warmup_steps = int(p.get("warmup_steps", 0))
method_str = str(p.get("method", "leapfrog")).lower().replace(" ", "_")
method_id = _METHOD_ID.get(method_str, 3)
G = p.get("G", [0.0, 0.0, -9.8])
B = p.get("B", [0.0, 0.0, 0.0])
gravity_field = int(p.get("gravity_field", 1))
elastic_force = int(p.get("elastic_force", 1))
damping_force = int(p.get("damping_force", 0))
gravity_strength = float(p.get("gravity_strength", 1.0))
driving_force = int(p.get("driving_force", 0))
save_traj = int(p.get("save_trajectory", 0))
# ── 读取原子数据 ──────────────────────────────────────────
coord_path = os.path.join(input_dir, "coord.txt")
atom_ids, masses, radii, positions, velocities, fixed = _c.load_coord_file(coord_path)
# ── 读取键数据 ────────────────────────────────────────────
conn_path = os.path.join(input_dir, "connection.txt")
bond_path = os.path.join(input_dir, "bond.txt")
bond_map = _c.load_bond_parameters(bond_path)
bond_pairs, bond_names, bond_stiffness, bond_rest_lengths = \
_c.load_bond_connections(conn_path, atom_ids, positions, bond_map)
n_bonds = len(bond_pairs)
# ── 读取驱动力 ────────────────────────────────────────────
drv_list = []
if driving_force:
driver_path = os.path.join(input_dir, "driver.txt")
raw_drivers = _c.load_driver_file(driver_path, atom_ids)
if raw_drivers:
atom_id_map = {int(aid): i for i, aid in enumerate(atom_ids)}
for d in raw_drivers:
aid = int(d["atom_id"])
if aid not in atom_id_map:
continue
lidx = atom_id_map[aid]
eq = positions[lidx].tolist()
d["eq_pos"] = np.array(eq)
pc = d.get("period_cycles")
nc = float(pc) if pc is not None else 0.0
hp = 1 if nc > 0 else 0
drv_list.append({
"local_idx": lidx,
"amp": d["amp"].tolist(),
"freq": d["freq"].tolist(),
"phi": d["phi"].tolist(), # radians
"eq": eq,
"nc": nc,
"hp": hp,
})
nd = len(drv_list)
if nd > 0:
drv_idx = np.array([d["local_idx"] for d in drv_list], dtype=np.int64)
drv_amp = np.array([d["amp"] for d in drv_list], dtype=np.float64)
drv_freq = np.array([d["freq"] for d in drv_list], dtype=np.float64)
drv_phi = np.array([d["phi"] for d in drv_list], dtype=np.float64)
drv_eq = np.array([d["eq"] for d in drv_list], dtype=np.float64)
drv_nc = np.array([d["nc"] for d in drv_list], dtype=np.float64)
drv_hp = np.array([d["hp"] for d in drv_list], dtype=np.int32)
else:
drv_idx = drv_amp = drv_freq = drv_phi = drv_eq = drv_nc = drv_hp = \
np.zeros(0, dtype=np.int64)
# ── 计算帧数 ──────────────────────────────────────────────
record_steps = NT - warmup_steps
n_frames = max(1, record_steps // NSTEP)
# ── 进度回调 ──────────────────────────────────────────────
def _progress(step, total):
pct = step * 100 // total
print(f"[python-engine] progress: {step}/{total} ({pct}%)", flush=True)
# ── 运行计算 ──────────────────────────────────────────────
t0 = time.time()
print(f"[python-engine] NT={NT} NSTEP={NSTEP} method={method_str} "
f"n_atoms={len(atom_ids)} n_bonds={n_bonds}")
out_x, out_y, out_z, out_vx, out_vy, out_vz = run_dynamics(
n_atoms=len(atom_ids),
pos_init=positions,
vel_init=velocities,
masses=masses,
fixed=fixed,
n_bonds=n_bonds,
bond_pairs=bond_pairs,
bond_k=bond_stiffness,
bond_r0=bond_rest_lengths,
box_a=box_a,
dt=dt,
NT=NT,
NSTEP=NSTEP,
warmup_steps=warmup_steps,
method_id=method_id,
Gx=float(G[0]), Gy=float(G[1]), Gz=float(G[2]),
Bx=float(B[0]), By=float(B[1]), Bz=float(B[2]),
gravity_field=gravity_field,
elastic_force=elastic_force,
damping_force=damping_force,
gravity_strength=gravity_strength,
n_drivers=nd,
drv_idx=drv_idx,
drv_amp=drv_amp,
drv_freq=drv_freq,
drv_phi=drv_phi,
drv_eq=drv_eq,
drv_ncycles=drv_nc,
drv_has_period=drv_hp,
n_frames=n_frames,
progress_cb=_progress,
)
elapsed = time.time() - t0
print(f"[python-engine] 完成: {n_frames} 帧 {elapsed:.3f} s")
# ── 构建 display header ───────────────────────────────────
ball_radius = float(p.get("ball_radius", 0.5))
ball_color = p.get("ball_color", [0.9, 0.2, 0.2])
box_color = p.get("box_color", [0.8, 0.8, 0.85])
use_marker = int(p.get("use_marker", 0))
alpha_val = p.get("alpha", 0.2)
cam_dist = float(p.get("camera_distance", 40.0))
cam_elev = float(p.get("camera_elevation", 0.0))
cam_azim = float(p.get("camera_azimuth", 0.0))
cam_cx = float(p.get("camera_center_x", 0.0))
cam_cy = float(p.get("camera_center_y", 0.0))
cam_cz = float(p.get("camera_center_z", 0.0))
header = {
"DT": str(dt),
"NSTEP": str(NSTEP),
"method": method_str,
"NT": str(NT),
"warmup_steps": str(warmup_steps),
"dynamic_steps": str(record_steps),
"T_total": str(NT * dt),
"box_a": str(box_a),
"gravity_field": str(gravity_field),
"elastic_force": str(elastic_force),
"damping_force": str(damping_force),
"driving_force": str(driving_force),
"gravity_strength": str(gravity_strength),
"G": json.dumps([float(v) for v in G]),
"B": json.dumps([float(v) for v in B]),
"number_of_frames": str(n_frames),
"number_of_particles": str(len(atom_ids)),
"use_marker": str(use_marker),
"ball_radius": str(ball_radius),
"ball_color_r": str(ball_color[0]),
"ball_color_g": str(ball_color[1]),
"ball_color_b": str(ball_color[2]),
"box_color_r": str(box_color[0]),
"box_color_g": str(box_color[1]),
"box_color_b": str(box_color[2]),
"alpha": str(alpha_val) if not isinstance(alpha_val, list)
else ",".join(str(a) for a in alpha_val),
"atom_radii": ",".join(str(r) for r in radii),
"atom_masses": json.dumps([float(m) for m in masses]),
"atom_positions": json.dumps(positions.tolist()),
"bond_pairs": json.dumps(bond_pairs.tolist() if n_bonds else []),
"bond_stiffness": json.dumps(bond_stiffness.tolist() if n_bonds else []),
"bond_rest_lengths": json.dumps(bond_rest_lengths.tolist() if n_bonds else []),
"X_MIN": str(-box_a), "X_MAX": str(box_a),
"Y_MIN": str(-box_a), "Y_MAX": str(box_a),
"Z_MIN": str(-box_a), "Z_MAX": str(box_a),
"camera_distance": str(cam_dist),
"camera_elevation": str(cam_elev),
"camera_azimuth": str(cam_azim),
"camera_center_x": str(cam_cx),
"camera_center_y": str(cam_cy),
"camera_center_z": str(cam_cz),
"camera_keyframes": "",
}
# ── 保存 display.txt + display.npz ───────────────────────
disp_txt = os.path.join(output_dir, "display.txt")
_c.save_display_txt(
disp_txt,
out_x, out_y, out_z, out_vx, out_vy, out_vz,
atom_ids, record_steps, len(atom_ids),
header_fields=header,
)
print(f"[python-engine] display.txt 已保存: {disp_txt}")
disp_npz = os.path.join(output_dir, "display.npz")
_c.save_display_npz(
disp_npz,
out_x, out_y, out_z, out_vx, out_vy, out_vz,
atom_ids, header_fields=header,
)
print(f"[python-engine] display.npz 已保存: {disp_npz}")
# ── 可选：保存 trajectory.txt ─────────────────────────────
if save_traj:
traj_payload = {
"traj_x": out_x, "traj_y": out_y, "traj_z": out_z,
"traj_vx": out_vx, "traj_vy": out_vy, "traj_vz": out_vz,
"NT": record_steps, "DT": dt, "NSTEP": NSTEP,
"method": method_str,
"atom_ids": atom_ids,
"atom_masses": masses,
"atom_radii": radii,
"atom_positions": positions,
"bond_pairs": bond_pairs,
"bond_stiffness": bond_stiffness,
"bond_rest_lengths": bond_rest_lengths,
"G": [float(v) for v in G],
"B": [float(v) for v in B],
}
traj_path = os.path.join(output_dir, "trajectory.txt")
_c.save_text_data(traj_path, traj_payload)
print(f"[python-engine] trajectory.txt 已保存: {traj_path}")
if __name__ == "__main__":
main()