diff --git a/compute.py b/compute.py index 975f858..5de7913 100644 --- a/compute.py +++ b/compute.py @@ -5,8 +5,8 @@ compute.py 功能： 1. 运行 NT 步物理模拟（ kinematics / dynamics 等运动模式） - 2. 将每一步的 (x, y, z, vx, vy, vz) 保存到 output/trajectory.txt - 3. 同时保存所有模拟参数元数据 + 2. 按 NSTEP 抽帧，输出 output/display.txt（新文本格式） + 3. 可选（save_trajectory=1）保留完整轨迹 output/trajectory.txt（JSON） """ import json @@ -118,7 +118,7 @@ def _from_text_value(value): def save_text_data(path, data): - """Save structured simulation data as formatted JSON text.""" + """Save structured simulation data as formatted JSON text (旧格式，保留兼容).""" os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(_to_text_value(data), f, ensure_ascii=False, indent=2) @@ -126,12 +126,150 @@ def save_text_data(path, data): def load_text_data(path): - """Load structured simulation data from JSON text.""" + """Load structured simulation data from JSON text (旧格式).""" with open(path, "r", encoding="utf-8") as f: data = json.load(f) return _from_text_value(data) +# ======================================================================== +# 新 display.txt 格式：纯文本，按帧分块 +# 第1行: number of frames: N +# 第2行: number of particles: M +# 第3行: frame: 1 +# 第4行: n x y z vx vy vz +# 第5+行: 数据行（每个原子一行） +# 重复第3-5行直到所有帧 +# ======================================================================== + +def save_display_txt(path, frames_x, frames_y, frames_z, + frames_vx, frames_vy, frames_vz, + atom_ids, n_total_frames, n_total_particles, + header_fields=None): + """Write display.txt in new text format. + + Args: + path: 输出文件路径 + frames_x/y/z/vx/vy/vz: (n_frames, n_atoms) 数组 + atom_ids: (n_atoms,) 原子编号数组 + n_total_frames: 总帧数（含未采样） + n_total_particles: 总粒子数 + header_fields: 可选的额外元数据字典（写入文件头之后） + """ + os.makedirs(os.path.dirname(path), exist_ok=True) + n_frames = frames_x.shape[0] + n_atoms = frames_x.shape[1] + + # 格式化辅助：固定宽度，6位小数 + def fmt(v): return f"{v:13.6f}" + + with open(path, "w", encoding="utf-8") as f: + f.write(f"number of frames: {n_total_frames}\n") + f.write(f"number of particles: {n_total_particles}\n") + # 写入额外元数据 + if header_fields: + for k, v in header_fields.items(): + f.write(f"{k}: {v}\n") + + for fr in range(n_frames): + f.write(f"\nframe: {fr + 1}\n") + f.write(f"n x y z vx vy vz\n") + for a in range(n_atoms): + f.write(f"{atom_ids[a]:d}" + f"{fmt(frames_x[fr, a])}" + f"{fmt(frames_y[fr, a])}" + f"{fmt(frames_z[fr, a])}" + f"{fmt(frames_vx[fr, a])}" + f"{fmt(frames_vy[fr, a])}" + f"{fmt(frames_vz[fr, a])}\n") + return path + + +def load_display_txt(path): + """Read display.txt new text format into numpy arrays. + + Returns dict with keys: frames_x/y/z/vx/vy/vz, atom_ids, + n_total_frames, n_total_particles, header_fields + """ + header_fields = {} + frames_x, frames_y, frames_z = [], [], [] + frames_vx, frames_vy, frames_vz = [], [], [] + atom_ids = None + n_total_frames = 0 + n_total_particles = 0 + + with open(path, "r", encoding="utf-8") as f: + lines = f.readlines() + + # Parse header + i = 0 + while i < len(lines): + line = lines[i].strip() + if line.startswith("number of frames:"): + n_total_frames = int(line.split(":")[1].strip()) + i += 1 + elif line.startswith("number of particles:"): + n_total_particles = int(line.split(":")[1].strip()) + i += 1 + elif line.startswith("frame:"): + break + else: + # Extra header field + if ":" in line: + k, v = line.split(":", 1) + header_fields[k.strip()] = v.strip() + i += 1 + + # Parse frames + while i < len(lines): + line = lines[i].strip() + if line.startswith("frame:"): + i += 1 # skip column header line + if i < len(lines): + i += 1 # skip "n x y..." + frame_x, frame_y, frame_z = [], [], [] + frame_vx, frame_vy, frame_vz = [], [], [] + cur_ids = [] + while i < len(lines) and not lines[i].strip().startswith("frame:") and lines[i].strip(): + parts = lines[i].strip().split() + if len(parts) >= 7: + cur_ids.append(int(parts[0])) + frame_x.append(float(parts[1])) + frame_y.append(float(parts[2])) + frame_z.append(float(parts[3])) + frame_vx.append(float(parts[4])) + frame_vy.append(float(parts[5])) + frame_vz.append(float(parts[6])) + i += 1 + if frame_x: + frames_x.append(frame_x) + frames_y.append(frame_y) + frames_z.append(frame_z) + frames_vx.append(frame_vx) + frames_vy.append(frame_vy) + frames_vz.append(frame_vz) + if atom_ids is None: + atom_ids = np.array(cur_ids) + else: + i += 1 + + if not frames_x: + raise ValueError(f"{path} 中没有有效帧数据") + + return { + "frames_x": np.array(frames_x), + "frames_y": np.array(frames_y), + "frames_z": np.array(frames_z), + "frames_vx": np.array(frames_vx), + "frames_vy": np.array(frames_vy), + "frames_vz": np.array(frames_vz), + "atom_ids": atom_ids, + "n_total_frames": n_total_frames, + "n_total_particles": n_total_particles, + "header_fields": header_fields, + } + + def get_output_dir(base_dir=None): """Return the output directory used for generated artifacts.""" override = os.environ.get("DYNAMICS_OUTPUT_DIR") @@ -597,7 +735,8 @@ def run_from_config(config, out_dir=None): return None, None, None, None, None, None t_start = time.time() t_start_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation() + traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation( + save_trajectory=int(config.get("save_trajectory", 0))) t_end = time.time() t_end_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") elapsed = t_end - t_start @@ -1213,12 +1352,13 @@ def wrap_position(x, y, z): # 主计算流程 # =========================================================================== -def run_simulation(): - """计算 NT 步轨迹，返回位置/速度数组。 +def run_simulation(save_trajectory=0): + """计算 NT 步轨迹，直接抽帧并保存 display.txt。 步骤控制： - warmup_steps: 预热步数，跳过不记录（用于稳定初始状态） - - 实际记录步数 = NT - warmup_steps + - 按 NSTEP 抽帧保存到 display.txt（新格式） + - save_trajectory=1 时额外保存完整 trajectory.txt（JSON 旧格式） """ # 预热阶段 x = ATOM_POSITIONS[:, 0].copy() @@ -1228,7 +1368,6 @@ def run_simulation(): vy = ATOM_VELOCITIES[:, 1].copy() vz = ATOM_VELOCITIES[:, 2].copy() x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) - # 初始时刻驱动力（t=0 时原子 1 的位置由驱动力决定而非 coord.txt） x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, 0.0, 0, DRIVER_DATA, DT) if warmup_steps is not None and warmup_steps > 0: @@ -1244,33 +1383,93 @@ def run_simulation(): f"({x[PLOT_ATOM_ROW]:.4f}, {y[PLOT_ATOM_ROW]:.4f}, {z[PLOT_ATOM_ROW]:.4f})" ) - # 记录阶段 + # 记录阶段 - 按 NSTEP 抽帧保存 record_steps = NT - (warmup_steps or 0) n_atoms = len(ATOM_IDS) - traj_x = np.zeros((record_steps, n_atoms), dtype=np.float64) - traj_y = np.zeros((record_steps, n_atoms), dtype=np.float64) - traj_z = np.zeros((record_steps, n_atoms), dtype=np.float64) - traj_vx = np.zeros((record_steps, n_atoms), dtype=np.float64) - traj_vy = np.zeros((record_steps, n_atoms), dtype=np.float64) - traj_vz = np.zeros((record_steps, n_atoms), dtype=np.float64) + 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) + sampled_z = np.zeros((n_frames, n_atoms), dtype=np.float64) + sampled_vx = np.zeros((n_frames, n_atoms), dtype=np.float64) + sampled_vy = np.zeros((n_frames, n_atoms), dtype=np.float64) + sampled_vz = np.zeros((n_frames, n_atoms), dtype=np.float64) + + # 如果 save_trajectory，准备完整轨迹缓冲区 + if save_trajectory: + traj_x = np.zeros((record_steps, n_atoms), dtype=np.float64) + traj_y = np.zeros((record_steps, n_atoms), dtype=np.float64) + traj_z = np.zeros((record_steps, n_atoms), dtype=np.float64) + traj_vx = np.zeros((record_steps, n_atoms), dtype=np.float64) + traj_vy = np.zeros((record_steps, n_atoms), dtype=np.float64) + traj_vz = np.zeros((record_steps, n_atoms), dtype=np.float64) for step in trange(record_steps, desc="[compute] 计算中"): t = (step + (warmup_steps or 0)) * DT - # 先施加驱动力（受驱原子的位置覆盖初始/固定约束，为弹簧力提供正确参考） x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, t, step, DRIVER_DATA, DT) - traj_x[step] = x - traj_y[step] = y - traj_z[step] = z - traj_vx[step] = vx - traj_vy[step] = vy - traj_vz[step] = vz + if save_trajectory: + traj_x[step] = x + traj_y[step] = y + traj_z[step] = z + traj_vx[step] = vx + traj_vy[step] = vy + traj_vz[step] = vz + + # 抽帧：NSTEP 间隔保存 + if step % NSTEP == 0: + fi = step // NSTEP + sampled_x[fi] = x + sampled_y[fi] = y + sampled_z[fi] = z + 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) x, y, z = wrap_position(x, y, z) x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) - return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz + # 写入 display.txt（新格式） + output_dir = get_output_dir() + disp_path = os.path.join(output_dir, "display.txt") + n_frames_actual = len(frame_indices) + 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), record_steps, n_atoms, + header_fields={"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD), + "warmup_steps": str(warmup_steps or 0), + "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)} + ) + print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)") + + # 可选：保存完整 trajectory.txt + if save_trajectory: + save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz) + print(f"[compute] trajectory.txt 已保存（完整轨迹）") + + return 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] def save_trajectory_table_txt(txt_path, traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT): @@ -1313,10 +1512,13 @@ def main(): output_dir = get_output_dir(script_dir) print(f"[compute] 开始计算 NT={NT} DT={DT} ") - traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation() + traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation(save_trajectory=0) print(f"[compute] 计算完成，共 {NT} 步") + print(f"[compute] display.txt 已在 run_simulation 中保存") - save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, script_dir) + # 如果需要完整轨迹，以上传 save_trajectory=1 重新运行 + # 以下旧函数保留兼容但不再自动调用 + # save_trajectory_txt(...) # 同时保存为逐行表格，便于直接查看 txt_path = os.path.join(output_dir, "trajectory_table.txt") diff --git a/draw.py b/draw.py index 82a16fd..65cd2e8 100644 --- a/draw.py +++ b/draw.py @@ -1,9 +1,8 @@ """VisPy 演示：加载预计算轨迹数据，驱动小球运动动画。 计算与显示完全分离： - 1. 先运行 compute.py → 生成 output/trajectory.txt（全量 NT 步轨迹） - 2. 再运行 sample.py → 从 output/trajectory.txt 抽帧生成 output/display.txt - 3. 本文件加载 output/display.txt，按帧播放动画 + 1. 运行 run_dynamics.py → 生成 output/display.txt（新格式，直接抽帧） + 2. 本文件加载 output/display.txt，按帧播放动画 用法： python draw.py # 使用 dynamics 根目录下的 output/ @@ -39,66 +38,67 @@ if not os.path.exists(disp_path): f"用法: python draw.py [案例输出目录]" ) -disp_data = compute.load_text_data(disp_path) +disp_data = compute.load_display_txt(disp_path) +h = disp_data["header_fields"] -# 单原子数据（plot_atom：用于信息显示） -DISP_X = disp_data["disp_x"] -DISP_Y = disp_data["disp_y"] -DISP_Z = disp_data["disp_z"] -DISP_VX = disp_data["disp_vx"] -DISP_VY = disp_data["disp_vy"] -DISP_VZ = disp_data["disp_vz"] +# 全原子帧数据 +DISP_ALL_X = disp_data["frames_x"] # (n_frames, n_atoms) +DISP_ALL_Y = disp_data["frames_y"] +DISP_ALL_Z = disp_data["frames_z"] +DISP_ALL_VX = disp_data["frames_vx"] +DISP_ALL_VY = disp_data["frames_vy"] +DISP_ALL_VZ = disp_data["frames_vz"] -# 全原子数据（用于多球绘制） -DISP_ALL_X = disp_data["disp_all_x"] # (n_frames, n_atoms) -DISP_ALL_Y = disp_data["disp_all_y"] -DISP_ALL_Z = disp_data["disp_all_z"] -DISP_ALL_VX = disp_data["disp_all_vx"] -DISP_ALL_VY = disp_data["disp_all_vy"] -DISP_ALL_VZ = disp_data["disp_all_vz"] +# 第一个原子的轨迹（用于信息显示） +DISP_X = DISP_ALL_X[:, 0] +DISP_Y = DISP_ALL_Y[:, 0] +DISP_Z = DISP_ALL_Z[:, 0] +DISP_VX = DISP_ALL_VX[:, 0] +DISP_VY = DISP_ALL_VY[:, 0] +DISP_VZ = DISP_ALL_VZ[:, 0] -DISP_T = disp_data["disp_t"] -DISP_STEP = disp_data["disp_step"] -N_FRAMES = int(disp_data["n_frames"]) -NT = int(disp_data["NT"]) -DT = float(disp_data["DT"]) -NSTEP = int(disp_data["NSTEP"]) +N_FRAMES = DISP_ALL_X.shape[0] +NT = int(disp_data["n_total_frames"]) +N_ATOMS = int(disp_data["n_total_particles"]) +DT = float(h.get("DT", 0.001)) +NSTEP = int(h.get("NSTEP", 1)) +DISP_STEP = np.arange(N_FRAMES) * NSTEP +DISP_T = DISP_STEP * DT # 原子信息 -ATOM_IDS = disp_data.get("atom_ids", np.array([1])) -ATOM_RADII = disp_data.get("atom_radii", np.array([float(disp_data["ball_radius"])])) -N_ATOMS = len(ATOM_IDS) -PLOT_ATOM_ROW = int(disp_data.get("plot_atom_row", 0)) -PLOT_ATOM_ID = int(disp_data.get("plot_atom_id", ATOM_IDS[0])) -BOND_PAIRS = disp_data.get("bond_pairs", []) +ATOM_IDS = disp_data["atom_ids"] +ATOM_RADII = np.full(N_ATOMS, float(h.get("ball_radius", 0.5))) +PLOT_ATOM_ROW = 0 +PLOT_ATOM_ID = int(ATOM_IDS[0]) +BOND_PAIRS = [] # display 格式不含成键信息，从原始数据加载 # 渲染方式：0=Sphere(网格球体), 1=Marker(GPU点精灵) -USE_MARKER = int(disp_data.get("use_marker", 0)) +USE_MARKER = 0 if N_FRAMES <= 0: raise ValueError( "output/display.txt 中没有可播放的帧，请检查 sample_start/sample_end/NSTEP 配置。") # 保留模拟边界常量（用于场景缩放、相机等），从 output/display.txt 中读取 -X_MIN = float(disp_data["X_MIN"]); X_MAX = float(disp_data["X_MAX"]) -Y_MIN = float(disp_data["Y_MIN"]); Y_MAX = float(disp_data["Y_MAX"]) -Z_MIN = float(disp_data["Z_MIN"]); Z_MAX = float(disp_data["Z_MAX"]) -X0 = float(disp_data["X0"]); Y0 = float(disp_data["Y0"]); Z0 = float(disp_data["Z0"]) -raw_alpha = disp_data["alpha"] -if isinstance(raw_alpha, (list, tuple, np.ndarray)): - alpha_list = [float(a) for a in raw_alpha] +X_MIN = float(h.get("X_MIN", -10)); X_MAX = float(h.get("X_MAX", 10)) +Y_MIN = float(h.get("Y_MIN", -10)); Y_MAX = float(h.get("Y_MAX", 10)) +Z_MIN = float(h.get("Z_MIN", -10)); Z_MAX = float(h.get("Z_MAX", 10)) +raw_alpha = h.get("alpha", "0.2") +try: + alpha_list = [float(x) for x in raw_alpha.split(",")] if len(alpha_list) != 6: - raise ValueError(f"alpha 数组长度须为 6，实际为 {len(alpha_list)}") -else: + alpha_list = alpha_list * 6 +except: alpha_list = [float(raw_alpha)] * 6 + # 绘图参数 -ball_radius = float(disp_data["ball_radius"]) -ball_color_r = float(disp_data["ball_color_r"]) -ball_color_g = float(disp_data["ball_color_g"]) -ball_color_b = float(disp_data["ball_color_b"]) -box_color_r = float(disp_data["box_color_r"]) -box_color_g = float(disp_data["box_color_g"]) -box_color_b = float(disp_data["box_color_b"]) +ball_radius = float(h.get("ball_radius", 0.5)) +ball_color_r = float(h.get("ball_color_r", 0.9)) +ball_color_g = float(h.get("ball_color_g", 0.2)) +ball_color_b = float(h.get("ball_color_b", 0.2)) +box_color_r = float(h.get("box_color_r", 0.8)) +box_color_g = float(h.get("box_color_g", 0.8)) +box_color_b = float(h.get("box_color_b", 0.85)) # =========================================================================== diff --git a/dynamics.py b/dynamics.py index 2f6777f..85e82b4 100644 --- a/dynamics.py +++ b/dynamics.py @@ -94,16 +94,6 @@ def build_sample_indices(total_steps, sample_step, sample_start, sample_end): return indices -def save_display_txt(data, out_dir=None): - """将抽帧数据保存到 output/display.txt（含所有参数元数据）。""" - if out_dir is None: - out_dir = os.path.dirname(os.path.abspath(__file__)) - disp_path = os.path.join(compute.get_output_dir(out_dir), "display.txt") - compute.save_text_data(disp_path, data) - print(f"[sample] 显示数组已保存至: {disp_path}") - return disp_path - - def run_case(config_path, runtime_base, input_dir="input", output_dir="output", no_plot=False): """Run one case with explicit program path, input path, and output path.""" runtime_base = Path(runtime_base).resolve() @@ -193,15 +183,15 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output", _t0 = _time.time() if engine == "python": - traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = compute.run_from_config(config, str(runtime_base)) - compute.save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, str(runtime_base)) + compute.run_from_config(config, str(runtime_base)) else: - # 外部引擎：先加载配置到全局变量，再运行引擎，再用 save_trajectory_txt 补全 metadata + # 外部引擎：先加载配置到全局变量，再运行引擎 config["_skip_run"] = True compute.run_from_config(config, str(runtime_base)) config.pop("_skip_run", None) input_dir_abs = str(input_dir_path.resolve()) output_dir_abs = str(output_dir_path.resolve()) + # 外部引擎写完整 trajectory.txt，后续抽帧 traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = compute.run_engine( engine, input_dir_abs, output_dir_abs, config) compute.save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, str(runtime_base)) @@ -209,123 +199,58 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output", _elapsed = _time.time() - _t0 print(f"[run] 引擎: {engine} 计算完成: {record_steps} 步 {_elapsed:.3f} s") else: - print("[run] 步骤 [模拟] 已跳过，直接加载已有轨迹") + print("[run] 步骤 [模拟] 已跳过") + + # 3. 检查 display.txt（新格式），不存在则从 trajectory.txt 抽帧生成 + disp_path_new = os.path.join(output_dir_abs, "display.txt") + if os.path.exists(disp_path_new): + print(f"[run] 发现已有 display.txt（引擎直接抽帧）") + else: + # 从 trajectory.txt 抽帧（外部引擎） + traj_path = os.path.join(output_dir_abs, "trajectory.txt") if not os.path.exists(traj_path): - print(f"[run] 错误: trajectory.txt 不存在，无法跳过模拟") + print(f"[run] 错误: 找不到 trajectory.txt 或 display.txt") sys.exit(1) - - # 3. 抽帧 → output/display.txt - traj_path = os.path.join(output_dir_abs, "trajectory.txt") - data = compute.load_text_data(traj_path) - - NT = int(data["NT"]); DT = float(data["DT"]); NSTEP = int(data["NSTEP"]) - warmup_steps = int(data.get("warmup_steps", 0)) - plot_atom_row = int(data["plot_atom_row"]) if "plot_atom_row" in data else 0 - plot_atom_id = int(data["plot_atom_id"]) if "plot_atom_id" in data else int(data["atom_ids"][plot_atom_row]) - - # 抽帧范围控制 - sample_start = read_optional_index(data, "sample_start", 0) - sample_end = read_optional_index(data, "sample_end", NT) - - indices = build_sample_indices(NT, NSTEP, sample_start, sample_end) - n_frames = len(indices) - - print(f"[run] 抽帧范围: [{sample_start}, {sample_end}), 共 {n_frames} 帧") - - traj_x = data["traj_x"] - traj_y = data["traj_y"] - traj_z = data["traj_z"] - traj_vx = data["traj_vx"] - traj_vy = data["traj_vy"] - traj_vz = data["traj_vz"] - - if traj_x.ndim == 1: - selected_x = traj_x - selected_y = traj_y - selected_z = traj_z - selected_vx = traj_vx - selected_vy = traj_vy - selected_vz = traj_vz - all_x = traj_x[:, None] - all_y = traj_y[:, None] - all_z = traj_z[:, None] - all_vx = traj_vx[:, None] - all_vy = traj_vy[:, None] - all_vz = traj_vz[:, None] - else: - selected_x = traj_x[:, plot_atom_row] - selected_y = traj_y[:, plot_atom_row] - selected_z = traj_z[:, plot_atom_row] - selected_vx = traj_vx[:, plot_atom_row] - selected_vy = traj_vy[:, plot_atom_row] - selected_vz = traj_vz[:, plot_atom_row] - all_x = traj_x - all_y = traj_y - all_z = traj_z - all_vx = traj_vx - all_vy = traj_vy - all_vz = traj_vz - - if config.get("step_sample", 1): - disp_data = { - "disp_x": selected_x[indices], - "disp_y": selected_y[indices], - "disp_z": selected_z[indices], - "disp_vx": selected_vx[indices], - "disp_vy": selected_vy[indices], - "disp_vz": selected_vz[indices], - "disp_all_x": all_x[indices], - "disp_all_y": all_y[indices], - "disp_all_z": all_z[indices], - "disp_all_vx": all_vx[indices], - "disp_all_vy": all_vy[indices], - "disp_all_vz": all_vz[indices], - "disp_t": indices * DT, - "disp_step": indices, - "n_frames": n_frames, - "NT": NT, "DT": DT, "NSTEP": NSTEP, - "plot_atom_id": plot_atom_id, - "plot_atom_row": plot_atom_row, - "method": str(data["method"]) if "method" in data else "explicit_euler", - "coord_file": str(data["coord_file"]) if "coord_file" in data else os.path.join("input", "coord.txt"), - "atom_ids": data["atom_ids"] if "atom_ids" in data else np.array([1]), - "atom_masses": data["atom_masses"] if "atom_masses" in data else np.array([float(data["M"])]), - "atom_radii": data["atom_radii"] if "atom_radii" in data else np.array([float(data["ball_radius"])]), - "atom_positions": data["atom_positions"] if "atom_positions" in data else np.array([[float(data["X0"]), float(data["Y0"]), float(data["Z0"])]]), - "atom_velocities": data["atom_velocities"] if "atom_velocities" in data else np.array([[float(data["VX0"]), float(data["VY0"]), float(data["VZ0"])]]), - "atom_fixed": data["atom_fixed"] if "atom_fixed" in data else np.array([[0, 0, 0]]), - "bond_pairs": data.get("bond_pairs", np.zeros((0, 2), dtype=np.int64)).tolist(), - "bond_stiffness": data.get("bond_stiffness", np.zeros(0, dtype=np.float64)).tolist(), - "bond_rest_lengths": data.get("bond_rest_lengths", np.zeros(0, dtype=np.float64)).tolist(), - "warmup_steps": warmup_steps, - "sample_start": sample_start, - "sample_end": sample_end, - "X_MIN": float(data["X_MIN"]), "X_MAX": float(data["X_MAX"]), - "Y_MIN": float(data["Y_MIN"]), "Y_MAX": float(data["Y_MAX"]), - "Z_MIN": float(data["Z_MIN"]), "Z_MAX": float(data["Z_MAX"]), - "X0": float(data["X0"]), "Y0": float(data["Y0"]), "Z0": float(data["Z0"]), - "VX0": float(data["VX0"]), "VY0": float(data["VY0"]), "VZ0": float(data["VZ0"]), - "M": float(data["M"]) if "M" in data else 1.0, - "alpha": data["alpha"], - "ball_radius": float(data["ball_radius"]), - "ball_color_r": float(data["ball_color_r"]), - "ball_color_g": float(data["ball_color_g"]), - "ball_color_b": float(data["ball_color_b"]), - "box_color_r": float(data["box_color_r"]), - "box_color_g": float(data["box_color_g"]), - "box_color_b": float(data["box_color_b"]), - "gravity_field": int(data.get("gravity_field", 1)), - "gravity_interaction": int(data.get("gravity_interaction", 0)), - "elastic_force": int(data.get("elastic_force", 1)), - "damping_force": int(data.get("damping_force", 0)), - "gravity_strength": float(data.get("gravity_strength", 1.0)), - "driving_force": int(config.get("driving_force", 0)), - "use_marker": int(config.get("use_marker", 0)), - } - save_display_txt(disp_data, str(runtime_base)) - print(f"[run] 抽帧完成: {sample_end - sample_start} 步 -> {n_frames} 帧") - else: - print("[run] 步骤 [抽帧] 已跳过") + data = compute.load_text_data(traj_path) + NT = int(data["NT"]); DT = float(data["DT"]); NSTEP = int(data.get("NSTEP", 1)) + record_steps = NT - int(data.get("warmup_steps", 0)) + n_atoms = len(data["atom_ids"]) + sample_start = 0 + sample_end = NT + indices = np.arange(0, record_steps, NSTEP, dtype=np.int64) + if len(indices) == 0: + indices = np.array([0]) + + traj_x = data["traj_x"]; traj_y = data["traj_y"]; traj_z = data["traj_z"] + traj_vx = data["traj_vx"]; traj_vy = data["traj_vy"]; traj_vz = data["traj_vz"] + + # 构建 header_fields + hf = {"DT": str(DT), "NSTEP": str(NSTEP), "method": str(data.get("method", "")), + "warmup_steps": str(data.get("warmup_steps", 0)), + "X_MAX": str(data.get("X_MAX", 10)), "X_MIN": str(data.get("X_MIN", -10)), + "Y_MAX": str(data.get("Y_MAX", 10)), "Y_MIN": str(data.get("Y_MIN", -10)), + "Z_MAX": str(data.get("Z_MAX", 10)), "Z_MIN": str(data.get("Z_MIN", -10)), + "ball_radius": str(data.get("ball_radius", 0.5)), + "ball_color_r": str(data.get("ball_color_r", 0.9)), + "ball_color_g": str(data.get("ball_color_g", 0.2)), + "ball_color_b": str(data.get("ball_color_b", 0.2)), + "box_color_r": str(data.get("box_color_r", 0.8)), + "box_color_g": str(data.get("box_color_g", 0.8)), + "box_color_b": str(data.get("box_color_b", 0.85)), + "gravity_field": str(data.get("gravity_field", 1)), + "gravity_interaction": str(data.get("gravity_interaction", 0)), + "elastic_force": str(data.get("elastic_force", 1)), + "damping_force": str(data.get("damping_force", 0)), + "gravity_strength": str(data.get("gravity_strength", 1.0)), + "driving_force": str(data.get("driving_force", 0))} + + compute.save_display_txt( + disp_path_new, + traj_x[indices], traj_y[indices], traj_z[indices], + traj_vx[indices], traj_vy[indices], traj_vz[indices], + np.array(data["atom_ids"]), record_steps, n_atoms, + header_fields=hf) + print(f"[run] 从 trajectory.txt 抽帧生成 display.txt ({len(indices)} 帧)") # 4. 绘图（可选） if not no_plot and config.get("step_plot", 1): diff --git a/examples/case06/input/input.txt b/examples/case06/input/input.txt index 43bd61c..ce4d38e 100644 --- a/examples/case06/input/input.txt +++ b/examples/case06/input/input.txt @@ -5,8 +5,8 @@ # ── 流程控制 ────────────────────────────────── # 每步用 0/1 单独开关，1=执行，0=跳过 # 依赖关系：抽帧依赖模拟结果，绘图依赖模拟+抽帧 -step_simulate: 1 # 运行物理模拟 → output/trajectory.txt -step_sample: 1 # 抽帧 → output/display.txt +step_simulate: 1 # 运行物理模拟 → output/display.txt（引擎直接抽帧） +step_sample: 0 # （旧版）从 trajectory.txt 重新抽帧，默认0=不执行 step_plot: 0 # 绘制轨迹/能量图 → output/trajectory_plots.png step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） step_plot_wave: 0 # 绘制波形能量动画 @@ -14,9 +14,12 @@ force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用 plot_wave_save_gif: 1 # 输出波形 GIF（需 step_plot_wave=1） plot_wave_save_mp4: 1 # 输出波形 MP4（需 step_plot_wave=1） +# ── 文件保存 ────────────────────────────────── +save_trajectory: 0 # 0=不保留完整轨迹文件, 1=保留 trajectory.txt（用于后续单独抽帧） + # ── 计算引擎 ────────────────────────────────── # 可选: python, c, cpp, fortran, java -engine: cpp # 默认使用 python 引擎 +engine: python # 默认使用 python 引擎 # ── 盒子 ────────────────────────────────────── box_a: 80.0 # 立方体半边长，粒子被限制在 [-box_a, box_a]³ 内