diff --git a/.gitignore b/.gitignore index ffe9ac7..8fe0180 100644 --- a/.gitignore +++ b/.gitignore @@ -100,3 +100,6 @@ desktop.ini *.tar.gz *.tar.bz2 output_test/ + +optimization +optimization/* \ No newline at end of file diff --git a/compute.py b/compute.py index bfc7a5a..862310d 100644 --- a/compute.py +++ b/compute.py @@ -66,6 +66,12 @@ box_color_g = None box_color_b = None use_marker = 0 camera_keyframes_raw = "" +camera_distance = 40.0 +camera_elevation = 0.0 +camera_azimuth = 0.0 +FIXED_MASK_X = None +FIXED_MASK_Y = None +FIXED_MASK_Z = None # 力开关 GRAVITY_FIELD = 1 # 均匀重力场 @@ -313,6 +319,72 @@ def load_display_txt(path): } +def save_display_npz(path, frames_x, frames_y, frames_z, + frames_vx, frames_vy, frames_vz, + atom_ids, header_fields=None): + """Write display.npz binary format alongside display.txt. + + Stores all frame arrays in compressed NumPy format. Metadata is + serialised as a JSON byte-string stored in the 'meta' array so the + file stays self-contained. + + Args: + path: output path ending in '.npz' + frames_x/y/z/vx/vy/vz: (n_frames, n_atoms) float64 arrays + atom_ids: (n_atoms,) int array + header_fields: dict of metadata key-value strings (same as for + save_display_txt) + """ + meta_bytes = np.bytes_(json.dumps(header_fields or {}, ensure_ascii=False)) + stem = path[:-4] if path.endswith(".npz") else path + np.savez_compressed( + stem, + frames_x=frames_x, + frames_y=frames_y, + frames_z=frames_z, + frames_vx=frames_vx, + frames_vy=frames_vy, + frames_vz=frames_vz, + atom_ids=atom_ids, + meta=np.array(meta_bytes), + ) + return path + + +def load_display_npz(path): + """Read display.npz and return the same dict as load_display_txt.""" + data = np.load(path, allow_pickle=False) + header_fields = json.loads(data["meta"].item().decode("utf-8")) + frames_x = data["frames_x"] + n_total_frames = int(header_fields.get("number_of_frames", + frames_x.shape[0])) + n_total_particles = frames_x.shape[1] + return { + "frames_x": frames_x, + "frames_y": data["frames_y"], + "frames_z": data["frames_z"], + "frames_vx": data["frames_vx"], + "frames_vy": data["frames_vy"], + "frames_vz": data["frames_vz"], + "atom_ids": data["atom_ids"], + "n_total_frames": n_total_frames, + "n_total_particles": n_total_particles, + "header_fields": header_fields, + } + + +def load_display(output_dir): + """Load display data, preferring binary .npz over text .txt.""" + npz_path = os.path.join(output_dir, "display.npz") + txt_path = os.path.join(output_dir, "display.txt") + if os.path.exists(npz_path): + return load_display_npz(npz_path) + if os.path.exists(txt_path): + return load_display_txt(txt_path) + raise FileNotFoundError( + f"找不到 display.npz 或 display.txt in {output_dir}") + + def get_output_dir(base_dir=None): """Return the output directory used for generated artifacts.""" override = os.environ.get("DYNAMICS_OUTPUT_DIR") @@ -627,9 +699,8 @@ def apply_driving_force(x, y, z, vx, vy, vz, t, step, drivers, dt): period_steps = None # 全程驱动 # 当前驱动力下的位置 / 速度 - t_vec = np.array([t, t, t], dtype=np.float64) - pos_drive = d["amp"] * np.cos(2.0 * np.pi * d["freq"] * t_vec + d["phi"]) - vel_drive = -d["amp"] * 2.0 * np.pi * d["freq"] * np.sin(2.0 * np.pi * d["freq"] * t_vec + d["phi"]) + pos_drive = d["amp"] * np.cos(2.0 * np.pi * d["freq"] * t + d["phi"]) + vel_drive = -d["amp"] * 2.0 * np.pi * d["freq"] * np.sin(2.0 * np.pi * d["freq"] * t + d["phi"]) x[idx] = pos_drive[0] y[idx] = pos_drive[1] @@ -692,7 +763,8 @@ def run_from_config(config, out_dir=None): global X_MIN, X_MAX, Y_MIN, Y_MAX, Z_MIN, Z_MAX global ball_radius, ball_color_r, ball_color_g, ball_color_b global box_color_r, box_color_g, box_color_b - global use_marker, camera_keyframes_raw + global use_marker, camera_keyframes_raw, camera_distance, camera_elevation, camera_azimuth + global FIXED_MASK_X, FIXED_MASK_Y, FIXED_MASK_Z global warmup_steps, sample_start, sample_end global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH global DRIVING_FORCE, DRIVER_DATA @@ -709,6 +781,9 @@ def run_from_config(config, out_dir=None): coord_path = os.path.join(out_dir, coord_path) (ATOM_IDS, ATOM_MASSES, ATOM_RADII, ATOM_POSITIONS, ATOM_VELOCITIES, ATOM_FIXED) = load_coord_file(coord_path) + FIXED_MASK_X = ATOM_FIXED[:, 0] != 0 + FIXED_MASK_Y = ATOM_FIXED[:, 1] != 0 + FIXED_MASK_Z = ATOM_FIXED[:, 2] != 0 BOND_CONNECTION_FILE = str(config.get("connection_file", os.path.join("input", "connection.txt"))) BOND_PARAMETER_FILE = str(config.get("bond_file", os.path.join("input", "bond.txt"))) connection_path = BOND_CONNECTION_FILE @@ -754,6 +829,9 @@ def run_from_config(config, out_dir=None): box_color_g = float(config.get("box_color_g", 0.8)) box_color_b = float(config.get("box_color_b", 0.85)) use_marker = int(config.get("use_marker", 0)) + camera_distance = float(config.get("camera_distance", 40.0)) + camera_elevation = float(config.get("camera_elevation", 0.0)) + camera_azimuth = float(config.get("camera_azimuth", 0.0)) # 力开关 global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH @@ -921,31 +999,55 @@ def run_engine(engine, input_dir, output_dir, config): print(f"[compute] input: {input_dir}") print(f"[compute] output: {output_dir}") - # ── 预校准:跑 NT=1000 步测速 ──────────────────────────────── + # ── 预校准:跑少量步测速(结果缓存,避免每次重跑)──────────── total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0)) _calib_nt = min(1000, max(100, total_steps // 10)) - _calib_param = dict(param_json) - _calib_param["NT"] = _calib_nt - _calib_path = os.path.join(script_dir, "engines", engine, "_calib.json") - with open(_calib_path, "w", encoding="utf-8") as _cf: - json.dump(_calib_param, _cf, indent=2) - _calib_outdir = os.path.join(script_dir, "engines", engine, "_calib_out") - os.makedirs(_calib_outdir, exist_ok=True) - _ct0 = time.time() - subprocess.run( - [engine_path, os.path.abspath(input_dir), _calib_outdir, _calib_path], - capture_output=True, timeout=60) - # 清理校准临时文件 - for _f in os.listdir(_calib_outdir): - try: os.remove(os.path.join(_calib_outdir, _f)) + n_atoms_calib = len(ATOM_IDS) if ATOM_IDS is not None else 0 + + # 尝试读取缓存;当 n_atoms 相同且 NT 在 50% 范围内时视为有效 + _cache_path = os.path.join(script_dir, "engines", engine, "_calib_cache.json") + _step_time = None + try: + with open(_cache_path, encoding="utf-8") as _cf: + _cache = json.load(_cf) + if (_cache.get("n_atoms") == n_atoms_calib and + abs(_cache.get("nt", 0) - total_steps) / max(total_steps, 1) < 0.5): + _step_time = float(_cache["step_time"]) + print(f"[compute] 使用校准缓存: {_step_time*1e6:.2f} μs/步") + except Exception: + pass + + if _step_time is None: + _calib_param = dict(param_json) + _calib_param["NT"] = _calib_nt + _calib_path = os.path.join(script_dir, "engines", engine, "_calib.json") + with open(_calib_path, "w", encoding="utf-8") as _cf: + json.dump(_calib_param, _cf, indent=2) + _calib_outdir = os.path.join(script_dir, "engines", engine, "_calib_out") + os.makedirs(_calib_outdir, exist_ok=True) + _ct0 = time.time() + subprocess.run( + [engine_path, os.path.abspath(input_dir), _calib_outdir, _calib_path], + capture_output=True, timeout=60) + # 清理校准临时文件 + for _f in os.listdir(_calib_outdir): + try: os.remove(os.path.join(_calib_outdir, _f)) + except OSError: pass + try: os.rmdir(_calib_outdir) except OSError: pass - try: os.rmdir(_calib_outdir) - except OSError: pass - os.remove(_calib_path) - _calib_elapsed = max(time.time() - _ct0, 0.001) - _overhead = _calib_elapsed * 0.15 - _step_time = max(_calib_elapsed - _overhead, 0.0001) / _calib_nt - _est_total = max(_calib_elapsed, _overhead + _step_time * total_steps) + os.remove(_calib_path) + _calib_elapsed = max(time.time() - _ct0, 0.001) + _overhead = _calib_elapsed * 0.15 + _step_time = max(_calib_elapsed - _overhead, 0.0001) / _calib_nt + # 保存缓存 + try: + with open(_cache_path, "w", encoding="utf-8") as _cf: + json.dump({"n_atoms": n_atoms_calib, "nt": total_steps, + "step_time": _step_time}, _cf) + except OSError: + pass + + _est_total = _step_time * total_steps t_start = time.time() t_start_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -1030,6 +1132,21 @@ def run_engine(engine, input_dir, output_dir, config): if not os.path.exists(disp_path): raise FileNotFoundError(f"引擎未生成 display.txt: {disp_path}") print(f"[compute] 引擎已生成 {disp_path}") + # 将 display.txt 转换为 display.npz(加快后续 draw.py / plot_wave.py 加载) + _npz_path = disp_path.replace("display.txt", "display.npz") + try: + _d = load_display_txt(disp_path) + save_display_npz( + _npz_path, + _d["frames_x"], _d["frames_y"], _d["frames_z"], + _d["frames_vx"], _d["frames_vy"], _d["frames_vz"], + _d["atom_ids"], + header_fields={**_d["header_fields"], + "number_of_frames": str(_d["n_total_frames"]), + "number_of_particles": str(_d["n_total_particles"])}) + print(f"[compute] display.npz 已生成: {_npz_path}") + except Exception as _e: + print(f"[compute] 警告: display.npz 生成失败({_e}),将使用 display.txt") return None, None, None, None, None, None # save_trajectory=1:加载完整 trajectory.txt @@ -1251,30 +1368,28 @@ def compute_force(x, y, z, vx, vy, vz, m, g, b): fz -= b[2] * vz if ELASTIC_FORCE and BOND_PAIRS is not None and len(BOND_PAIRS) > 0: - for bond_idx, (idx_1, idx_2) in enumerate(BOND_PAIRS): - dx = x[idx_2] - x[idx_1] - dy = y[idx_2] - y[idx_1] - dz = z[idx_2] - z[idx_1] - dist = np.sqrt(dx * dx + dy * dy + dz * dz) - if dist <= 1e-12: - continue + idx_1 = BOND_PAIRS[:, 0] + idx_2 = BOND_PAIRS[:, 1] + dx = x[idx_2] - x[idx_1] + dy = y[idx_2] - y[idx_1] + dz = z[idx_2] - z[idx_1] + dist = np.sqrt(dx * dx + dy * dy + dz * dz) + valid = dist > 1e-12 + if np.any(valid): + force_scale = np.zeros_like(dist) + stretch = dist[valid] - BOND_REST_LENGTHS[valid] + force_scale[valid] = BOND_STIFFNESS[valid] * stretch / dist[valid] - stretch = dist - BOND_REST_LENGTHS[bond_idx] - force_mag = BOND_STIFFNESS[bond_idx] * stretch - ux = dx / dist - uy = dy / dist - uz = dz / dist + fx_bond = force_scale * dx + fy_bond = force_scale * dy + fz_bond = force_scale * dz - fx_bond = force_mag * ux - fy_bond = force_mag * uy - fz_bond = force_mag * uz - - fx[idx_1] += fx_bond - fy[idx_1] += fy_bond - fz[idx_1] += fz_bond - fx[idx_2] -= fx_bond - fy[idx_2] -= fy_bond - fz[idx_2] -= fz_bond + np.add.at(fx, idx_1, fx_bond) + np.add.at(fx, idx_2, -fx_bond) + np.add.at(fy, idx_1, fy_bond) + np.add.at(fy, idx_2, -fy_bond) + np.add.at(fz, idx_1, fz_bond) + np.add.at(fz, idx_2, -fz_bond) if GRAVITY_INTERACTION: n = len(m) @@ -1304,6 +1419,50 @@ def compute_acceleration(x, y, z, vx, vy, vz, m, g, b): return fx / m, fy / m, fz / m +def compute_accel_conservative(x, y, z, m, g): + """保守力加速度(弹簧 + 重力 + 原子间引力),不含阻尼。 + 供真蛙跳法使用:阻尼项由 leapfrog_staggered_step 半隐式处理。 + """ + _v0 = np.zeros_like(x) + _b0 = np.zeros(3) + fx, fy, fz = compute_force(x, y, z, _v0, _v0, _v0, m, g, _b0) + return fx / m, fy / m, fz / m + + +def leapfrog_staggered_step(x, y, z, vx_h, vy_h, vz_h, dt, m, g, b): + """真蛙跳一步:x(t), v(t-dt/2) → x(t+dt), v(t+dt/2) + + - 无阻尼 (DAMPING_FORCE=False 或 b≈0): + 纯保守蛙跳,每步仅 1 次力计算,辛积分器。 + v(t+dt/2) = v(t-dt/2) + a_c(t)·dt + + - 有阻尼 (DAMPING_FORCE=True 且 b≠0): + 半隐式处理阻尼,仍只 1 次力计算,对任意阻尼无条件稳定。 + 利用 v(t) ≈ [v(t-dt/2) + v(t+dt/2)] / 2 解析求解: + v(t+dt/2) = [v(t-dt/2)·(1 - α) + a_c(t)·dt] / (1 + α) + 其中 α = b·dt / (2·m),每个方向独立。 + """ + ax_c, ay_c, az_c = compute_accel_conservative(x, y, z, m, g) + + has_damping = DAMPING_FORCE and np.any(b != 0) + if has_damping: + alpha_x = b[0] * dt / (2.0 * m) + alpha_y = b[1] * dt / (2.0 * m) + alpha_z = b[2] * dt / (2.0 * m) + vx_h_new = (vx_h * (1.0 - alpha_x) + ax_c * dt) / (1.0 + alpha_x) + vy_h_new = (vy_h * (1.0 - alpha_y) + ay_c * dt) / (1.0 + alpha_y) + vz_h_new = (vz_h * (1.0 - alpha_z) + az_c * dt) / (1.0 + alpha_z) + else: + vx_h_new = vx_h + ax_c * dt + vy_h_new = vy_h + ay_c * dt + vz_h_new = vz_h + az_c * dt + + x_new = x + vx_h_new * dt + y_new = y + vy_h_new * dt + z_new = z + vz_h_new * dt + return x_new, y_new, z_new, vx_h_new, vy_h_new, vz_h_new + + def Explicit_Euler_Method(x, y, z, vx, vy, vz, dt, m, g, b): ax, ay, az = compute_acceleration(x, y, z, vx, vy, vz, m, g, b) x = x + vx * dt @@ -1407,15 +1566,16 @@ def apply_motion_update(x, y, z, vx, vy, vz, dt, m, g, b): def apply_fixed_constraints(x, y, z, vx, vy, vz): """Keep fixed degrees of freedom at their initial coordinate with zero speed.""" - fixed = ATOM_FIXED != 0 - positions = np.column_stack((x, y, z)) - velocities = np.column_stack((vx, vy, vz)) - positions = np.where(fixed, ATOM_POSITIONS, positions) - velocities = np.where(fixed, 0.0, velocities) - return ( - positions[:, 0], positions[:, 1], positions[:, 2], - velocities[:, 0], velocities[:, 1], velocities[:, 2], - ) + if FIXED_MASK_X is not None and np.any(FIXED_MASK_X): + x[FIXED_MASK_X] = ATOM_POSITIONS[FIXED_MASK_X, 0] + vx[FIXED_MASK_X] = 0.0 + if FIXED_MASK_Y is not None and np.any(FIXED_MASK_Y): + y[FIXED_MASK_Y] = ATOM_POSITIONS[FIXED_MASK_Y, 1] + vy[FIXED_MASK_Y] = 0.0 + if FIXED_MASK_Z is not None and np.any(FIXED_MASK_Z): + z[FIXED_MASK_Z] = ATOM_POSITIONS[FIXED_MASK_Z, 2] + vz[FIXED_MASK_Z] = 0.0 + return x, y, z, vx, vy, vz def wrap_position(x, y, z): """边界回绕( dynamics 模式)。""" @@ -1450,12 +1610,25 @@ def run_simulation(save_trajectory=0): x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, 0.0, 0, DRIVER_DATA, DT) + # 真蛙跳初始化:从 v(0) 反推 v(-dt/2),使后续每步只需 1 次力计算。 + # 其他方法(euler/midpoint)vx/vy/vz 始终存整步速度,不受影响。 + if METHOD == "leapfrog": + _ax0, _ay0, _az0 = compute_accel_conservative(x, y, z, ATOM_MASSES, G) + vx = vx - 0.5 * _ax0 * DT + vy = vy - 0.5 * _ay0 * DT + vz = vz - 0.5 * _az0 * DT + # 从此 vx/vy/vz 存 v(t-dt/2)(蛙跳半步速度) + if warmup_steps is not None and warmup_steps > 0: print(f"[compute] 预热阶段: 前 {warmup_steps} 步不记录") for step in trange(warmup_steps, desc="[compute] 预热"): t = (step + 1) * DT x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, t, step, DRIVER_DATA, DT) - x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) + if METHOD == "leapfrog": + x, y, z, vx, vy, vz = leapfrog_staggered_step( + x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) + else: + x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) x, y, z = wrap_position(x, y, z) x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) print( @@ -1467,8 +1640,6 @@ def run_simulation(save_trajectory=0): record_steps = NT - (warmup_steps or 0) n_atoms = len(ATOM_IDS) n_frames = (record_steps + NSTEP - 1) // NSTEP - frame_indices = [] - # 按 NSTEP 抽帧的临时缓冲区(远小于全量轨迹) sampled_x = np.zeros((n_frames, n_atoms), dtype=np.float64) sampled_y = np.zeros((n_frames, n_atoms), dtype=np.float64) @@ -1507,21 +1678,23 @@ def run_simulation(save_trajectory=0): sampled_vx[fi] = vx sampled_vy[fi] = vy sampled_vz[fi] = vz - frame_indices.append(step) - - x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) + if METHOD == "leapfrog": + x, y, z, vx, vy, vz = leapfrog_staggered_step( + x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) + else: + x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B) x, y, z = wrap_position(x, y, z) x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) # 写入 display.txt(新格式) output_dir = get_output_dir() disp_path = os.path.join(output_dir, "display.txt") - n_frames_actual = len(frame_indices) + n_frames_actual = (record_steps + NSTEP - 1) // NSTEP save_display_txt( disp_path, sampled_x[:n_frames_actual], sampled_y[:n_frames_actual], sampled_z[:n_frames_actual], sampled_vx[:n_frames_actual], sampled_vy[:n_frames_actual], sampled_vz[:n_frames_actual], - np.array(ATOM_IDS), n_frames_actual, n_atoms, + np.array(ATOM_IDS), record_steps, n_atoms, header_fields={"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD), "warmup_steps": str(warmup_steps or 0), "dynamic_steps": str(record_steps), @@ -1544,14 +1717,65 @@ def run_simulation(save_trajectory=0): "driving_force": str(DRIVING_FORCE), "use_marker": str(use_marker), "alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])), + "atom_masses": json.dumps([float(v) for v in ATOM_MASSES]), + "atom_positions": json.dumps(ATOM_POSITIONS.tolist()), + "bond_pairs": json.dumps(BOND_PAIRS.tolist() if BOND_PAIRS is not None else []), + "bond_stiffness": json.dumps(BOND_STIFFNESS.tolist() if BOND_STIFFNESS is not None else []), + "bond_rest_lengths": json.dumps(BOND_REST_LENGTHS.tolist() if BOND_REST_LENGTHS is not None else []), + "G": json.dumps(G.tolist() if G is not None else [0.0, 0.0, 0.0]), "atom_radii": ",".join(str(r) for r in ATOM_RADII), - "camera_distance": str(config.get("camera_distance", 40.0)), - "camera_elevation": str(config.get("camera_elevation", 0)), - "camera_azimuth": str(config.get("camera_azimuth", 0)), + "camera_distance": str(camera_distance), + "camera_elevation": str(camera_elevation), + "camera_azimuth": str(camera_azimuth), "camera_keyframes": str(camera_keyframes_raw)} ) print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)") + # 同时保存二进制 display.npz(I/O 速度提升 5-10x) + _npz_path = os.path.join(output_dir, "display.npz") + _hdr = {"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD), + "warmup_steps": str(warmup_steps or 0), + "dynamic_steps": str(record_steps), + "T_total": str(NT * DT), + "X_MAX": str(X_MAX), "X_MIN": str(X_MIN), + "Y_MAX": str(Y_MAX), "Y_MIN": str(Y_MIN), + "Z_MAX": str(Z_MAX), "Z_MIN": str(Z_MIN), + "ball_radius": str(ball_radius), + "ball_color_r": str(ball_color_r), + "ball_color_g": str(ball_color_g), + "ball_color_b": str(ball_color_b), + "box_color_r": str(box_color_r), + "box_color_g": str(box_color_g), + "box_color_b": str(box_color_b), + "gravity_field": str(GRAVITY_FIELD), + "gravity_interaction": str(GRAVITY_INTERACTION), + "elastic_force": str(ELASTIC_FORCE), + "damping_force": str(DAMPING_FORCE), + "gravity_strength": str(GRAVITY_STRENGTH), + "driving_force": str(DRIVING_FORCE), + "use_marker": str(use_marker), + "alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])), + "atom_masses": json.dumps([float(v) for v in ATOM_MASSES]), + "atom_positions": json.dumps(ATOM_POSITIONS.tolist()), + "bond_pairs": json.dumps(BOND_PAIRS.tolist() if BOND_PAIRS is not None else []), + "bond_stiffness": json.dumps(BOND_STIFFNESS.tolist() if BOND_STIFFNESS is not None else []), + "bond_rest_lengths": json.dumps(BOND_REST_LENGTHS.tolist() if BOND_REST_LENGTHS is not None else []), + "G": json.dumps(G.tolist() if G is not None else [0.0, 0.0, 0.0]), + "atom_radii": ",".join(str(r) for r in ATOM_RADII), + "camera_distance": str(camera_distance), + "camera_elevation": str(camera_elevation), + "camera_azimuth": str(camera_azimuth), + "camera_keyframes": str(camera_keyframes_raw), + "number_of_frames": str(record_steps), + "number_of_particles": str(n_atoms)} + save_display_npz(_npz_path, + sampled_x[:n_frames_actual], sampled_y[:n_frames_actual], + sampled_z[:n_frames_actual], + sampled_vx[:n_frames_actual], sampled_vy[:n_frames_actual], + sampled_vz[:n_frames_actual], + np.array(ATOM_IDS), header_fields=_hdr) + print(f"[compute] display.npz 已保存至: {_npz_path}") + # 可选:保存完整 trajectory.txt if save_trajectory: save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz) diff --git a/draw.py b/draw.py index 574712d..2c5d08f 100644 --- a/draw.py +++ b/draw.py @@ -30,16 +30,21 @@ else: output_dir = compute.get_output_dir(script_dir) os.environ["DYNAMICS_OUTPUT_DIR"] = output_dir disp_path = os.path.join(output_dir, "display.txt") +npz_path = os.path.join(output_dir, "display.npz") -if not os.path.exists(disp_path): +if not os.path.exists(npz_path) and not os.path.exists(disp_path): raise FileNotFoundError( - f"找不到 display.txt!\n" - f"期望路径: {disp_path}\n" + f"找不到 display.npz 或 display.txt!\n" + f"期望路径: {output_dir}\n" f"请先运行 compute.py 计算轨迹,再运行 sample.py 生成显示数组。\n" f"用法: python draw.py [案例输出目录]" ) -disp_data = compute.load_display_txt(disp_path) +# 优先读二进制 npz(加载速度约快 5-10x) +if os.path.exists(npz_path): + disp_data = compute.load_display_npz(npz_path) +else: + disp_data = compute.load_display_txt(disp_path) h = disp_data["header_fields"] # 全原子帧数据 @@ -94,7 +99,7 @@ try: alpha_list = [float(x) for x in raw_alpha.split(",")] if len(alpha_list) != 6: alpha_list = alpha_list * 6 -except: +except (ValueError, AttributeError): alpha_list = [float(raw_alpha)] * 6 # 绘图参数 @@ -514,8 +519,9 @@ def handle_mouse_press(event): def _update_atom_positions(f_idx): """更新所有原子到第 f_idx 帧的位置。""" if USE_MARKER: - for i in range(N_ATOMS): - marker_pos[i] = [DISP_ALL_X[f_idx, i], DISP_ALL_Y[f_idx, i], DISP_ALL_Z[f_idx, i]] + marker_pos[:, 0] = DISP_ALL_X[f_idx] + marker_pos[:, 1] = DISP_ALL_Y[f_idx] + marker_pos[:, 2] = DISP_ALL_Z[f_idx] balls.set_data(pos=marker_pos) else: for i in range(N_ATOMS): diff --git a/dynamics.py b/dynamics.py index d601b51..b6a83fb 100644 --- a/dynamics.py +++ b/dynamics.py @@ -14,6 +14,7 @@ import sys import subprocess import time import argparse +import json from contextlib import contextmanager from pathlib import Path @@ -32,6 +33,13 @@ def _fmt_alpha(v): return str(float(v)) +def _json_field(value): + """Serialize arrays/lists for display header metadata.""" + if isinstance(value, np.ndarray): + value = value.tolist() + return json.dumps(value, ensure_ascii=False) + + def _load_camera_kf(config, runtime_base): """加载 move_camera.txt(速度段格式)→ JSON 字符串。""" import re, json @@ -295,6 +303,12 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output", "driving_force": str(data.get("driving_force", 0)), "use_marker": str(config.get("use_marker", 0)), "alpha": _fmt_alpha(data.get("alpha", 0.2)), + "atom_masses": _json_field(data.get("atom_masses", [])), + "atom_positions": _json_field(data.get("atom_positions", [])), + "bond_pairs": _json_field(data.get("bond_pairs", [])), + "bond_stiffness": _json_field(data.get("bond_stiffness", [])), + "bond_rest_lengths": _json_field(data.get("bond_rest_lengths", [])), + "G": _json_field(data.get("G", [0.0, 0.0, 0.0])), "atom_radii": _fmt_alpha(data.get("atom_radii", [])), "camera_distance": str(config.get("camera_distance", 40.0)), "camera_elevation": str(config.get("camera_elevation", 0)), @@ -320,137 +334,8 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output", # 4. 绘图(可选) if not no_plot and config.get("step_plot", 1): - try: - import matplotlib.pyplot as plt - - # 配置中文字体支持 - plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans'] - plt.rcParams['axes.unicode_minus'] = False - - time_arr = np.arange(NT) * DT - n_atoms = all_x.shape[1] - atom_ids_list = data.get("atom_ids", np.arange(n_atoms) + 1) - - fig, axes = plt.subplots(3, 3, figsize=(15, 13)) - fig.suptitle("轨迹与能量分析", fontsize=16) - - # ── 位置 / 速度 6 子图(前 2 行,每行 3 列) ── - plot_configs = [ - (axes[0, 0], all_x, "x - 时间"), - (axes[0, 1], all_y, "y - 时间"), - (axes[0, 2], all_z, "z - 时间"), - (axes[1, 0], all_vx, "vx - 时间"), - (axes[1, 1], all_vy, "vy - 时间"), - (axes[1, 2], all_vz, "vz - 时间"), - ] - - colors = plt.cm.tab10(np.linspace(0, 1, n_atoms)) - - for ax, data_arr, title in plot_configs: - for i in range(n_atoms): - atom_id = int(atom_ids_list[i]) - ax.plot(time_arr, data_arr[:, i], color=colors[i], linewidth=1.5, label=f"原子 {atom_id}") - ax.set_title(title) - ax.set_xlabel("时间 (s)") - ax.grid(True, alpha=0.3) - ax.legend() - - # ── 能量计算 ───────────────────────────────────── - masses = np.array(data["atom_masses"]) # (n_atoms,) - G_vec = np.array(data.get("G", [0.0, 0.0, -9.8])) # [gx, gy, gz] - gravity_field_enabled = int(data.get("gravity_field", 1)) - gravity_interaction_enabled = int(data.get("gravity_interaction", 0)) - gravity_strength = float(data.get("gravity_strength", 1.0)) - elastic_force_enabled = int(data.get("elastic_force", 1)) - damping_force_enabled = int(data.get("damping_force", 0)) - - # 动能 Ek = ½ m v² - ek = 0.5 * masses[np.newaxis, :] * (all_vx**2 + all_vy**2 + all_vz**2) - - # 均匀重力场势能 Ug = -m G·r - ug = np.zeros_like(ek) - if gravity_field_enabled: - ug = -masses[np.newaxis, :] * ( - G_vec[0] * all_x + G_vec[1] * all_y + G_vec[2] * all_z - ) - - # 弹性势能 Us = ½ k (d - d₀)² - us = np.zeros_like(ek) - bond_pairs = data.get("bond_pairs") - bond_stiffness = data.get("bond_stiffness") - bond_rest_lengths = data.get("bond_rest_lengths") - if elastic_force_enabled and bond_pairs is not None and len(bond_pairs) > 0: - for b_idx in range(len(bond_pairs)): - i, j = bond_pairs[b_idx] - dx = all_x[:, j] - all_x[:, i] - dy = all_y[:, j] - all_y[:, i] - dz = all_z[:, j] - all_z[:, i] - dist = np.sqrt(dx**2 + dy**2 + dz**2) - stretch = dist - bond_rest_lengths[b_idx] - us_each = 0.5 * bond_stiffness[b_idx] * stretch**2 - us[:, i] += us_each # 将整根键的势能记给 i - - # 万有引力势能 Ug_grav = -G_grav * m_i * m_j / r - ug_grav = np.zeros_like(ek) - if gravity_interaction_enabled: - n_atoms_en = len(masses) - for i in range(n_atoms_en): - for j in range(i + 1, n_atoms_en): - dx = all_x[:, j] - all_x[:, i] - dy = all_y[:, j] - all_y[:, i] - dz = all_z[:, j] - all_z[:, i] - dist = np.sqrt(dx**2 + dy**2 + dz**2) - dist = np.maximum(dist, 1e-12) - pair_pe = -gravity_strength * masses[i] * masses[j] / dist - ug_grav[:, i] += 0.5 * pair_pe - ug_grav[:, j] += 0.5 * pair_pe - - # 各原子总能量 - e_total = ek + ug + us + ug_grav # (NT, n_atoms) - - # 系统能量分量 - ek_sys = np.sum(ek, axis=1) - ug_sys = np.sum(ug, axis=1) - us_sys = np.sum(us, axis=1) - ug_grav_sys = np.sum(ug_grav, axis=1) - e_sys = ek_sys + ug_sys + us_sys + ug_grav_sys - - # ── 第 3 行左:各原子总能量 ── - ax_e = axes[2, 0] - for i in range(n_atoms): - aid = int(atom_ids_list[i]) - ax_e.plot(time_arr, e_total[:, i], color=colors[i], linewidth=1.5, label=f"原子 {aid}") - ax_e.set_title("各原子总能量") - ax_e.set_xlabel("时间 (s)") - ax_e.set_ylabel("能量") - ax_e.grid(True, alpha=0.3) - ax_e.legend(loc="upper right") - - # ── 第 3 行右:系统总能量 ── - ax_sys = axes[2, 1] - ax_sys.plot(time_arr, ek_sys, 'b-', linewidth=1.5, label="系统动能") - ax_sys.plot(time_arr, ug_sys, 'g-', linewidth=1.5, label="均匀重力势能") - if elastic_force_enabled and bond_pairs is not None and len(bond_pairs) > 0: - ax_sys.plot(time_arr, us_sys, color='orange', linewidth=1.5, label="系统弹性势能") - if gravity_interaction_enabled: - ax_sys.plot(time_arr, ug_grav_sys, color='purple', linewidth=1.5, label="万有引力势能") - ax_sys.plot(time_arr, e_sys, 'r--', linewidth=1.5, label="系统总能量") - ax_sys.set_title("系统总能量") - ax_sys.set_xlabel("时间 (s)") - ax_sys.set_ylabel("能量") - ax_sys.grid(True, alpha=0.3) - ax_sys.legend(loc="upper right") - - # 隐藏第 3 行第 3 列空子图 - axes[2, 2].set_visible(False) - - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - plot_path = os.path.join(output_dir_abs, "trajectory_plots.png") - plt.savefig(plot_path, dpi=300, bbox_inches="tight") - print(f"[run] 轨迹图表已保存至: {plot_path}") - plt.show() - except ImportError: - print("[run] 警告: 未安装 matplotlib,跳过绘图") + print("[run] 注意: 旧版 step_plot 绘图路径依赖完整轨迹局部变量,当前已暂时跳过。") + print("[run] 如需波形/能量动画,请使用 step_plot_wave: 1。") print(f"[run] 完成!输出目录: {output_dir_abs}") diff --git a/engines/c/main.c b/engines/c/main.c index 0e4c0af..06f30b5 100644 --- a/engines/c/main.c +++ b/engines/c/main.c @@ -533,6 +533,20 @@ static void compute_acceleration( } } +/* 保守力加速度(不含阻尼),供真蛙跳法专用。 + 通过传入零速度调用 compute_acceleration,阻尼项 -B*v/m 自动为零。 */ +static void compute_accel_conservative( + int n, const double *x, const double *y, const double *z, + const double *m, const double G[3], + const BondData *bonds, + double *ax, double *ay, double *az) +{ + double *v0 = (double*)alloca(n * sizeof(double)); + for (int i = 0; i < n; i++) v0[i] = 0.0; + double Bzero[3] = {0.0, 0.0, 0.0}; + compute_acceleration(n, x, y, z, v0, v0, v0, m, G, Bzero, bonds, ax, ay, az); +} + /* 边界条件:clamp 位置 + 速度反转 ——与 Python Limit_in_box 一致 */ static void limit_in_box(double *pos, double *vel, double lo, double hi) { if (*pos > hi) { *pos = hi; *vel = -*vel; } @@ -650,6 +664,15 @@ static void midpoint_step( } /* ── 蛙跳法(Velocity-Verlet)── */ +/* 真蛙跳一步:x(t), v(t-dt/2) → x(t+dt), v(t+dt/2) + * + * 无阻尼:纯保守蛙跳,每步 1 次力计算,辛积分器。 + * v(t+dt/2) = v(t-dt/2) + a_c(t)·dt + * + * 有阻尼:半隐式处理,仍 1 次力计算,对任意阻尼无条件稳定。 + * 利用 v(t) ≈ [v(t-dt/2) + v(t+dt/2)] / 2 解析求解: + * v(t+dt/2) = [v(t-dt/2)·(1-α) + a_c(t)·dt] / (1+α),α = B·dt/(2m) + */ static void leapfrog_step( int n, double *x, double *y, double *z, double *vx, double *vy, double *vz, @@ -659,47 +682,30 @@ static void leapfrog_step( double *ax = (double*)alloca(n * sizeof(double)); double *ay = (double*)alloca(n * sizeof(double)); double *az = (double*)alloca(n * sizeof(double)); - compute_acceleration(n, x, y, z, vx, vy, vz, m, G, B, bonds, ax, ay, az); - /* 半推速度:v_half = v + 0.5*a*dt */ + /* 1 次保守力计算(不含阻尼) */ + compute_accel_conservative(n, x, y, z, m, G, bonds, ax, ay, az); + + int has_damping = g_damping_force && (B[0] != 0.0 || B[1] != 0.0 || B[2] != 0.0); + for (int i = 0; i < n; i++) { if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - vx[i] += ax[i] * dt * 0.5; - vy[i] += ay[i] * dt * 0.5; - vz[i] += az[i] * dt * 0.5; - } - - /* 全推位置(不含边界)*/ - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - x[i] += vx[i] * dt; /* vx 此时是 v_half */ + if (has_damping) { + double alphax = B[0] * dt / (2.0 * m[i]); + double alphay = B[1] * dt / (2.0 * m[i]); + double alphaz = B[2] * dt / (2.0 * m[i]); + vx[i] = (vx[i] * (1.0 - alphax) + ax[i] * dt) / (1.0 + alphax); + vy[i] = (vy[i] * (1.0 - alphay) + ay[i] * dt) / (1.0 + alphay); + vz[i] = (vz[i] * (1.0 - alphaz) + az[i] * dt) / (1.0 + alphaz); + } else { + vx[i] += ax[i] * dt; + vy[i] += ay[i] * dt; + vz[i] += az[i] * dt; + } + x[i] += vx[i] * dt; y[i] += vy[i] * dt; z[i] += vz[i] * dt; } - - /* 显式预测器:v_pred = v_half + 0.5*a_old*dt,用第一次加速度外推半步 - 包含重力+阻尼+弹簧的所有贡献(标准 Velocity-Verlet 预测步)*/ - double *pred_vx = (double*)alloca(n * sizeof(double)); - double *pred_vy = (double*)alloca(n * sizeof(double)); - double *pred_vz = (double*)alloca(n * sizeof(double)); - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - pred_vx[i] = vx[i] + 0.5 * ax[i] * dt; - pred_vy[i] = vy[i] + 0.5 * ay[i] * dt; - pred_vz[i] = vz[i] + 0.5 * az[i] * dt; - } - - /* 用新位置 + 预测速度重算加速度 */ - compute_acceleration(n, x, y, z, pred_vx, pred_vy, pred_vz, m, G, B, bonds, ax, ay, az); - - /* 速度后半步:v = v_half + 0.5*a_next*dt - vx 仍为 v_half(未被覆盖)*/ - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - vx[i] += ax[i] * dt * 0.5; - vy[i] += ay[i] * dt * 0.5; - vz[i] += az[i] * dt * 0.5; - } } /* ── 驱动力(与 Python apply_driving_force 一致)──────────────── */ @@ -895,12 +901,13 @@ static void write_display_txt(const char *path, const Trajectory *traj, FILE *f = fopen(path, "w"); if (!f) die("无法写入 display.txt"); - int n_frames = traj->n_steps; + int n_frames = traj->n_steps; /* 实际采样帧数,用于下面的帧循环 */ int n_particles = traj->n_atoms; int dynamic_steps = params->NT - params->warmup_steps; double T_total = dynamic_steps * params->DT; - fprintf(f, "number of frames: %d\n", n_frames); + /* number of frames 写总积分步数(与 draw.py NT 对应),不是采样帧数 */ + fprintf(f, "number of frames: %d\n", dynamic_steps); fprintf(f, "number of particles: %d\n", n_particles); fprintf(f, "DT: %.16g\n", params->DT); fprintf(f, "NSTEP: %d\n", params->NSTEP); @@ -1015,6 +1022,20 @@ int main(int argc, char **argv) { traj.vz = traj.vy + sampled_steps * n; } + /* 真蛙跳初始化:v(0) 反推 v(-dt/2) = v(0) - 0.5*a_c(0)*dt */ + if (strcmp(params.method, "leapfrog") == 0) { + double *ax0 = (double*)alloca(n * sizeof(double)); + double *ay0 = (double*)alloca(n * sizeof(double)); + double *az0 = (double*)alloca(n * sizeof(double)); + compute_accel_conservative(n, x, y, z, atoms.masses, params.G, &bonds, ax0, ay0, az0); + for (int i = 0; i < n; i++) { + if (atoms.fixed[i*3+0] && atoms.fixed[i*3+1] && atoms.fixed[i*3+2]) continue; + vx[i] -= 0.5 * ax0[i] * params.DT; + vy[i] -= 0.5 * ay0[i] * params.DT; + vz[i] -= 0.5 * az0[i] * params.DT; + } + } + /* 预热 */ /* 初始时刻 t=0 驱动力(与 Python run_simulation 一致)*/ if (params.driving_force) apply_driving_force(n, x, y, z, vx, vy, vz, 0.0, 0, params.DT, &drivers); diff --git a/engines/cpp/main.cpp b/engines/cpp/main.cpp index 03e7e30..b9602a9 100644 --- a/engines/cpp/main.cpp +++ b/engines/cpp/main.cpp @@ -420,6 +420,26 @@ static void compute_acceleration( } } +/* 保守力加速度(不含阻尼),供真蛙跳法专用。 + 传入 damping_force=0 使 compute_acceleration 跳过阻尼项。 */ +static void compute_accel_conservative( + int n, + const double *x, const double *y, const double *z, + const double *m, const double G[3], + const BondData &bonds, + int gravity_field, int gravity_interaction, + int elastic_force, double gravity_strength, + double *ax, double *ay, double *az) +{ + std::vector v0(n, 0.0); + double Bzero[3] = {0.0, 0.0, 0.0}; + compute_acceleration(n, x, y, z, v0.data(), v0.data(), v0.data(), + m, G, Bzero, bonds, + gravity_field, gravity_interaction, + elastic_force, 0, gravity_strength, + ax, ay, az); +} + /* 边界条件:clamp 位置 + 速度反转 ——与 Python Limit_in_box 一致 */ static void limit_in_box(double &pos, double &vel, double lo, double hi) { if (pos > hi) { pos = hi; vel = -vel; } @@ -543,6 +563,14 @@ static void midpoint_step( } /* ── 蛙跳法(Velocity-Verlet)——与 Python Leapfrog_Method 一致 ── */ +/* 真蛙跳一步:x(t), v(t-dt/2) → x(t+dt), v(t+dt/2) + * + * 无阻尼:纯保守蛙跳,每步 1 次力计算,辛积分器。 + * v(t+dt/2) = v(t-dt/2) + a_c(t)·dt + * + * 有阻尼:半隐式处理,仍 1 次力计算,对任意阻尼无条件稳定。 + * v(t+dt/2) = [v(t-dt/2)·(1-α) + a_c(t)·dt] / (1+α),α = B·dt/(2m) + */ static void leapfrog_full_step( int n, double *x, double *y, double *z, double *vx, double *vy, double *vz, @@ -552,54 +580,33 @@ static void leapfrog_full_step( int elastic_force, int damping_force, double gravity_strength) { - // 第一次加速度 std::vector ax(n), ay(n), az(n); - compute_acceleration(n, x, y, z, vx, vy, vz, m, G, B, bonds, - gravity_field, gravity_interaction, - elastic_force, damping_force, gravity_strength, - ax.data(), ay.data(), az.data()); - // 半推速度:v_half = v + 0.5*a*dt (存入 vx, vy, vz) + // 1 次保守力计算(不含阻尼) + compute_accel_conservative(n, x, y, z, m, G, bonds, + gravity_field, gravity_interaction, + elastic_force, gravity_strength, + ax.data(), ay.data(), az.data()); + + bool has_damping = damping_force && (B[0] != 0.0 || B[1] != 0.0 || B[2] != 0.0); for (int i = 0; i < n; i++) { if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - vx[i] += ax[i] * dt * 0.5; - vy[i] += ay[i] * dt * 0.5; - vz[i] += az[i] * dt * 0.5; - } - - // 全推位置(不含边界,边界在外层统一处理) - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - x[i] += vx[i] * dt; // vx 此时是 v_half + if (has_damping) { + double alphax = B[0] * dt / (2.0 * m[i]); + double alphay = B[1] * dt / (2.0 * m[i]); + double alphaz = B[2] * dt / (2.0 * m[i]); + vx[i] = (vx[i] * (1.0 - alphax) + ax[i] * dt) / (1.0 + alphax); + vy[i] = (vy[i] * (1.0 - alphay) + ay[i] * dt) / (1.0 + alphay); + vz[i] = (vz[i] * (1.0 - alphaz) + az[i] * dt) / (1.0 + alphaz); + } else { + vx[i] += ax[i] * dt; + vy[i] += ay[i] * dt; + vz[i] += az[i] * dt; + } + x[i] += vx[i] * dt; y[i] += vy[i] * dt; z[i] += vz[i] * dt; } - - // 显式预测器:v_pred = v_half + 0.5*a_old*dt,用第一次加速度外推半步 - // 包含所有力的贡献(标准 Velocity-Verlet 预测步) - std::vector pred_vx(n), pred_vy(n), pred_vz(n); - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - pred_vx[i] = vx[i] + 0.5 * ax[i] * dt; - pred_vy[i] = vy[i] + 0.5 * ay[i] * dt; - pred_vz[i] = vz[i] + 0.5 * az[i] * dt; - } - - // 用新位置 + 预测速度重算加速度 - compute_acceleration(n, x, y, z, pred_vx.data(), pred_vy.data(), pred_vz.data(), - m, G, B, bonds, - gravity_field, gravity_interaction, - elastic_force, damping_force, gravity_strength, - ax.data(), ay.data(), az.data()); - - // 速度后半步:v = v_half + 0.5*a_next*dt - // vx 仍为 v_half(未被覆盖),直接加上 0.5*a_next*dt - for (int i = 0; i < n; i++) { - if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; - vx[i] += ax[i] * dt * 0.5; - vy[i] += ay[i] * dt * 0.5; - vz[i] += az[i] * dt * 0.5; - } } /* ── 分发器:调用对应积分方法 + 边界条件(与 Python apply_motion_update 一致)── */ @@ -737,15 +744,17 @@ static void write_display_txt( if (!f) die("无法写入 " + path); std::cout << "[Cpp-engine] 正在写入显示数据…" << std::endl; - double T_total = params.NT * params.DT; + int dynamic_steps = params.NT - params.warmup_steps; + double T_total = dynamic_steps * params.DT; - f << "number of frames: " << n_steps << "\n"; + /* number of frames 写总积分步数(与 draw.py NT 对应),不是采样帧数 */ + f << "number of frames: " << dynamic_steps << "\n"; f << "number of particles: " << n_atoms << "\n"; f << "DT: " << params.DT << "\n"; f << "NSTEP: " << params.NSTEP << "\n"; f << "method: " << params.method << "\n"; f << "warmup_steps: " << params.warmup_steps << "\n"; - f << "dynamic_steps: " << (params.NT - params.warmup_steps) << "\n"; + f << "dynamic_steps: " << dynamic_steps << "\n"; f << "T_total: " << T_total << "\n"; f << "box_a: " << params.box_a << "\n"; f << "driving_force: " << params.driving_force << "\n"; @@ -925,6 +934,21 @@ int main(int argc, char **argv) { std::vector traj_vy(buf_steps * n); std::vector traj_vz(buf_steps * n); + // 真蛙跳初始化:v(0) 反推 v(-dt/2) = v(0) - 0.5*a_c(0)*dt + if (params.method == "leapfrog") { + std::vector ax0(n), ay0(n), az0(n); + compute_accel_conservative(n, x.data(), y.data(), z.data(), atoms.masses.data(), params.G, bonds, + params.gravity_field, params.gravity_interaction, + params.elastic_force, params.gravity_strength, + ax0.data(), ay0.data(), az0.data()); + for (int i = 0; i < n; i++) { + if (atoms.fixed[i*3+0] && atoms.fixed[i*3+1] && atoms.fixed[i*3+2]) continue; + vx[i] -= 0.5 * ax0[i] * params.DT; + vy[i] -= 0.5 * ay0[i] * params.DT; + vz[i] -= 0.5 * az0[i] * params.DT; + } + } + // 预热 // 初始时刻 t=0 驱动力(与 Python run_simulation 一致) if (params.driving_force) diff --git a/engines/fortran/main.f90 b/engines/fortran/main.f90 index fa26955..46e20d7 100644 --- a/engines/fortran/main.f90 +++ b/engines/fortran/main.f90 @@ -107,6 +107,24 @@ program dynamics_f90 allocate(traj_x(record_steps, n), traj_y(record_steps, n), traj_z(record_steps, n)) allocate(traj_vx(record_steps, n), traj_vy(record_steps, n), traj_vz(record_steps, n)) + ! 真蛙跳初始化:v(0) 反推 v(-dt/2) = v(0) - 0.5*a_c(0)*dt + if (trim(method) == 'leapfrog') then + block + double precision :: ax0(n), ay0(n), az0(n) + integer :: ii + call accel_conservative(n, x, y, z, masses, G, & + n_bonds, bond_pairs, bond_stiffness, bond_rest_lengths, & + gravity_field, gravity_interaction, & + elastic_force, gravity_strength, ax0, ay0, az0) + do ii = 1, n + if (fixed(ii,1) /= 0 .and. fixed(ii,2) /= 0 .and. fixed(ii,3) /= 0) cycle + vx(ii) = vx(ii) - 0.5d0 * ax0(ii) * DT + vy(ii) = vy(ii) - 0.5d0 * ay0(ii) * DT + vz(ii) = vz(ii) - 0.5d0 * az0(ii) * DT + end do + end block + end if + ! 预热 ! 初始时刻 t=0 驱动力(与 Python run_simulation 一致) if (driving_force /= 0 .and. n_drivers > 0) then @@ -524,6 +542,24 @@ pure subroutine accel(n, x, y, z, vx, vy, vz, m, G, B, & end if end subroutine accel +! 保守力加速度(不含阻尼),供真蛙跳法专用。 +! 传入零速度、零 B 调用 accel,阻尼项 -B*v/m 自动为零。 +subroutine accel_conservative(n, x, y, z, m, G, nb, bp, bk, br, & + gravity_field, gravity_interaction, & + elastic_force, gravity_strength, ax, ay, az) + integer, intent(in) :: n, nb, bp(nb, 2) + integer, intent(in) :: gravity_field, gravity_interaction, elastic_force + double precision, intent(in) :: x(n), y(n), z(n), m(n), G(3) + double precision, intent(in) :: bk(nb), br(nb), gravity_strength + double precision, intent(out) :: ax(n), ay(n), az(n) + double precision :: v0(n), B0(3) + v0 = 0.0d0 + B0 = 0.0d0 + call accel(n, x, y, z, v0, v0, v0, m, G, B0, nb, bp, bk, br, & + gravity_field, gravity_interaction, & + elastic_force, 0, gravity_strength, ax, ay, az) +end subroutine accel_conservative + ! 边界条件:clamp 位置 + 速度反转 subroutine limit_in_box(pos, vel, lo, hi) double precision, intent(inout) :: pos, vel @@ -652,7 +688,13 @@ subroutine midpoint_step(n, x, y, z, vx, vy, vz, m, G, B, & end do end subroutine midpoint_step -! ── 蛙跳法(Velocity-Verlet)── +! 真蛙跳一步:x(t), v(t-dt/2) → x(t+dt), v(t+dt/2) +! +! 无阻尼:纯保守蛙跳,每步 1 次力计算,辛积分器。 +! v(t+dt/2) = v(t-dt/2) + a_c(t)*dt +! +! 有阻尼:半隐式处理,仍 1 次力计算,对任意阻尼无条件稳定。 +! v(t+dt/2) = [v(t-dt/2)*(1-α) + a_c(t)*dt] / (1+α),α = B*dt/(2m) subroutine leapfrog_full(n, x, y, z, vx, vy, vz, m, G, B, & nb, bp, bk, br, fixed, dt, & gravity_field, gravity_interaction, & @@ -662,53 +704,36 @@ subroutine leapfrog_full(n, x, y, z, vx, vy, vz, m, G, B, & double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n) double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), dt, gravity_strength double precision :: ax(n), ay(n), az(n) - double precision :: dmp_vx(n), dmp_vy(n), dmp_vz(n) - double precision :: gx, gy, gz + double precision :: alphax, alphay, alphaz + logical :: has_damping integer :: i - call accel(n, x, y, z, vx, vy, vz, m, G, B, nb, bp, bk, br, & - gravity_field, gravity_interaction, & - elastic_force, damping_force, gravity_strength, ax, ay, az) + ! 1 次保守力计算(不含阻尼) + call accel_conservative(n, x, y, z, m, G, nb, bp, bk, br, & + gravity_field, gravity_interaction, & + elastic_force, gravity_strength, ax, ay, az) + + has_damping = (damping_force /= 0) .and. & + (B(1) /= 0.0d0 .or. B(2) /= 0.0d0 .or. B(3) /= 0.0d0) - ! 速度半步推 - do i = 1, n - if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle - vx(i) = vx(i) + ax(i) * dt * 0.5d0 - vy(i) = vy(i) + ay(i) * dt * 0.5d0 - vz(i) = vz(i) + az(i) * dt * 0.5d0 - end do - - ! 全推位置(不含边界) do i = 1, n if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle + if (has_damping) then + alphax = B(1) * dt / (2.0d0 * m(i)) + alphay = B(2) * dt / (2.0d0 * m(i)) + alphaz = B(3) * dt / (2.0d0 * m(i)) + vx(i) = (vx(i) * (1.0d0 - alphax) + ax(i) * dt) / (1.0d0 + alphax) + vy(i) = (vy(i) * (1.0d0 - alphay) + ay(i) * dt) / (1.0d0 + alphay) + vz(i) = (vz(i) * (1.0d0 - alphaz) + az(i) * dt) / (1.0d0 + alphaz) + else + vx(i) = vx(i) + ax(i) * dt + vy(i) = vy(i) + ay(i) * dt + vz(i) = vz(i) + az(i) * dt + end if x(i) = x(i) + vx(i) * dt y(i) = y(i) + vy(i) * dt z(i) = z(i) + vz(i) * dt end do - - ! 隐式阻尼处理(用临时数组 dmp_v,不覆盖 vx/vy/vz) - do i = 1, n - if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) then - dmp_vx(i) = 0; dmp_vy(i) = 0; dmp_vz(i) = 0; cycle - end if - gx = B(1) / m(i); gy = B(2) / m(i); gz = B(3) / m(i) - dmp_vx(i) = (vx(i) + 0.5d0 * G(1) * dt) / (1.0d0 + 0.5d0 * gx * dt) - dmp_vy(i) = (vy(i) + 0.5d0 * G(2) * dt) / (1.0d0 + 0.5d0 * gy * dt) - dmp_vz(i) = (vz(i) + 0.5d0 * G(3) * dt) / (1.0d0 + 0.5d0 * gz * dt) - end do - - ! 用新位置 + 阻尼速度重算加速度 - call accel(n, x, y, z, dmp_vx, dmp_vy, dmp_vz, m, G, B, nb, bp, bk, br, & - gravity_field, gravity_interaction, & - elastic_force, damping_force, gravity_strength, ax, ay, az) - - ! 速度后半步:v = v_half + 0.5*a_next*dt(vx 仍为 v_half) - do i = 1, n - if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle - vx(i) = vx(i) + ax(i) * dt * 0.5d0 - vy(i) = vy(i) + ay(i) * dt * 0.5d0 - vz(i) = vz(i) + az(i) * dt * 0.5d0 - end do end subroutine leapfrog_full ! ── 分发器 + 边界条件 + 自由度约束 ── diff --git a/examples/case06/input/bond.txt b/examples/case06/input/bond.txt index 2a14f8d..b841478 100644 --- a/examples/case06/input/bond.txt +++ b/examples/case06/input/bond.txt @@ -1,2 +1,2 @@ bond_name k rest_length -k1 10.0 1.0 +k1 50.0 1.0 diff --git a/examples/case06/input/driver.txt b/examples/case06/input/driver.txt index 92c6a38..540ce07 100644 --- a/examples/case06/input/driver.txt +++ b/examples/case06/input/driver.txt @@ -1,2 +1,2 @@ n amp_x amp_y amp_z freq_x freq_y freq_z phi_x phi_y phi_z period -1 0 0 2.0 0 0 0.1 0 0 90 all +1 0 0 2.0 0 0 0.05 0 0 90 all diff --git a/examples/case06/input/input.txt b/examples/case06/input/input.txt index fa039e5..d7ffffe 100644 --- a/examples/case06/input/input.txt +++ b/examples/case06/input/input.txt @@ -19,7 +19,7 @@ save_trajectory: 0 # 0=不保留完整轨迹文件, 1=保留 trajectory.txt( # ── 计算引擎 ────────────────────────────────── # 可选: python, c, cpp, fortran, java -engine: c # 默认使用 python 引擎 +engine: python # 默认使用 python 引擎 # ── 盒子 ────────────────────────────────────── box_a: 80.0 # 立方体半边长,粒子被限制在 [-box_a, box_a]³ 内 @@ -69,10 +69,10 @@ warmup_steps: 0 # 默认 0(立即开始记录) T_total: 100.0 # 抽帧间隔(每 NSTEP 步取一帧用于动画) -NSTEP: 100 +NSTEP: 10 # ── 时间步长 ────────────────────────────────── -DT: 0.001 # 时间步长 (s) +DT: 0.01 # 时间步长 (s) # 抽帧范围:只保存 [sample_start, sample_end) 区间内的帧 sample_start: null # null 表示从头开始(帧索引从 0 起) diff --git a/examples/case06/run_dynamics.py b/examples/case06/run_dynamics.py index f438f63..3a3f2b5 100644 --- a/examples/case06/run_dynamics.py +++ b/examples/case06/run_dynamics.py @@ -31,7 +31,7 @@ def load_dynamics_module(module_path: Path): def main(): - parser = argparse.ArgumentParser(description="运行 Dynamics 示例案例 case01") + parser = argparse.ArgumentParser(description="运行 Dynamics 示例案例 case06") parser.add_argument("--no-plot", action="store_true", help="跳过 matplotlib 绘图") args = parser.parse_args() diff --git a/plot_wave.py b/plot_wave.py index f7ee65f..ca34bf7 100644 --- a/plot_wave.py +++ b/plot_wave.py @@ -14,17 +14,101 @@ import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation import os import sys +import json sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import compute def load_disp_data(output_dir): - """加载 display.txt""" - disp_path = os.path.join(output_dir, "display.txt") - if not os.path.exists(disp_path): - raise FileNotFoundError(f"找不到 {disp_path}") - return compute.load_text_data(disp_path) + """加载 display.npz(优先)或 display.txt。""" + npz_path = os.path.join(output_dir, "display.npz") + txt_path = os.path.join(output_dir, "display.txt") + if os.path.exists(npz_path): + return compute.load_display_npz(npz_path) + if os.path.exists(txt_path): + return compute.load_display_txt(txt_path) + raise FileNotFoundError(f"找不到 display.npz 或 display.txt in {output_dir}") + + +def _header_json(header_fields, key, default): + raw = header_fields.get(key, "") + if not raw: + return default + try: + return json.loads(raw) + except json.JSONDecodeError: + return default + + +def _load_wave_dataset(output_dir): + """Load wave/energy plotting data from display metadata or sibling input files.""" + disp_data = load_disp_data(output_dir) + header = disp_data["header_fields"] + + x = disp_data["frames_x"] + y = disp_data["frames_y"] + z = disp_data["frames_z"] + vx = disp_data["frames_vx"] + vy = disp_data["frames_vy"] + vz = disp_data["frames_vz"] + atom_ids = np.array(disp_data["atom_ids"], dtype=np.int64) + n_frames = x.shape[0] + dt = float(header.get("DT", 0.001)) + nstep = int(header.get("NSTEP", 1)) + t = np.arange(n_frames, dtype=np.float64) * dt * nstep + + masses = np.array(_header_json(header, "atom_masses", []), dtype=np.float64) + pos_0 = np.array(_header_json(header, "atom_positions", []), dtype=np.float64) + bond_pairs = np.array(_header_json(header, "bond_pairs", []), dtype=np.int64) + bond_stiffness = np.array(_header_json(header, "bond_stiffness", []), dtype=np.float64) + bond_rest_lengths = np.array(_header_json(header, "bond_rest_lengths", []), dtype=np.float64) + gravity_vec = _header_json(header, "G", [0.0, 0.0, 0.0]) + + # Backward-compatible fallback for older display.txt outputs. + if masses.size == 0 or pos_0.size == 0: + input_dir = os.path.join(os.path.dirname(output_dir), "input") + coord_path = os.path.join(input_dir, "coord.txt") + if os.path.exists(coord_path): + (_, masses_fb, _, positions_fb, _, _) = compute.load_coord_file(coord_path) + masses = np.array(masses_fb, dtype=np.float64) + pos_0 = np.array(positions_fb, dtype=np.float64) + else: + raise ValueError("display.txt 缺少 atom_masses/atom_positions 元数据,且未找到 input/coord.txt") + + if bond_pairs.size == 0: + input_dir = os.path.join(os.path.dirname(output_dir), "input") + connection_path = os.path.join(input_dir, "connection.txt") + bond_path = os.path.join(input_dir, "bond.txt") + if os.path.exists(connection_path) and os.path.exists(bond_path): + bond_map = compute.load_bond_parameters(bond_path) + pairs_fb, _, stiffness_fb, rest_lengths_fb = compute.load_bond_connections( + connection_path, atom_ids, pos_0, bond_map) + bond_pairs = np.array(pairs_fb, dtype=np.int64) + bond_stiffness = np.array(stiffness_fb, dtype=np.float64) + bond_rest_lengths = np.array(rest_lengths_fb, dtype=np.float64) + + return { + "n_frames": n_frames, + "t": t, + "x": x, + "y": y, + "z": z, + "vx": vx, + "vy": vy, + "vz": vz, + "pos_0": pos_0, + "masses": masses, + "atom_ids": atom_ids, + "bond_pairs": bond_pairs, + "bond_stiffness": bond_stiffness, + "bond_rest_lengths": bond_rest_lengths, + "gravity_field": int(header.get("gravity_field", 0)), + "gravity_interaction": int(header.get("gravity_interaction", 0)), + "gravity_strength": float(header.get("gravity_strength", 1.0)), + "G": gravity_vec, + "driving_force": int(header.get("driving_force", 0)), + } def compute_energy(x, y, z, vx, vy, vz, masses, mass_arr, @@ -91,36 +175,36 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False): save_gif: 是否保存 GIF save_mp4: 是否保存 MP4 """ - data = load_disp_data(output_dir) + data = _load_wave_dataset(output_dir) n_frames = int(data["n_frames"]) - t = np.array(data["disp_t"]) + t = np.array(data["t"]) # 位置 / 速度 - x = np.array(data["disp_all_x"]) - y = np.array(data["disp_all_y"]) - z = np.array(data["disp_all_z"]) - vx = np.array(data["disp_all_vx"]) - vy = np.array(data["disp_all_vy"]) - vz = np.array(data["disp_all_vz"]) + x = np.array(data["x"]) + y = np.array(data["y"]) + z = np.array(data["z"]) + vx = np.array(data["vx"]) + vy = np.array(data["vy"]) + vz = np.array(data["vz"]) # 原子信息 - pos_0 = np.array(data["atom_positions"]) # (n_atoms, 3) - masses = np.array(data["atom_masses"]) + pos_0 = np.array(data["pos_0"]) + masses = np.array(data["masses"]) atom_ids = np.array(data["atom_ids"]) - n_atoms = len(atom_ids) + n_atoms = len(atom_ids) # 成键 - bond_pairs = data.get("bond_pairs", []) - bond_stiffness = np.array(data.get("bond_stiffness", [])) - bond_rest_lengths = np.array(data.get("bond_rest_lengths", [])) + bond_pairs = np.array(data.get("bond_pairs", []), dtype=np.int64) + bond_stiffness = np.array(data.get("bond_stiffness", []), dtype=np.float64) + bond_rest_lengths = np.array(data.get("bond_rest_lengths", []), dtype=np.float64) # 物理开关 - gravity_field = int(data.get("gravity_field", 0)) + gravity_field = int(data.get("gravity_field", 0)) gravity_interaction = int(data.get("gravity_interaction", 0)) - G = data.get("G", [0, 0, 0]) - gravity_strength = float(data.get("gravity_strength", 1.0)) - driving_force = int(data.get("driving_force", 0)) + G = data.get("G", [0, 0, 0]) + gravity_strength = float(data.get("gravity_strength", 1.0)) + driving_force = int(data.get("driving_force", 0)) # ── 位移(偏离初始平衡位形)── dx = x - pos_0[np.newaxis, :, 0] # 纵波(沿链方向 x)