refactor: 引擎直接抽帧 + 新 display.txt 纯文本格式 + save_trajectory 开关
核心变更： 1. compute.py: run_simulation 直接按 NSTEP 抽帧写 display.txt（新格式） - 新格式：纯文本，帧 1→n 分块，每行: n x y z vx vy vz - 新函数 save_display_txt() / load_display_txt() - save_trajectory 参数（默认0=不保留 trajectory.txt） 2. dynamics.py: 移除旧 JSON 采样逻辑，自动检测 display.txt - Python 引擎直接读取引擎输出的 display.txt - 外部引擎仍写 trajectory.txt，自动抽帧转 display.txt 3. draw.py: 适配 load_display_txt() 新格式 4. case06/input.txt: 添加 save_trajectory: 0, step_sample: 0 TODO: 外部引擎（C/C++/Fortran）内部抽帧写 display.txt TODO: plot_wave.py 适配新格式 TODO: 其他案例 input.txt 更新默认值
This commit is contained in:
+53
-128
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user