185 lines
7.2 KiB
Python
185 lines
7.2 KiB
Python
"""
|
|
sample.py
|
|
---------
|
|
纯 Python 抽帧脚本，不依赖 VisPy。
|
|
|
|
功能：
|
|
1. 从 output/trajectory.txt 读取完整轨迹
|
|
2. 每隔 NSTEP 帧抽取一帧，生成 output/display.txt
|
|
3. output/display.txt 可直接被 draw.py 加载驱动动画
|
|
"""
|
|
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import compute
|
|
|
|
|
|
def build_sample_indices(total_steps, sample_step, sample_start, sample_end):
|
|
"""Validate sampling settings and build frame indices."""
|
|
if sample_step <= 0:
|
|
raise ValueError(f"NSTEP 必须为正整数，实际为 {sample_step}")
|
|
if sample_start < 0:
|
|
raise ValueError(f"sample_start 不能小于 0，实际为 {sample_start}")
|
|
if sample_end > total_steps:
|
|
raise ValueError(
|
|
f"sample_end 不能大于记录步数 {total_steps}，实际为 {sample_end}")
|
|
if sample_start >= sample_end:
|
|
raise ValueError(
|
|
f"sample_start 必须小于 sample_end，实际为 [{sample_start}, {sample_end})")
|
|
|
|
n_frames = (sample_end - sample_start) // sample_step
|
|
if n_frames <= 0:
|
|
raise ValueError(
|
|
f"抽帧范围 [{sample_start}, {sample_end}) 过短，按 NSTEP={sample_step} 无法抽出任何帧")
|
|
|
|
return np.arange(n_frames, dtype=np.int64) * sample_step + sample_start
|
|
|
|
|
|
def read_optional_index(data, key, default_value):
|
|
"""Read an optional integer index from txt metadata."""
|
|
if key not in data:
|
|
return default_value
|
|
value = data[key]
|
|
if value is None or int(value) < 0:
|
|
return default_value
|
|
return int(value)
|
|
|
|
|
|
def main():
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
output_dir = compute.get_output_dir(script_dir)
|
|
traj_path = os.path.join(output_dir, "trajectory.txt")
|
|
disp_path = os.path.join(output_dir, "display.txt")
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 1. 加载完整轨迹
|
|
# -------------------------------------------------------------------------
|
|
data = compute.load_text_data(traj_path)
|
|
NT = int(data["NT"])
|
|
DT = float(data["DT"])
|
|
NSTEP = int(data["NSTEP"])
|
|
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"]
|
|
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])
|
|
|
|
print(f"[sample] 加载轨迹数据: NT={NT}, DT={DT}, NSTEP={NSTEP}")
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 2. 抽帧：支持配置文件中保存的 [sample_start, sample_end) 区间
|
|
# -------------------------------------------------------------------------
|
|
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)
|
|
|
|
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
|
|
|
|
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_t = indices * DT # 物理时间 = step * DT
|
|
disp_step = indices # 对应保存轨迹的步编号（0-based）
|
|
|
|
print(f"[sample] 抽帧完成: [{sample_start}, {sample_end}) -> {n_frames} 帧")
|
|
print(f"[sample] 每帧时间跨度: {NSTEP*DT:.3f} s (即每隔 {NSTEP} 步取一帧)")
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 3. 保存显示数组
|
|
# -------------------------------------------------------------------------
|
|
payload = {
|
|
"disp_x": disp_x,
|
|
"disp_y": disp_y,
|
|
"disp_z": disp_z,
|
|
"disp_vx": disp_vx,
|
|
"disp_vy": disp_vy,
|
|
"disp_vz": disp_vz,
|
|
"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": disp_t,
|
|
"disp_step": disp_step,
|
|
"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]]),
|
|
"warmup_steps": int(data["warmup_steps"]) if "warmup_steps" in data else 0,
|
|
"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": float(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"]),
|
|
}
|
|
compute.save_text_data(disp_path, payload)
|
|
print(f"[sample] 显示数组已保存至: {disp_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|