Files
dynamics/draw.py
T
admin 854f00ae44 feat: 增加驱动力系统、Marker渲染模式、动画防闪退、案例文档
- 新增 driving_force 驱动力系统(driver.txt 定义,支持周期控制)
- 新增 use_marker 渲染开关(GPU实例化点精灵,提升大量原子性能)
- 修复动画闪退:独立控制台、错误日志、启动存活检测
- 重绘 draw.py 架构:双渲染模式 + 预分配键线缓冲区
- 修复 raw trajectory 采样时间变量遮蔽 bug
- 重构 case05: 60原子一维链 + 驱动力 + 完整案例文档
- 修复所有案例 Readme.md 编码(GBK → UTF-8)
- 所有 input.txt 新增 driver_file / driving_force / use_marker 参数
2026-06-10 15:34:53 +08:00

520 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""VisPy 演示:加载预计算轨迹数据,驱动小球运动动画。
计算与显示完全分离:
1. 先运行 compute.py → 生成 output/trajectory.txt(全量 NT 步轨迹)
2. 再运行 sample.py → 从 output/trajectory.txt 抽帧生成 output/display.txt
3. 本文件加载 output/display.txt,按帧播放动画
用法:
python draw.py # 使用 dynamics 根目录下的 output/
python draw.py examples/case01/output # 指定案例输出目录
"""
import numpy as np
import os
import sys
from vispy import app, scene
from vispy.visuals.transforms import STTransform
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import compute
# ===========================================================================
# 加载预计算轨迹
# ===========================================================================
script_dir = os.path.dirname(os.path.abspath(__file__))
if len(sys.argv) > 1:
# 用户指定了输出目录
output_dir = os.path.abspath(sys.argv[1])
else:
output_dir = compute.get_output_dir(script_dir)
os.environ["DYNAMICS_OUTPUT_DIR"] = output_dir
disp_path = os.path.join(output_dir, "display.txt")
if not os.path.exists(disp_path):
raise FileNotFoundError(
f"找不到 display.txt\n"
f"期望路径: {disp_path}\n"
f"请先运行 compute.py 计算轨迹,再运行 sample.py 生成显示数组。\n"
f"用法: python draw.py [案例输出目录]"
)
disp_data = compute.load_text_data(disp_path)
# 单原子数据(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["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_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"])
# 原子信息
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", [])
# 渲染方式:0=Sphere(网格球体), 1=Marker(GPU点精灵)
USE_MARKER = int(disp_data.get("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]
if len(alpha_list) != 6:
raise ValueError(f"alpha 数组长度须为 6,实际为 {len(alpha_list)}")
else:
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"])
# ===========================================================================
# 图形界面无关的几何参数(不参与物理计算,仅用于场景外观)
# ===========================================================================
info_margin = 36
axis_length = 10.0
initial_camera = {
"distance": 40.0,
"elevation": 0,
"azimuth": 0,
"center": (0, 0, 0),
}
# ===========================================================================
# 创建画布与相机
# ===========================================================================
canvas = scene.SceneCanvas(
keys="interactive",
size=(1000, 700),
bgcolor=(0.08, 0.08, 0.10, 1.0),
show=True,
)
view = canvas.central_widget.add_view()
view.camera = "turntable"
view.camera.distance = initial_camera["distance"]
view.camera.elevation = initial_camera["elevation"]
view.camera.azimuth = initial_camera["azimuth"]
view.camera.center = initial_camera["center"]
# ===========================================================================
# 场景对象
# ===========================================================================
axis_width = 3
arrow_size = 14
axes_visible = True
axes_group = []
axes_group.append(scene.visuals.Arrow(
pos=np.array([[0, 0, 0], [axis_length, 0, 0]], dtype=np.float32),
color=(1.0, 0.2, 0.2, 1.0),
width=axis_width,
arrows=np.array([[0, 0, 0, axis_length, 0, 0]], dtype=np.float32),
arrow_size=arrow_size,
parent=view.scene,
))
axes_group.append(scene.visuals.Arrow(
pos=np.array([[0, 0, 0], [0, axis_length, 0]], dtype=np.float32),
color=(0.2, 1.0, 0.2, 1.0),
width=axis_width,
arrows=np.array([[0, 0, 0, 0, axis_length, 0]], dtype=np.float32),
arrow_size=arrow_size,
parent=view.scene,
))
axes_group.append(scene.visuals.Arrow(
pos=np.array([[0, 0, 0], [0, 0, axis_length]], dtype=np.float32),
color=(0.3, 0.6, 1.0, 1.0),
width=axis_width,
arrows=np.array([[0, 0, 0, 0, 0, axis_length]], dtype=np.float32),
arrow_size=arrow_size,
parent=view.scene,
))
axes_group.append(scene.visuals.Text(text="x", color=(1.0, 0.2, 0.2, 1.0), font_size=18,
pos=(axis_length + 0.2, 0, 0), anchor_x="left", anchor_y="center", parent=view.scene))
axes_group.append(scene.visuals.Text(text="y", color=(0.2, 1.0, 0.2, 1.0), font_size=18,
pos=(0, axis_length + 0.2, 0), anchor_x="left", anchor_y="bottom", parent=view.scene))
axes_group.append(scene.visuals.Text(text="z", color=(0.3, 0.6, 1.0, 1.0), font_size=18,
pos=(0, 0, axis_length + 0.2), anchor_x="left", anchor_y="bottom", parent=view.scene))
# ── 原子渲染 ──────────────────────────────────
# 两种渲染模式:
# Sphere = 独立网格球体(精细,适合少量原子)
# Marker = GPU 实例化点精灵(高效,适合大量原子)
TAB10_RGB = np.array([
[0.1216, 0.4667, 0.7059], # 蓝
[0.8902, 0.4667, 0.1137], # 橙
[0.1725, 0.6275, 0.1725], # 绿
[0.8392, 0.1529, 0.1569], # 红
[0.5804, 0.4039, 0.7412], # 紫
[0.5490, 0.3373, 0.2941], # 棕
[0.8902, 0.4667, 0.7608], # 粉
[0.4980, 0.4980, 0.4980], # 灰
[0.7373, 0.7412, 0.1333], # 黄绿
[0.0902, 0.7451, 0.8118], # 青
])
# 每个原子的颜色(循环使用 tab10 色板)
atom_colors = np.zeros((N_ATOMS, 4), dtype=np.float32)
for i in range(N_ATOMS):
r, g, b = TAB10_RGB[i % len(TAB10_RGB)]
atom_colors[i] = [r, g, b, 1.0]
if USE_MARKER:
# ── Marker 模式:GPU 实例化,一次 draw call ──
marker_pos = np.zeros((N_ATOMS, 3), dtype=np.float32)
marker_sizes = np.array([float(ATOM_RADII[i]) * 40 for i in range(N_ATOMS)], dtype=np.float32)
balls = scene.visuals.Markers(parent=view.scene)
balls.set_data(
pos=marker_pos, face_color=atom_colors,
size=marker_sizes, symbol='disc', edge_width=0,
)
else:
# ── Sphere 模式:每个原子一个独立网格球体 ──
balls = []
for i in range(N_ATOMS):
r, g, b, _ = atom_colors[i]
s = scene.visuals.Sphere(
radius=float(ATOM_RADII[i]), method="latitude",
color=(r, g, b, 1.0), edge_color=None, parent=view.scene)
s.mesh.shading = "smooth"
balls.append(s)
# 成键线(原子之间的连接)
if len(BOND_PAIRS) > 0:
n_bonds = len(BOND_PAIRS)
bond_pos_buffer = np.zeros((n_bonds * 2, 3), dtype=np.float32) # 预分配,后续原地修改
bond_lines = scene.visuals.Line(
pos=bond_pos_buffer, color=(0.7, 0.7, 0.7, 0.8), width=3,
connect="segments", parent=view.scene)
else:
bond_pos_buffer = None
bond_lines = None
# 六个面形成立方体边界(每个面独立透明度,alpha<=0 时隐藏该面)
box_size = X_MAX - X_MIN
faces = [
((X_MAX, 0, 0), "-x"),
((X_MIN, 0, 0), "+x"),
((0, Y_MAX, 0), "-y"),
((0, Y_MIN, 0), "+y"),
((0, 0, Z_MAX), "-z"),
((0, 0, Z_MIN), "+z"),
]
for f_idx, (pos, direction) in enumerate(faces):
a = alpha_list[f_idx]
if a <= 0:
continue
face_color = (box_color_r, box_color_g, box_color_b, a)
plane = scene.visuals.Plane(
width=box_size, height=box_size, width_segments=1, height_segments=1,
direction=direction, color=face_color, parent=view.scene)
plane.set_gl_state(depth_test=False, blend=True)
plane.transform = STTransform(translate=pos)
# 右上角:相机信息
camera_info = scene.visuals.Text(
text="", color="white", font_size=14,
pos=(0, 0), anchor_x="right", anchor_y="top", parent=canvas.scene)
# 左上角:小球信息
ball_info = scene.visuals.Text(
text="", color=(0.2, 1.0, 0.2, 1.0), font_size=28,
pos=(0, 0), anchor_x="left", anchor_y="top",
face="黑体", bold=True, parent=canvas.scene)
# 左上角:reset 按钮(在 info 上方)
reset_btn_size = (60, 30)
reset_button = scene.visuals.Rectangle(
center=(reset_btn_size[0] / 2 + 8, reset_btn_size[1] / 2 + 8),
width=reset_btn_size[0], height=reset_btn_size[1],
radius=6, color=(0.18, 0.35, 0.65, 0.85),
border_color="white", parent=canvas.scene)
reset_button_label = scene.visuals.Text(
text="reset", color="white", font_size=16,
pos=(reset_btn_size[0] / 2 + 8, reset_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
# 信息显示/隐藏 切换按钮(左上角小方块,在 reset 下方)
info_btn_size = (60, 30)
info_toggle_visible = True
info_button = scene.visuals.Rectangle(
center=(info_btn_size[0] / 2 + 8, info_btn_size[1] / 2 + 8),
width=info_btn_size[0], height=info_btn_size[1],
radius=6, color=(0.9, 0.3, 0.3, 0.9),
border_color="white", parent=canvas.scene)
info_button_label = scene.visuals.Text(
text="info", color="white", font_size=16,
pos=(info_btn_size[0] / 2 + 8, info_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
# 强制在初始化时定位到左上角(canvas.scene坐标系:左下角为原点)
_cw, _ch = canvas.size
reset_button.center = (reset_btn_size[0] / 2 + 8, _ch - reset_btn_size[1] / 2 - 8 - info_btn_size[1] - 4)
reset_button_label.pos = reset_button.center
info_button.center = (info_btn_size[0] / 2 + 8, _ch - info_btn_size[1] / 2 - 8)
info_button_label.pos = info_button.center
# axes 显示/隐藏 按钮(在 info 下方)
axes_btn_size = (60, 30)
axes_button = scene.visuals.Rectangle(
center=(axes_btn_size[0] / 2 + 8, axes_btn_size[1] / 2 + 8),
width=axes_btn_size[0], height=axes_btn_size[1],
radius=6, color=(0.3, 0.7, 0.3, 0.9),
border_color="white", parent=canvas.scene)
axes_button_label = scene.visuals.Text(
text="axes", color="white", font_size=16,
pos=(axes_btn_size[0] / 2 + 8, axes_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
axes_button.center = (axes_btn_size[0] / 2 + 8, _ch - axes_btn_size[1] / 2 - 8 - info_btn_size[1] - 4 - axes_btn_size[1] - 4)
axes_button_label.pos = axes_button.center
# ===========================================================================
# 回调函数
# ===========================================================================
def update_camera_info(event=None):
c = view.camera
camera_info.text = (
"Camera\n"
f"center = ({c.center[0]:.2f}, {c.center[1]:.2f}, {c.center[2]:.2f})\n"
f"distance = {c.distance:.2f}\n"
f"elevation = {c.elevation:.2f}\n"
f"azimuth = {c.azimuth:.2f}"
)
def update_ball_info(frame_idx, x, y, z, vx, vy, vz):
step = int(DISP_STEP[frame_idx])
t = float(DISP_T[frame_idx])
ball_info.text = (
f"原子 {PLOT_ATOM_ID} (共 {N_ATOMS} 个原子)\n"
f"frame {frame_idx+1}/{N_FRAMES} | saved step {step}/{NT-1}\n"
f"t = {t:.2f} s | dt = {DT:.3f} s | nstep = {NSTEP}\n"
f"Position: ({x:.2f}, {y:.2f}, {z:.2f})\n"
f"Velocity: ({vx:.2f}, {vy:.2f}, {vz:.2f})\n"
)
def reposition_camera_info(event=None):
width, height = canvas.size
camera_info.pos = (width - 20, height - 20)
# reset → info → axes 三按钮纵向排列
gap = 4
info_y = info_btn_size[1] / 2 + 8
reset_y = info_y + info_btn_size[1] + gap
axes_y = info_y + info_btn_size[1] + gap + axes_btn_size[1] + gap
reset_button.center = (reset_btn_size[0] / 2 + 8, height - reset_y)
reset_button_label.pos = reset_button.center
info_button.center = (info_btn_size[0] / 2 + 8, height - info_y)
info_button_label.pos = info_button.center
axes_button.center = (axes_btn_size[0] / 2 + 8, height - axes_y)
axes_button_label.pos = axes_button.center
# info 文字放在所有按钮下方
buttons_bottom = height - (axes_y + axes_btn_size[1] / 2)
ball_info.pos = (info_margin, buttons_bottom - 10)
update_ball_info(frame_idx, DISP_X[frame_idx], DISP_Y[frame_idx], DISP_Z[frame_idx],
DISP_VX[frame_idx], DISP_VY[frame_idx], DISP_VZ[frame_idx])
update_camera_info()
def handle_view_interaction(event):
update_camera_info()
def rotate_about_screen_normal(angle):
if hasattr(view.camera, "roll"):
view.camera.roll = (view.camera.roll + angle) % 360
else:
view.camera.azimuth = (view.camera.azimuth + angle) % 360
update_camera_info()
def handle_key_press(event):
key_name = ""
if getattr(event, "text", None):
key_name = event.text.lower()
elif getattr(event, "key", None) is not None:
key_name = str(event.key).lower()
if key_name == "q":
rotate_about_screen_normal(-90)
elif key_name == "e":
rotate_about_screen_normal(90)
def reset_camera_view():
global frame_idx
frame_idx = 0
# 立即复位所有原子到第 0 帧
_update_atom_positions(frame_idx)
if bond_lines is not None and len(BOND_PAIRS) > 0:
_update_bond_positions(frame_idx)
# 复位相机
view.camera.distance = initial_camera["distance"]
view.camera.elevation = initial_camera["elevation"]
view.camera.azimuth = initial_camera["azimuth"]
view.camera.center = initial_camera["center"]
if hasattr(view.camera, "roll"):
view.camera.roll = 0
update_camera_info()
def handle_mouse_press(event):
global info_toggle_visible, axes_visible
if event.pos is None:
return
ex, ey = event.pos
# reset 按钮
bx, by = reset_button.center
lw = reset_btn_size[0] / 2; lh = reset_btn_size[1] / 2
if (bx - lw) <= ex <= (bx + lw) and (by - lh) <= ey <= (by + lh):
reset_camera_view()
return
# info 按钮
bx, by = info_button.center
if (bx - info_btn_size[0]/2) <= ex <= (bx + info_btn_size[0]/2) and (by - info_btn_size[1]/2) <= ey <= (by + info_btn_size[1]/2):
info_toggle_visible = not info_toggle_visible
ball_info.visible = info_toggle_visible
camera_info.visible = info_toggle_visible
return
# axes 按钮
bx, by = axes_button.center
if (bx - axes_btn_size[0]/2) <= ex <= (bx + axes_btn_size[0]/2) and (by - axes_btn_size[1]/2) <= ey <= (by + axes_btn_size[1]/2):
axes_visible = not axes_visible
for axis in axes_group:
axis.visible = axes_visible
# ===========================================================================
# 位置/键线更新辅助函数(两种渲染模式共用)
# ===========================================================================
def _update_atom_positions(f_idx):
"""更新所有原子到第 f_idx 帧的位置。"""
if USE_MARKER:
for i in range(N_ATOMS):
marker_pos[i] = [DISP_ALL_X[f_idx, i], DISP_ALL_Y[f_idx, i], DISP_ALL_Z[f_idx, i]]
balls.set_data(pos=marker_pos)
else:
for i in range(N_ATOMS):
balls[i].transform = STTransform(translate=(
float(DISP_ALL_X[f_idx, i]),
float(DISP_ALL_Y[f_idx, i]),
float(DISP_ALL_Z[f_idx, i]),
))
def _update_bond_positions(f_idx):
"""原地更新键线顶点位置,避免重新分配内存。"""
for b_idx, (i, j) in enumerate(BOND_PAIRS):
b2 = b_idx * 2
bond_pos_buffer[b2] = [DISP_ALL_X[f_idx, i], DISP_ALL_Y[f_idx, i], DISP_ALL_Z[f_idx, i]]
bond_pos_buffer[b2 + 1] = [DISP_ALL_X[f_idx, j], DISP_ALL_Y[f_idx, j], DISP_ALL_Z[f_idx, j]]
bond_lines.set_data(pos=bond_pos_buffer)
# ===========================================================================
# 动画初始化
# ===========================================================================
frame_idx = 0
# 初始帧:摆放所有原子并刷新 UI
_update_atom_positions(frame_idx)
if bond_lines is not None and len(BOND_PAIRS) > 0:
_update_bond_positions(frame_idx)
reposition_camera_info()
update_ball_info(frame_idx,
float(DISP_X[frame_idx]), float(DISP_Y[frame_idx]), float(DISP_Z[frame_idx]),
float(DISP_VX[frame_idx]), float(DISP_VY[frame_idx]), float(DISP_VZ[frame_idx]))
mode_str = "Marker (GPU 实例化)" if USE_MARKER else "Sphere (独立网格)"
print(f"[draw] 加载 output/display.txt: {N_FRAMES} 帧, {N_ATOMS} 个原子, NT={NT}, DT={DT}, NSTEP={NSTEP}")
print(f"[draw] 渲染方式: {mode_str}")
print(f"[draw] 绘图参数: ball_radius={ball_radius}, box_color=({box_color_r:.2f},{box_color_g:.2f},{box_color_b:.2f}), alpha={alpha_list}")
# ===========================================================================
# 每帧回调:仅推进帧索引,从预存数组读取位置,零物理计算
# ===========================================================================
def update(event):
global frame_idx
frame_idx = (frame_idx + 1) % N_FRAMES # 循环播放
# 更新所有原子位置(两种模式共用逻辑)
_update_atom_positions(frame_idx)
# 更新成键线(原地修改,无内存分配)
if bond_lines is not None and len(BOND_PAIRS) > 0:
_update_bond_positions(frame_idx)
# 信息面板显示 plot_atom 的数据
x = float(DISP_X[frame_idx])
y = float(DISP_Y[frame_idx])
z = float(DISP_Z[frame_idx])
vx = float(DISP_VX[frame_idx])
vy = float(DISP_VY[frame_idx])
vz = float(DISP_VZ[frame_idx])
update_ball_info(frame_idx, x, y, z, vx, vy, vz)
update_camera_info()
# ===========================================================================
# 事件绑定与启动
# ===========================================================================
timer = app.Timer(interval=0.02, connect=update, start=True)
canvas.events.mouse_move.connect(handle_view_interaction)
canvas.events.mouse_press.connect(handle_mouse_press)
canvas.events.mouse_wheel.connect(handle_view_interaction)
canvas.events.resize.connect(reposition_camera_info)
canvas.events.key_press.connect(handle_key_press)
if hasattr(canvas, "native") and hasattr(canvas.native, "setFocus"):
canvas.native.setFocus()
if __name__ == "__main__":
app.run()