admin
0e636e275d
- 覆盖全部 10 个案例（原 Readme 只到 case06） - 新增案例选择指南表格 - Readme.html 为深色主题独立 HTML 页面 （含卡片布局、标签分类、代码高亮、响应式设计） - 各案例详情对齐最新配置参数
284 lines
12 KiB
Python
284 lines
12 KiB
Python
"""
|
||
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()
|