admin
757a891a43
之前所有原子使用统一的 ball_radius（来自 input.txt）， 现在 display.txt header 包含 atom_radii 字段（逗号分隔 的半径数组），draw.py 据此为每个原子设置独立半径。 fallback：若没有 atom_radii 字段，仍使用 ball_radius。
598 lines
22 KiB
Python
598 lines
22 KiB
Python
"""VisPy 演示：加载预计算轨迹数据，驱动小球运动动画。
|
||
|
||
计算与显示完全分离：
|
||
1. 运行 run_dynamics.py → 生成 output/display.txt（新格式，直接抽帧）
|
||
2. 本文件加载 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_display_txt(disp_path)
|
||
h = disp_data["header_fields"]
|
||
|
||
# 全原子帧数据
|
||
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_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]
|
||
|
||
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["atom_ids"]
|
||
# 优先使用 per-atom 半径，否则用统一的 ball_radius
|
||
_raw_radii = h.get("atom_radii", "")
|
||
if _raw_radii.strip():
|
||
ATOM_RADII = np.array([float(x) for x in _raw_radii.split(",")])
|
||
else:
|
||
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(h.get("use_marker", 0))
|
||
|
||
if N_FRAMES <= 0:
|
||
raise ValueError(
|
||
"output/display.txt 中没有可播放的帧，请检查 sample_start/sample_end/NSTEP 配置。")
|
||
|
||
# 保留模拟边界常量（用于场景缩放、相机等），从 output/display.txt 中读取
|
||
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:
|
||
alpha_list = alpha_list * 6
|
||
except:
|
||
alpha_list = [float(raw_alpha)] * 6
|
||
|
||
# 绘图参数
|
||
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))
|
||
|
||
|
||
# ===========================================================================
|
||
# 图形界面无关的几何参数（不参与物理计算，仅用于场景外观）
|
||
# ===========================================================================
|
||
|
||
info_margin = 8
|
||
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=18,
|
||
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]:.1f}, {c.center[1]:.1f}, {c.center[2]:.1f})\n"
|
||
f"distance = {c.distance:.1f}\n"
|
||
f"elevation = {c.elevation:.1f}\n"
|
||
f"azimuth = {c.azimuth:.1f}\n"
|
||
f"step = {PAN_SPEED:.1f} | {'透视' if _PERSPECTIVE else '正交'}"
|
||
)
|
||
|
||
|
||
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"
|
||
f"W/S 沿Z轴 | A/D 左右 | Q/E 升降 | C/X 步长 | V 透视/正交"
|
||
)
|
||
|
||
|
||
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()
|
||
|
||
|
||
# ── 平移速度 & 投影模式（全局变量）──
|
||
PAN_SPEED = 1.0
|
||
_PERSPECTIVE = True # True=透视, False=正交
|
||
|
||
|
||
def handle_view_interaction(event):
|
||
update_camera_info()
|
||
|
||
|
||
def handle_key_press(event):
|
||
global PAN_SPEED, _PERSPECTIVE
|
||
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()
|
||
|
||
c = view.camera
|
||
|
||
# ── 投影切换 ──
|
||
if key_name == "v":
|
||
_PERSPECTIVE = not _PERSPECTIVE
|
||
try:
|
||
if _PERSPECTIVE:
|
||
c.fov = 60.0
|
||
else:
|
||
c.fov = 0.0
|
||
print(f"[draw] 投影模式: {'透视' if _PERSPECTIVE else '正交'}")
|
||
except Exception:
|
||
print("[draw] 当前相机不支持 fov 切换")
|
||
update_camera_info()
|
||
return
|
||
|
||
# ── 步长控制 ──
|
||
if key_name == "c":
|
||
PAN_SPEED = min(PAN_SPEED * 1.5, 50.0)
|
||
print(f"[draw] 步长: {PAN_SPEED:.1f}")
|
||
update_camera_info()
|
||
return
|
||
elif key_name == "x":
|
||
PAN_SPEED = max(PAN_SPEED / 1.5, 0.05)
|
||
print(f"[draw] 步长: {PAN_SPEED:.1f}")
|
||
update_camera_info()
|
||
return
|
||
|
||
# ── 计算相机方向向量 ──
|
||
import math as _math
|
||
azim_rad = _math.radians(c.azimuth)
|
||
elev_rad = _math.radians(c.elevation)
|
||
|
||
# 视线方向（从 center 指向相机）
|
||
vd = np.array([
|
||
_math.cos(elev_rad) * _math.sin(azim_rad),
|
||
_math.sin(elev_rad),
|
||
_math.cos(elev_rad) * _math.cos(azim_rad),
|
||
])
|
||
vd /= np.linalg.norm(vd)
|
||
|
||
# 屏幕右方向
|
||
world_up = np.array([0.0, 1.0, 0.0])
|
||
right = np.cross(vd, world_up)
|
||
rn = np.linalg.norm(right)
|
||
if rn > 1e-10:
|
||
right /= rn
|
||
else:
|
||
right = np.array([1.0, 0.0, 0.0])
|
||
|
||
# 屏幕上方向
|
||
up = np.cross(right, vd)
|
||
un = np.linalg.norm(up)
|
||
if un > 1e-10:
|
||
up /= un
|
||
else:
|
||
up = np.array([0.0, 1.0, 0.0])
|
||
|
||
pan = PAN_SPEED * 0.3
|
||
|
||
if key_name == "a": # 右移（屏幕右方向）
|
||
c.center = tuple(np.array(c.center) + right * pan)
|
||
elif key_name == "d": # 左移（屏幕右方向负向）
|
||
c.center = tuple(np.array(c.center) - right * pan)
|
||
elif key_name == "e": # 上升（屏幕上方向，原 W 的功能）
|
||
c.center = tuple(np.array(c.center) + up * pan)
|
||
elif key_name == "q": # 下降（屏幕上方向负向，原 S 的功能）
|
||
c.center = tuple(np.array(c.center) - up * pan)
|
||
elif key_name == "w": # 相机沿 Z 轴上移（靠近场景）
|
||
c.center = (c.center[0], c.center[1], c.center[2] + pan)
|
||
elif key_name == "s": # 相机沿 Z 轴下移（远离场景）
|
||
c.center = (c.center[0], c.center[1], c.center[2] - pan)
|
||
else:
|
||
return
|
||
|
||
update_camera_info()
|
||
|
||
|
||
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()
|