feat: 真蛙跳法重构(Python/C/C++/Fortran 四引擎统一)

- 新增 compute_accel_conservative / accel_conservative:
  保守力加速度(弹簧+重力+原子间引力),不含阻尼,供蛙跳专用
- 重写 leapfrog_step / leapfrog_full:
  - 无阻尼:纯辛积分器,每步 1 次力计算(原 Velocity-Verlet 需 2 次)
  - 有阻尼:半隐式处理 v(t+dt/2)=[v(t-dt/2)*(1-α)+a_c*dt]/(1+α),无条件稳定
- 主循环加初始化反向半步 v(-dt/2)=v(0)-0.5*a_c(0)*dt
- 修复 C/C++ number of frames 字段写采样帧数而非总积分步数的 bug
- Python 引擎:新增 display.npz 二进制格式,draw.py/plot_wave.py 优先读取
- 编译参数统一为 -O3 -march=native -ffast-math
This commit is contained in:
2026-06-12 18:36:37 +08:00
parent e1ade53fff
commit e62e536cee
12 changed files with 627 additions and 355 deletions
+3
View File
@@ -100,3 +100,6 @@ desktop.ini
*.tar.gz *.tar.gz
*.tar.bz2 *.tar.bz2
output_test/ output_test/
optimization
optimization/*
+266 -42
View File
@@ -66,6 +66,12 @@ box_color_g = None
box_color_b = None box_color_b = None
use_marker = 0 use_marker = 0
camera_keyframes_raw = "" 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 # 均匀重力场 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): def get_output_dir(base_dir=None):
"""Return the output directory used for generated artifacts.""" """Return the output directory used for generated artifacts."""
override = os.environ.get("DYNAMICS_OUTPUT_DIR") 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 # 全程驱动 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 + d["phi"])
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 + d["phi"])
vel_drive = -d["amp"] * 2.0 * np.pi * d["freq"] * np.sin(2.0 * np.pi * d["freq"] * t_vec + d["phi"])
x[idx] = pos_drive[0] x[idx] = pos_drive[0]
y[idx] = pos_drive[1] 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 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 ball_radius, ball_color_r, ball_color_g, ball_color_b
global box_color_r, box_color_g, box_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 warmup_steps, sample_start, sample_end
global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH
global DRIVING_FORCE, DRIVER_DATA 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) coord_path = os.path.join(out_dir, coord_path)
(ATOM_IDS, ATOM_MASSES, ATOM_RADII, ATOM_POSITIONS, (ATOM_IDS, ATOM_MASSES, ATOM_RADII, ATOM_POSITIONS,
ATOM_VELOCITIES, ATOM_FIXED) = load_coord_file(coord_path) 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_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"))) BOND_PARAMETER_FILE = str(config.get("bond_file", os.path.join("input", "bond.txt")))
connection_path = BOND_CONNECTION_FILE 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_g = float(config.get("box_color_g", 0.8))
box_color_b = float(config.get("box_color_b", 0.85)) box_color_b = float(config.get("box_color_b", 0.85))
use_marker = int(config.get("use_marker", 0)) 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 global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH
@@ -921,9 +999,25 @@ def run_engine(engine, input_dir, output_dir, config):
print(f"[compute] input: {input_dir}") print(f"[compute] input: {input_dir}")
print(f"[compute] output: {output_dir}") print(f"[compute] output: {output_dir}")
# ── 预校准:跑 NT=1000 步测速 ──────────────────────────────── # ── 预校准:跑少量步测速(结果缓存,避免每次重跑)────────────
total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0)) total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0))
_calib_nt = min(1000, max(100, total_steps // 10)) _calib_nt = min(1000, max(100, total_steps // 10))
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 = dict(param_json)
_calib_param["NT"] = _calib_nt _calib_param["NT"] = _calib_nt
_calib_path = os.path.join(script_dir, "engines", engine, "_calib.json") _calib_path = os.path.join(script_dir, "engines", engine, "_calib.json")
@@ -945,7 +1039,15 @@ def run_engine(engine, input_dir, output_dir, config):
_calib_elapsed = max(time.time() - _ct0, 0.001) _calib_elapsed = max(time.time() - _ct0, 0.001)
_overhead = _calib_elapsed * 0.15 _overhead = _calib_elapsed * 0.15
_step_time = max(_calib_elapsed - _overhead, 0.0001) / _calib_nt _step_time = max(_calib_elapsed - _overhead, 0.0001) / _calib_nt
_est_total = max(_calib_elapsed, _overhead + _step_time * total_steps) # 保存缓存
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 = time.time()
t_start_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") 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): if not os.path.exists(disp_path):
raise FileNotFoundError(f"引擎未生成 display.txt: {disp_path}") raise FileNotFoundError(f"引擎未生成 display.txt: {disp_path}")
print(f"[compute] 引擎已生成 {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 return None, None, None, None, None, None
# save_trajectory=1:加载完整 trajectory.txt # 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 fz -= b[2] * vz
if ELASTIC_FORCE and BOND_PAIRS is not None and len(BOND_PAIRS) > 0: 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): idx_1 = BOND_PAIRS[:, 0]
idx_2 = BOND_PAIRS[:, 1]
dx = x[idx_2] - x[idx_1] dx = x[idx_2] - x[idx_1]
dy = y[idx_2] - y[idx_1] dy = y[idx_2] - y[idx_1]
dz = z[idx_2] - z[idx_1] dz = z[idx_2] - z[idx_1]
dist = np.sqrt(dx * dx + dy * dy + dz * dz) dist = np.sqrt(dx * dx + dy * dy + dz * dz)
if dist <= 1e-12: valid = dist > 1e-12
continue 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] fx_bond = force_scale * dx
force_mag = BOND_STIFFNESS[bond_idx] * stretch fy_bond = force_scale * dy
ux = dx / dist fz_bond = force_scale * dz
uy = dy / dist
uz = dz / dist
fx_bond = force_mag * ux np.add.at(fx, idx_1, fx_bond)
fy_bond = force_mag * uy np.add.at(fx, idx_2, -fx_bond)
fz_bond = force_mag * uz np.add.at(fy, idx_1, fy_bond)
np.add.at(fy, idx_2, -fy_bond)
fx[idx_1] += fx_bond np.add.at(fz, idx_1, fz_bond)
fy[idx_1] += fy_bond np.add.at(fz, idx_2, -fz_bond)
fz[idx_1] += fz_bond
fx[idx_2] -= fx_bond
fy[idx_2] -= fy_bond
fz[idx_2] -= fz_bond
if GRAVITY_INTERACTION: if GRAVITY_INTERACTION:
n = len(m) 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 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): 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) ax, ay, az = compute_acceleration(x, y, z, vx, vy, vz, m, g, b)
x = x + vx * dt 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): def apply_fixed_constraints(x, y, z, vx, vy, vz):
"""Keep fixed degrees of freedom at their initial coordinate with zero speed.""" """Keep fixed degrees of freedom at their initial coordinate with zero speed."""
fixed = ATOM_FIXED != 0 if FIXED_MASK_X is not None and np.any(FIXED_MASK_X):
positions = np.column_stack((x, y, z)) x[FIXED_MASK_X] = ATOM_POSITIONS[FIXED_MASK_X, 0]
velocities = np.column_stack((vx, vy, vz)) vx[FIXED_MASK_X] = 0.0
positions = np.where(fixed, ATOM_POSITIONS, positions) if FIXED_MASK_Y is not None and np.any(FIXED_MASK_Y):
velocities = np.where(fixed, 0.0, velocities) y[FIXED_MASK_Y] = ATOM_POSITIONS[FIXED_MASK_Y, 1]
return ( vy[FIXED_MASK_Y] = 0.0
positions[:, 0], positions[:, 1], positions[:, 2], if FIXED_MASK_Z is not None and np.any(FIXED_MASK_Z):
velocities[:, 0], velocities[:, 1], velocities[:, 2], 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): def wrap_position(x, y, z):
"""边界回绕( dynamics 模式)。""" """边界回绕( dynamics 模式)。"""
@@ -1450,11 +1610,24 @@ 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_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) 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/midpointvx/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: if warmup_steps is not None and warmup_steps > 0:
print(f"[compute] 预热阶段: 前 {warmup_steps} 步不记录") print(f"[compute] 预热阶段: 前 {warmup_steps} 步不记录")
for step in trange(warmup_steps, desc="[compute] 预热"): for step in trange(warmup_steps, desc="[compute] 预热"):
t = (step + 1) * DT 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_driving_force(x, y, z, vx, vy, vz, t, step, DRIVER_DATA, DT)
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, 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 = wrap_position(x, y, z)
x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)
@@ -1467,8 +1640,6 @@ def run_simulation(save_trajectory=0):
record_steps = NT - (warmup_steps or 0) record_steps = NT - (warmup_steps or 0)
n_atoms = len(ATOM_IDS) n_atoms = len(ATOM_IDS)
n_frames = (record_steps + NSTEP - 1) // NSTEP n_frames = (record_steps + NSTEP - 1) // NSTEP
frame_indices = []
# 按 NSTEP 抽帧的临时缓冲区(远小于全量轨迹) # 按 NSTEP 抽帧的临时缓冲区(远小于全量轨迹)
sampled_x = np.zeros((n_frames, n_atoms), dtype=np.float64) sampled_x = np.zeros((n_frames, n_atoms), dtype=np.float64)
sampled_y = np.zeros((n_frames, n_atoms), dtype=np.float64) sampled_y = np.zeros((n_frames, n_atoms), dtype=np.float64)
@@ -1507,8 +1678,10 @@ def run_simulation(save_trajectory=0):
sampled_vx[fi] = vx sampled_vx[fi] = vx
sampled_vy[fi] = vy sampled_vy[fi] = vy
sampled_vz[fi] = vz sampled_vz[fi] = vz
frame_indices.append(step) 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, 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 = wrap_position(x, y, z)
x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz) x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)
@@ -1516,12 +1689,12 @@ def run_simulation(save_trajectory=0):
# 写入 display.txt(新格式) # 写入 display.txt(新格式)
output_dir = get_output_dir() output_dir = get_output_dir()
disp_path = os.path.join(output_dir, "display.txt") 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( save_display_txt(
disp_path, disp_path,
sampled_x[:n_frames_actual], sampled_y[:n_frames_actual], sampled_z[:n_frames_actual], 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], 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), header_fields={"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD),
"warmup_steps": str(warmup_steps or 0), "warmup_steps": str(warmup_steps or 0),
"dynamic_steps": str(record_steps), "dynamic_steps": str(record_steps),
@@ -1544,14 +1717,65 @@ def run_simulation(save_trajectory=0):
"driving_force": str(DRIVING_FORCE), "driving_force": str(DRIVING_FORCE),
"use_marker": str(use_marker), "use_marker": str(use_marker),
"alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])), "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), "atom_radii": ",".join(str(r) for r in ATOM_RADII),
"camera_distance": str(config.get("camera_distance", 40.0)), "camera_distance": str(camera_distance),
"camera_elevation": str(config.get("camera_elevation", 0)), "camera_elevation": str(camera_elevation),
"camera_azimuth": str(config.get("camera_azimuth", 0)), "camera_azimuth": str(camera_azimuth),
"camera_keyframes": str(camera_keyframes_raw)} "camera_keyframes": str(camera_keyframes_raw)}
) )
print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)") print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)")
# 同时保存二进制 display.npzI/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 # 可选:保存完整 trajectory.txt
if save_trajectory: if save_trajectory:
save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz) save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz)
+12 -6
View File
@@ -30,15 +30,20 @@ else:
output_dir = compute.get_output_dir(script_dir) output_dir = compute.get_output_dir(script_dir)
os.environ["DYNAMICS_OUTPUT_DIR"] = output_dir os.environ["DYNAMICS_OUTPUT_DIR"] = output_dir
disp_path = os.path.join(output_dir, "display.txt") 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( raise FileNotFoundError(
f"找不到 display.txt\n" f"找不到 display.npz 或 display.txt\n"
f"期望路径: {disp_path}\n" f"期望路径: {output_dir}\n"
f"请先运行 compute.py 计算轨迹,再运行 sample.py 生成显示数组。\n" f"请先运行 compute.py 计算轨迹,再运行 sample.py 生成显示数组。\n"
f"用法: python draw.py [案例输出目录]" f"用法: python draw.py [案例输出目录]"
) )
# 优先读二进制 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) disp_data = compute.load_display_txt(disp_path)
h = disp_data["header_fields"] h = disp_data["header_fields"]
@@ -94,7 +99,7 @@ try:
alpha_list = [float(x) for x in raw_alpha.split(",")] alpha_list = [float(x) for x in raw_alpha.split(",")]
if len(alpha_list) != 6: if len(alpha_list) != 6:
alpha_list = alpha_list * 6 alpha_list = alpha_list * 6
except: except (ValueError, AttributeError):
alpha_list = [float(raw_alpha)] * 6 alpha_list = [float(raw_alpha)] * 6
# 绘图参数 # 绘图参数
@@ -514,8 +519,9 @@ def handle_mouse_press(event):
def _update_atom_positions(f_idx): def _update_atom_positions(f_idx):
"""更新所有原子到第 f_idx 帧的位置。""" """更新所有原子到第 f_idx 帧的位置。"""
if USE_MARKER: if USE_MARKER:
for i in range(N_ATOMS): marker_pos[:, 0] = DISP_ALL_X[f_idx]
marker_pos[i] = [DISP_ALL_X[f_idx, i], DISP_ALL_Y[f_idx, i], DISP_ALL_Z[f_idx, i]] marker_pos[:, 1] = DISP_ALL_Y[f_idx]
marker_pos[:, 2] = DISP_ALL_Z[f_idx]
balls.set_data(pos=marker_pos) balls.set_data(pos=marker_pos)
else: else:
for i in range(N_ATOMS): for i in range(N_ATOMS):
+16 -131
View File
@@ -14,6 +14,7 @@ import sys
import subprocess import subprocess
import time import time
import argparse import argparse
import json
from contextlib import contextmanager from contextlib import contextmanager
from pathlib import Path from pathlib import Path
@@ -32,6 +33,13 @@ def _fmt_alpha(v):
return str(float(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): def _load_camera_kf(config, runtime_base):
"""加载 move_camera.txt(速度段格式)→ JSON 字符串。""" """加载 move_camera.txt(速度段格式)→ JSON 字符串。"""
import re, 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)), "driving_force": str(data.get("driving_force", 0)),
"use_marker": str(config.get("use_marker", 0)), "use_marker": str(config.get("use_marker", 0)),
"alpha": _fmt_alpha(data.get("alpha", 0.2)), "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", [])), "atom_radii": _fmt_alpha(data.get("atom_radii", [])),
"camera_distance": str(config.get("camera_distance", 40.0)), "camera_distance": str(config.get("camera_distance", 40.0)),
"camera_elevation": str(config.get("camera_elevation", 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. 绘图(可选) # 4. 绘图(可选)
if not no_plot and config.get("step_plot", 1): if not no_plot and config.get("step_plot", 1):
try: print("[run] 注意: 旧版 step_plot 绘图路径依赖完整轨迹局部变量,当前已暂时跳过。")
import matplotlib.pyplot as plt print("[run] 如需波形/能量动画,请使用 step_plot_wave: 1。")
# 配置中文字体支持
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(f"[run] 完成!输出目录: {output_dir_abs}") print(f"[run] 完成!输出目录: {output_dir_abs}")
+57 -36
View File
@@ -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 一致 */ /* 边界条件:clamp 位置 + 速度反转 ——与 Python Limit_in_box 一致 */
static void limit_in_box(double *pos, double *vel, double lo, double hi) { static void limit_in_box(double *pos, double *vel, double lo, double hi) {
if (*pos > hi) { *pos = hi; *vel = -*vel; } if (*pos > hi) { *pos = hi; *vel = -*vel; }
@@ -650,6 +664,15 @@ static void midpoint_step(
} }
/* ── 蛙跳法(Velocity-Verlet)── */ /* ── 蛙跳法(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( static void leapfrog_step(
int n, double *x, double *y, double *z, int n, double *x, double *y, double *z,
double *vx, double *vy, double *vz, double *vx, double *vy, double *vz,
@@ -659,47 +682,30 @@ static void leapfrog_step(
double *ax = (double*)alloca(n * sizeof(double)); double *ax = (double*)alloca(n * sizeof(double));
double *ay = (double*)alloca(n * sizeof(double)); double *ay = (double*)alloca(n * sizeof(double));
double *az = (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++) { for (int i = 0; i < n; i++) {
if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue;
vx[i] += ax[i] * dt * 0.5; if (has_damping) {
vy[i] += ay[i] * dt * 0.5; double alphax = B[0] * dt / (2.0 * m[i]);
vz[i] += az[i] * dt * 0.5; 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;
/* 全推位置(不含边界)*/
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 */
y[i] += vy[i] * dt; y[i] += vy[i] * dt;
z[i] += vz[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 一致)──────────────── */ /* ── 驱动力(与 Python apply_driving_force 一致)──────────────── */
@@ -895,12 +901,13 @@ static void write_display_txt(const char *path, const Trajectory *traj,
FILE *f = fopen(path, "w"); FILE *f = fopen(path, "w");
if (!f) die("无法写入 display.txt"); if (!f) die("无法写入 display.txt");
int n_frames = traj->n_steps; int n_frames = traj->n_steps; /* 实际采样帧数,用于下面的帧循环 */
int n_particles = traj->n_atoms; int n_particles = traj->n_atoms;
int dynamic_steps = params->NT - params->warmup_steps; int dynamic_steps = params->NT - params->warmup_steps;
double T_total = dynamic_steps * params->DT; 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, "number of particles: %d\n", n_particles);
fprintf(f, "DT: %.16g\n", params->DT); fprintf(f, "DT: %.16g\n", params->DT);
fprintf(f, "NSTEP: %d\n", params->NSTEP); fprintf(f, "NSTEP: %d\n", params->NSTEP);
@@ -1015,6 +1022,20 @@ int main(int argc, char **argv) {
traj.vz = traj.vy + sampled_steps * n; 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 一致)*/ /* 初始时刻 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); if (params.driving_force) apply_driving_force(n, x, y, z, vx, vy, vz, 0.0, 0, params.DT, &drivers);
+65 -41
View File
@@ -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<double> 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 一致 */ /* 边界条件:clamp 位置 + 速度反转 ——与 Python Limit_in_box 一致 */
static void limit_in_box(double &pos, double &vel, double lo, double hi) { static void limit_in_box(double &pos, double &vel, double lo, double hi) {
if (pos > hi) { pos = hi; vel = -vel; } if (pos > hi) { pos = hi; vel = -vel; }
@@ -543,6 +563,14 @@ static void midpoint_step(
} }
/* ── 蛙跳法(Velocity-Verlet)——与 Python Leapfrog_Method 一致 ── */ /* ── 蛙跳法(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( static void leapfrog_full_step(
int n, double *x, double *y, double *z, int n, double *x, double *y, double *z,
double *vx, double *vy, double *vz, double *vx, double *vy, double *vz,
@@ -552,54 +580,33 @@ static void leapfrog_full_step(
int elastic_force, int damping_force, int elastic_force, int damping_force,
double gravity_strength) double gravity_strength)
{ {
// 第一次加速度
std::vector<double> ax(n), ay(n), az(n); std::vector<double> ax(n), ay(n), az(n);
compute_acceleration(n, x, y, z, vx, vy, vz, m, G, B, bonds,
// 1 次保守力计算(不含阻尼)
compute_accel_conservative(n, x, y, z, m, G, bonds,
gravity_field, gravity_interaction, gravity_field, gravity_interaction,
elastic_force, damping_force, gravity_strength, elastic_force, gravity_strength,
ax.data(), ay.data(), az.data()); ax.data(), ay.data(), az.data());
// 半推速度:v_half = v + 0.5*a*dt (存入 vx, vy, vz) bool has_damping = damping_force && (B[0] != 0.0 || B[1] != 0.0 || B[2] != 0.0);
for (int i = 0; i < n; i++) { for (int i = 0; i < n; i++) {
if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue; if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue;
vx[i] += ax[i] * dt * 0.5; if (has_damping) {
vy[i] += ay[i] * dt * 0.5; double alphax = B[0] * dt / (2.0 * m[i]);
vz[i] += az[i] * dt * 0.5; 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;
// 全推位置(不含边界,边界在外层统一处理)
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
y[i] += vy[i] * dt; y[i] += vy[i] * dt;
z[i] += vz[i] * dt; z[i] += vz[i] * dt;
} }
// 显式预测器:v_pred = v_half + 0.5*a_old*dt,用第一次加速度外推半步
// 包含所有力的贡献(标准 Velocity-Verlet 预测步)
std::vector<double> 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 一致)── */ /* ── 分发器:调用对应积分方法 + 边界条件(与 Python apply_motion_update 一致)── */
@@ -737,15 +744,17 @@ static void write_display_txt(
if (!f) die("无法写入 " + path); if (!f) die("无法写入 " + path);
std::cout << "[Cpp-engine] 正在写入显示数据…" << std::endl; 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 << "number of particles: " << n_atoms << "\n";
f << "DT: " << params.DT << "\n"; f << "DT: " << params.DT << "\n";
f << "NSTEP: " << params.NSTEP << "\n"; f << "NSTEP: " << params.NSTEP << "\n";
f << "method: " << params.method << "\n"; f << "method: " << params.method << "\n";
f << "warmup_steps: " << params.warmup_steps << "\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 << "T_total: " << T_total << "\n";
f << "box_a: " << params.box_a << "\n"; f << "box_a: " << params.box_a << "\n";
f << "driving_force: " << params.driving_force << "\n"; f << "driving_force: " << params.driving_force << "\n";
@@ -925,6 +934,21 @@ int main(int argc, char **argv) {
std::vector<double> traj_vy(buf_steps * n); std::vector<double> traj_vy(buf_steps * n);
std::vector<double> traj_vz(buf_steps * n); std::vector<double> traj_vz(buf_steps * n);
// 真蛙跳初始化:v(0) 反推 v(-dt/2) = v(0) - 0.5*a_c(0)*dt
if (params.method == "leapfrog") {
std::vector<double> 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 一致) // 初始时刻 t=0 驱动力(与 Python run_simulation 一致)
if (params.driving_force) if (params.driving_force)
+63 -38
View File
@@ -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_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)) 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 ! t=0 Python run_simulation
if (driving_force /= 0 .and. n_drivers > 0) then 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 if
end subroutine accel 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 + ! clamp +
subroutine limit_in_box(pos, vel, lo, hi) subroutine limit_in_box(pos, vel, lo, hi)
double precision, intent(inout) :: pos, vel 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 do
end subroutine midpoint_step 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, & subroutine leapfrog_full(n, x, y, z, vx, vy, vz, m, G, B, &
nb, bp, bk, br, fixed, dt, & nb, bp, bk, br, fixed, dt, &
gravity_field, gravity_interaction, & 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(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, 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 :: ax(n), ay(n), az(n)
double precision :: dmp_vx(n), dmp_vy(n), dmp_vz(n) double precision :: alphax, alphay, alphaz
double precision :: gx, gy, gz logical :: has_damping
integer :: i integer :: i
call accel(n, x, y, z, vx, vy, vz, m, G, B, nb, bp, bk, br, & ! 1
call accel_conservative(n, x, y, z, m, G, nb, bp, bk, br, &
gravity_field, gravity_interaction, & gravity_field, gravity_interaction, &
elastic_force, damping_force, gravity_strength, ax, ay, az) 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 do i = 1, n
if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle 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 x(i) = x(i) + vx(i) * dt
y(i) = y(i) + vy(i) * dt y(i) = y(i) + vy(i) * dt
z(i) = z(i) + vz(i) * dt z(i) = z(i) + vz(i) * dt
end do 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*dtvx 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 end subroutine leapfrog_full
! + + ! + +
+1 -1
View File
@@ -1,2 +1,2 @@
bond_name k rest_length bond_name k rest_length
k1 10.0 1.0 k1 50.0 1.0
+1 -1
View File
@@ -1,2 +1,2 @@
n amp_x amp_y amp_z freq_x freq_y freq_z phi_x phi_y phi_z period 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
+3 -3
View File
@@ -19,7 +19,7 @@ save_trajectory: 0 # 0=不保留完整轨迹文件, 1=保留 trajectory.txt
# ── 计算引擎 ────────────────────────────────── # ── 计算引擎 ──────────────────────────────────
# 可选: python, c, cpp, fortran, java # 可选: python, c, cpp, fortran, java
engine: c # 默认使用 python 引擎 engine: python # 默认使用 python 引擎
# ── 盒子 ────────────────────────────────────── # ── 盒子 ──────────────────────────────────────
box_a: 80.0 # 立方体半边长,粒子被限制在 [-box_a, box_a]³ 内 box_a: 80.0 # 立方体半边长,粒子被限制在 [-box_a, box_a]³ 内
@@ -69,10 +69,10 @@ warmup_steps: 0 # 默认 0(立即开始记录)
T_total: 100.0 T_total: 100.0
# 抽帧间隔(每 NSTEP 步取一帧用于动画) # 抽帧间隔(每 NSTEP 步取一帧用于动画)
NSTEP: 100 NSTEP: 10
# ── 时间步长 ────────────────────────────────── # ── 时间步长 ──────────────────────────────────
DT: 0.001 # 时间步长 (s) DT: 0.01 # 时间步长 (s)
# 抽帧范围:只保存 [sample_start, sample_end) 区间内的帧 # 抽帧范围:只保存 [sample_start, sample_end) 区间内的帧
sample_start: null # null 表示从头开始(帧索引从 0 起) sample_start: null # null 表示从头开始(帧索引从 0 起)
+1 -1
View File
@@ -31,7 +31,7 @@ def load_dynamics_module(module_path: Path):
def main(): def main():
parser = argparse.ArgumentParser(description="运行 Dynamics 示例案例 case01") parser = argparse.ArgumentParser(description="运行 Dynamics 示例案例 case06")
parser.add_argument("--no-plot", action="store_true", help="跳过 matplotlib 绘图") parser.add_argument("--no-plot", action="store_true", help="跳过 matplotlib 绘图")
args = parser.parse_args() args = parser.parse_args()
+102 -18
View File
@@ -14,17 +14,101 @@ import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation from matplotlib.animation import FuncAnimation
import os import os
import sys import sys
import json
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import compute import compute
def load_disp_data(output_dir): def load_disp_data(output_dir):
"""加载 display.txt""" """加载 display.npz(优先)或 display.txt"""
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): txt_path = os.path.join(output_dir, "display.txt")
raise FileNotFoundError(f"找不到 {disp_path}") if os.path.exists(npz_path):
return compute.load_text_data(disp_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, def compute_energy(x, y, z, vx, vy, vz, masses, mass_arr,
@@ -91,29 +175,29 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False):
save_gif: 是否保存 GIF save_gif: 是否保存 GIF
save_mp4: 是否保存 MP4 save_mp4: 是否保存 MP4
""" """
data = load_disp_data(output_dir) data = _load_wave_dataset(output_dir)
n_frames = int(data["n_frames"]) n_frames = int(data["n_frames"])
t = np.array(data["disp_t"]) t = np.array(data["t"])
# 位置 / 速度 # 位置 / 速度
x = np.array(data["disp_all_x"]) x = np.array(data["x"])
y = np.array(data["disp_all_y"]) y = np.array(data["y"])
z = np.array(data["disp_all_z"]) z = np.array(data["z"])
vx = np.array(data["disp_all_vx"]) vx = np.array(data["vx"])
vy = np.array(data["disp_all_vy"]) vy = np.array(data["vy"])
vz = np.array(data["disp_all_vz"]) vz = np.array(data["vz"])
# 原子信息 # 原子信息
pos_0 = np.array(data["atom_positions"]) # (n_atoms, 3) pos_0 = np.array(data["pos_0"])
masses = np.array(data["atom_masses"]) masses = np.array(data["masses"])
atom_ids = np.array(data["atom_ids"]) atom_ids = np.array(data["atom_ids"])
n_atoms = len(atom_ids) n_atoms = len(atom_ids)
# 成键 # 成键
bond_pairs = data.get("bond_pairs", []) bond_pairs = np.array(data.get("bond_pairs", []), dtype=np.int64)
bond_stiffness = np.array(data.get("bond_stiffness", [])) bond_stiffness = np.array(data.get("bond_stiffness", []), dtype=np.float64)
bond_rest_lengths = np.array(data.get("bond_rest_lengths", [])) 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))