""" engines/python/main.py ----------------------- 独立 Python 计算引擎。 与 main.c / main.cpp / main.f90 结构一致： 输入: /coord.txt, connection.txt, bond.txt, [driver.txt] (同 engines/c/param.json 格式) 输出: /display.txt (+ display.npz) /trajectory.txt (若 save_trajectory=1) 用法: python main.py 内部调用 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 ") 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), "display_color": json.dumps(p.get("display_color", {"x":[0,[1.0,0.0,0.0]],"y":[0,[0.0,1.0,0.0]],"z":[0,[0.0,0.0,1.0]], "xy":[0,[1.0,1.0,0.0]],"yz":[0,[0.0,1.0,1.0]],"zx":[0,[1.0,0.0,1.0]], "xyz":[1,[1.0,1.0,1.0]]})), "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()