ylniu/dynamics
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 引擎直接抽帧 + 新 display.txt 纯文本格式 + save_trajectory 开关

Browse Source 
核心变更：
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:
Ying-Li Niu
2026-06-12 06:36:50 +08:00
parent c158c74609
commit 9d1f84d2bf
4 changed files with 334 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compute.py
+216 -14
View File
@@ -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,9 +1383,22 @@ 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)
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)
@@ -1256,9 +1408,9 @@ def run_simulation():
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)
if save_trajectory:
traj_x[step] = x
traj_y[step] = y
traj_z[step] = z
@@ -1266,11 +1418,58 @@ def run_simulation():
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")
draw.py
+47 -47
View File
@@ -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))
# ===========================================================================
dynamics.py
+49 -124
View File
@@ -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] 步骤 [模拟] 已跳过，直接加载已有轨迹")
if not os.path.exists(traj_path):
print(f"[run] 错误: trajectory.txt 不存在，无法跳过模拟")
sys.exit(1)
print("[run] 步骤 [模拟] 已跳过")
# 3. 抽帧 → output/display.txt
# 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 或 display.txt")
sys.exit(1)
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])
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])
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"]
# 抽帧范围控制
sample_start = read_optional_index(data, "sample_start", 0)
sample_end = read_optional_index(data, "sample_end", NT)
# 构建 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))}
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] 步骤 [抽帧] 已跳过")
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):
examples/case06/input/input.txt
+6 -3
View File
@@ -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]³ 内