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

fix: color_xrange 范围限定 d_max + color_driver + color 范围改为 [0,1] + bug 修复

Browse Source
This commit is contained in:
Ying-Li Niu
2026-06-24 16:09:39 +08:00
parent b50c55d0a3
commit e415bf5ba5
87 changed files with 360440 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
draw.py
+280 -60
View File
@@ -45,7 +45,65 @@ if os.path.exists(npz_path):
disp_data = compute.load_display_npz(npz_path)
else:
disp_data = compute.load_display_txt(disp_path)
h = disp_data["header_fields"]
# ── 从 input.txt 读取参数（替代 display.npz 中的 meta）──
try:
import yaml
_have_yaml = True
except ImportError:
_have_yaml = False
input_dir = os.path.join(os.path.dirname(output_dir), "input")
input_path = os.path.join(input_dir, "input.txt")
if _have_yaml and os.path.exists(input_path):
try:
with open(input_path, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
except Exception:
config = {}
else:
config = {}
# 兼容旧版：若 input.txt 不存在或解析失败，降级到 display.npz 的 meta
if not config:
config = disp_data.get("header_fields", {})
# ── 从 coord.txt 读取平衡位置 ─────────────────
# 注意: config 中的 *_file 可能带 "input/" 前缀，但 input_dir 已指向 input/ 目录
_coord_file_raw = config.get("coord_file", "coord.txt")
coord_path = os.path.join(input_dir, os.path.basename(_coord_file_raw))
if os.path.exists(coord_path):
try:
_ids, _masses, _radii, _pos, _vel, _fixed = compute.load_coord_file(coord_path)
EQ_POS = _pos # (n_atoms, 3)
ATOM_FIXED = _fixed # (n_atoms, 3)
except Exception:
EQ_POS = None
ATOM_FIXED = None
else:
EQ_POS = None
ATOM_FIXED = None
# ── 从 connection.txt 读取成键信息（若 meta 中没有）──
BOND_PAIRS = disp_data.get("bond_pairs", [])
if not BOND_PAIRS and 'bond_pairs' not in disp_data:
conn_path = os.path.join(input_dir, os.path.basename(config.get("connection_file", "connection.txt")))
if os.path.exists(conn_path):
try:
_pairs = []
with open(conn_path, "r", encoding="utf-8") as _f:
_f.readline() # skip header
for _line in _f:
_line = _line.strip()
if not _line or _line.startswith("#"):
continue
_parts = _line.split()
if len(_parts) >= 2:
_pairs.append([int(_parts[0]) - 1, int(_parts[1]) - 1])
BOND_PAIRS = np.array(_pairs, dtype=np.int32) if _pairs else []
except Exception:
BOND_PAIRS = []
BOND_PAIRS = BOND_PAIRS.tolist() if hasattr(BOND_PAIRS, 'tolist') else BOND_PAIRS
# 全原子帧数据
DISP_ALL_X = disp_data["frames_x"] # (n_frames, n_atoms)
@@ -66,64 +124,220 @@ 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))
DT = float(config.get("DT", 0.001))
# 视觉位移放大：display_amp: [ax, ay, az]，对偏离第0帧的位移乘以倍数
_damp_raw = h.get("display_amp", "")
if _damp_raw.strip():
# 视觉位移放大：display_amp: [ax, ay, az]，对偏离平衡位置的位移乘以倍数
_damp_raw = config.get("display_amp", "")
if isinstance(_damp_raw, (list, tuple)):
_damp = np.array(_damp_raw, dtype=np.float64)
elif isinstance(_damp_raw, str) and _damp_raw.strip():
import ast as _ast
_damp_vals = _ast.literal_eval(_damp_raw.strip())
_damp = np.array(_damp_vals, dtype=np.float64)
if _damp.shape == (3,) and not np.allclose(_damp, 1.0):
_eq_x = DISP_ALL_X[0:1, :] # 第0帧作为平衡位置参考
_eq_y = DISP_ALL_Y[0:1, :]
_eq_z = DISP_ALL_Z[0:1, :]
_damp = np.array(_ast.literal_eval(_damp_raw.strip()), dtype=np.float64)
else:
_damp = None
if _damp is not None and _damp.shape == (3,) and not np.allclose(_damp, 1.0):
if EQ_POS is not None:
_eq_x = EQ_POS[None, :, 0] # coord.txt 平衡位置
_eq_y = EQ_POS[None, :, 1]
_eq_z = EQ_POS[None, :, 2]
else:
_eq_x = DISP_ALL_X[0:1, :] # 第0帧作为平衡位置参考
_eq_y = DISP_ALL_Y[0:1, :]
_eq_z = DISP_ALL_Z[0:1, :]
DISP_ALL_X = _eq_x + (DISP_ALL_X - _eq_x) * _damp[0]
DISP_ALL_Y = _eq_y + (DISP_ALL_Y - _eq_y) * _damp[1]
DISP_ALL_Z = _eq_z + (DISP_ALL_Z - _eq_z) * _damp[2]
NSTEP = int(h.get("NSTEP", 1))
# ── 位移颜色映射 ──────────────────────────────
# display_color: { mode: [enabled, [R,G,B]], ... }
# mode: x, y, z, xy, yz, zx, xyz
# enabled: 0=off, 1=on
# [R,G,B]: 最大位移时的颜色（渐变起点为白色）
# 例：xyz: [1, [1.0,1.0,1.0]] → 三方向合成位移 → 白色渐变
FRAME_COLORS = None
_DC_MODE = None
_DC_COLOR = None
_raw_dc = config.get("display_color", "")
if _raw_dc:
try:
dc = json.loads(_raw_dc) if isinstance(_raw_dc, str) else _raw_dc
if isinstance(dc, dict):
for mode, (enabled, color) in dc.items():
if int(enabled) and mode in ("x","y","z","xy","yz","zx","xyz"):
_DC_MODE = mode
_DC_COLOR = np.array(color, dtype=np.float32)
break
except Exception:
pass
if _DC_MODE is not None and _DC_COLOR is not None:
if EQ_POS is not None:
eq_x = EQ_POS[None, :, 0] # coord.txt 平衡位置
eq_y = EQ_POS[None, :, 1]
eq_z = EQ_POS[None, :, 2]
else:
eq_x = DISP_ALL_X[0:1, :] # 第0帧
eq_y = DISP_ALL_Y[0:1, :]
eq_z = DISP_ALL_Z[0:1, :]
dx = DISP_ALL_X - eq_x
dy = DISP_ALL_Y - eq_y
dz = DISP_ALL_Z - eq_z
if _DC_MODE == "x":
disp_mag = np.abs(dx)
elif _DC_MODE == "y":
disp_mag = np.abs(dy)
elif _DC_MODE == "z":
disp_mag = np.abs(dz)
elif _DC_MODE == "xy":
disp_mag = np.sqrt(dx**2 + dy**2)
elif _DC_MODE == "yz":
disp_mag = np.sqrt(dy**2 + dz**2)
elif _DC_MODE == "zx":
disp_mag = np.sqrt(dz**2 + dx**2)
else: # xyz
disp_mag = np.sqrt(dx**2 + dy**2 + dz**2)
# ── color_xrange: 限定归一化基准的原子范围 ───────────
# 格式: [['min'|'mid'|'max'|数值, 'min'|'mid'|'max'|数值], ...] 对应 x,y,z
# 在此范围内的原子计算最大位移 d_max，所有原子以此基准归一化着色
_cxr = config.get("color_xrange", None)
_d_max_source = disp_mag # 默认：所有原子
_range_label = "全部原子"
if _cxr and isinstance(_cxr, (list, tuple)) and len(_cxr) == 3:
try:
# 获取坐标极值
if EQ_POS is not None:
_eq_all = EQ_POS
else:
_eq_all = np.column_stack([
DISP_ALL_X[0], DISP_ALL_Y[0], DISP_ALL_Z[0]])
_cmin = _eq_all.min(axis=0)
_cmax = _eq_all.max(axis=0)
_cmid = (_cmin + _cmax) / 2
_key_map = {"min": _cmin, "mid": _cmid, "max": _cmax}
_range_lo = np.zeros(3, dtype=np.float64)
_range_hi = np.zeros(3, dtype=np.float64)
for _i in range(3):
_lo = _cxr[_i][0]
_hi = _cxr[_i][1]
_range_lo[_i] = _key_map[_lo][_i] if _lo in _key_map else float(_lo)
_range_hi[_i] = _key_map[_hi][_i] if _hi in _key_map else float(_hi)
_in_x = (_eq_all[:, 0] >= _range_lo[0]) & (_eq_all[:, 0] <= _range_hi[0])
_in_y = (_eq_all[:, 1] >= _range_lo[1]) & (_eq_all[:, 1] <= _range_hi[1])
_in_z = (_eq_all[:, 2] >= _range_lo[2]) & (_eq_all[:, 2] <= _range_hi[2])
_color_mask = _in_x & _in_y & _in_z
_n_in_range = _color_mask.sum()
if _n_in_range > 0:
_d_max_source = disp_mag[:, _color_mask] # 仅在范围内找最大位移
_range_label = (f"x[{_range_lo[0]:.1f},{_range_hi[0]:.1f}] "
f"y[{_range_lo[1]:.1f},{_range_hi[1]:.1f}] "
f"z[{_range_lo[2]:.1f},{_range_hi[2]:.1f}]")
except Exception as _e:
print(f"[draw] color_xrange 解析失败: {_e}")
# 用（范围限定的）最大位移归一化，所有原子统一着色
d_max = _d_max_source.max()
if d_max > 1e-12:
disp_norm = disp_mag / d_max
else:
disp_norm = disp_mag
FRAME_COLORS = np.ones((N_FRAMES, N_ATOMS, 4), dtype=np.float32)
t = disp_norm[..., None]
FRAME_COLORS[..., :3] = 1.0 + (_DC_COLOR - 1.0) * t
print(f"[draw] 位移颜色映射: mode={_DC_MODE}, color={_DC_COLOR.tolist()}, "
f"d_max={d_max:.4f}, 范围: {_range_label}")
# ── color_driver: 驱动原子 → 统一颜色 ──
_cd_raw = config.get("color_driver", None)
if _cd_raw is not None:
try:
_cd = np.array(_cd_raw, dtype=np.float32)
# 读取 driver.txt 获取驱动原子 ID
_driver_file = config.get("driver_file", "driver.txt")
_driver_path = os.path.join(input_dir, os.path.basename(_driver_file))
if os.path.exists(_driver_path):
_driver_ids = []
with open(_driver_path, "r", encoding="utf-8") as _df:
_df.readline() # skip header
for _line in _df:
_line = _line.strip()
if not _line or _line.startswith("#"):
continue
_parts = _line.split()
if _parts:
_driver_ids.append(int(_parts[0]) - 1) # 1-based → 0-based
if _driver_ids:
FRAME_COLORS[:, _driver_ids, :3] = _cd
print(f"[draw] color_driver: {len(_driver_ids)} 驱动原子 → RGB{_cd_raw}")
except Exception as _e:
print(f"[draw] color_driver 解析失败: {_e}")
# ── color_fix: 全固定原子 (fix_x=fix_y=fix_z=1) → 统一颜色 ──
_cf_raw = config.get("color_fix", None)
if _cf_raw is not None and ATOM_FIXED is not None:
try:
_cf = np.array(_cf_raw, dtype=np.float32)
_full_fixed = (ATOM_FIXED[:, 0] == 1) & (ATOM_FIXED[:, 1] == 1) & (ATOM_FIXED[:, 2] == 1)
_n_fix = _full_fixed.sum()
if _n_fix > 0:
FRAME_COLORS[:, _full_fixed, :3] = _cf
print(f"[draw] color_fix: {_n_fix} 全固定原子 → RGB{_cf_raw}")
except Exception as _e:
print(f"[draw] color_fix 解析失败: {_e}")
NSTEP = int(config.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", "")
_raw_radii = config.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)))
ATOM_RADII = np.full(N_ATOMS, float(config.get("ball_radius", 0.5)))
PLOT_ATOM_ROW = 0
PLOT_ATOM_ID = int(ATOM_IDS[0])
BOND_PAIRS = [] # display 格式不含成键信息，从原始数据加载
# 成键信息已在上面从 connection.txt 加载
# 渲染方式：0=Sphere(网格球体), 1=Marker(GPU点精灵)
USE_MARKER = int(h.get("use_marker", 0))
USE_MARKER = int(config.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 (ValueError, AttributeError):
alpha_list = [float(raw_alpha)] * 6
# 模拟边界（从 input.txt 的 box_a 计算）
_box_a = float(config.get("box_a", 10.0))
X_MIN = -_box_a; X_MAX = _box_a
Y_MIN = -_box_a; Y_MAX = _box_a
Z_MIN = -_box_a; Z_MAX = _box_a
raw_alpha = config.get("alpha", "0.2")
if isinstance(raw_alpha, (list, tuple)):
alpha_list = [float(x) for x in raw_alpha]
else:
try:
alpha_list = [float(x) for x in raw_alpha.split(",")]
except (ValueError, AttributeError):
alpha_list = [float(raw_alpha)] * 6
if len(alpha_list) != 6:
alpha_list = (alpha_list * 6)[: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))
ball_radius = float(config.get("ball_radius", 0.5))
ball_color_r = float(config.get("ball_color_r", 0.9))
ball_color_g = float(config.get("ball_color_g", 0.2))
ball_color_b = float(config.get("ball_color_b", 0.2))
box_color_r = float(config.get("box_color_r", 0.8))
box_color_g = float(config.get("box_color_g", 0.8))
box_color_b = float(config.get("box_color_b", 0.85))
# ===========================================================================
@@ -135,23 +349,23 @@ axis_length = 10.0
import math as _math_cam
_cx = float(h.get("camera_center_x", 0.0))
_cy = float(h.get("camera_center_y", 0.0))
_cz = float(h.get("camera_center_z", 0.0))
_cx = float(config.get("camera_center_x", 0.0))
_cy = float(config.get("camera_center_y", 0.0))
_cz = float(config.get("camera_center_z", 0.0))
# 若 input.txt 指定了摄像机自身坐标，则由坐标反推 distance/elevation/azimuth
if h.get("camera_pos_x") is not None:
_px = float(h["camera_pos_x"])
_py = float(h["camera_pos_y"])
_pz = float(h["camera_pos_z"])
if config.get("camera_pos_x") is not None:
_px = float(config.get("camera_pos_x"))
_py = float(config.get("camera_pos_y"))
_pz = float(config.get("camera_pos_z"))
_dx, _dy, _dz = _px - _cx, _py - _cy, _pz - _cz
_dist = _math_cam.sqrt(_dx*_dx + _dy*_dy + _dz*_dz) or 1.0
_elev = _math_cam.degrees(_math_cam.asin(max(-1.0, min(1.0, _dy / _dist))))
_azim = _math_cam.degrees(_math_cam.atan2(_dx, _dz))
else:
_dist = float(h.get("camera_distance", 40.0))
_elev = float(h.get("camera_elevation", 0))
_azim = float(h.get("camera_azimuth", 0))
_dist = float(config.get("camera_distance", 40.0))
_elev = float(config.get("camera_elevation", 0))
_azim = float(config.get("camera_azimuth", 0))
initial_camera = {
"distance": _dist,
@@ -212,11 +426,11 @@ axes_group.append(scene.visuals.Arrow(
parent=view.scene,
))
axes_group.append(scene.visuals.Text(text="x", color=(1.0, 0.2, 0.2, 1.0), font_size=18,
axes_group.append(scene.visuals.Text(text="x", color=(1.0, 0.2, 0.2, 1.0), font_size=14,
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,
axes_group.append(scene.visuals.Text(text="y", color=(0.2, 1.0, 0.2, 1.0), font_size=14,
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,
axes_group.append(scene.visuals.Text(text="z", color=(0.3, 0.6, 1.0, 1.0), font_size=14,
pos=(0, 0, axis_length + 0.2), anchor_x="left", anchor_y="bottom", parent=view.scene))
# ── 原子渲染 ──────────────────────────────────
@@ -235,11 +449,14 @@ TAB10_RGB = np.array([
[0.7373, 0.7412, 0.1333], # 黄绿
[0.0902, 0.7451, 0.8118], # 青
])
# 每个原子的颜色（循环使用 tab10 色板）
# 每个原子的颜色（循环使用 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 FRAME_COLORS is not None:
atom_colors[:] = FRAME_COLORS[0] # 初始帧颜色
else:
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 ──
@@ -295,12 +512,12 @@ for f_idx, (pos, direction) in enumerate(faces):
# 右上角：相机信息
camera_info = scene.visuals.Text(
text="", color="white", font_size=14,
text="", color="white", font_size=12,
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,
text="", color=(0.2, 1.0, 0.2, 1.0), font_size=14,
pos=(0, 0), anchor_x="left", anchor_y="top",
face="黑体", bold=True, parent=canvas.scene)
@@ -312,7 +529,7 @@ reset_button = scene.visuals.Rectangle(
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,
text="reset", color="white", font_size=13,
pos=(reset_btn_size[0] / 2 + 8, reset_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
@@ -326,7 +543,7 @@ info_button = scene.visuals.Rectangle(
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,
text="info", color="white", font_size=13,
pos=(info_btn_size[0] / 2 + 8, info_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
@@ -346,7 +563,7 @@ axes_button = scene.visuals.Rectangle(
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,
text="axes", color="white", font_size=13,
pos=(axes_btn_size[0] / 2 + 8, axes_btn_size[1] / 2 + 8),
anchor_x="center", anchor_y="center",
bold=True, parent=canvas.scene)
@@ -551,12 +768,15 @@ def handle_mouse_press(event):
# ===========================================================================
def _update_atom_positions(f_idx):
"""更新所有原子到第 f_idx 帧的位置。"""
"""更新所有原子到第 f_idx 帧的位置，必要时更新颜色。"""
if USE_MARKER:
marker_pos[:, 0] = DISP_ALL_X[f_idx]
marker_pos[:, 1] = DISP_ALL_Y[f_idx]
marker_pos[:, 2] = DISP_ALL_Z[f_idx]
balls.set_data(pos=marker_pos)
if FRAME_COLORS is not None:
balls.set_data(pos=marker_pos, face_color=FRAME_COLORS[f_idx])
else:
balls.set_data(pos=marker_pos)
else:
for i in range(N_ATOMS):
balls[i].transform = STTransform(translate=(
@@ -630,10 +850,10 @@ def _load_move_camera_txt():
# 先试 move_camera.txt 直读，没有则用 display.txt 缓存
# header 中 camera_keyframes 为空字符串表示 move_camera=0（开关关闭），跳过文件加载
_camera_motion_enabled = bool(h.get("camera_keyframes", ""))
_camera_motion_enabled = bool(config.get("camera_keyframes", ""))
_CAM_MOTION = _load_move_camera_txt() if _camera_motion_enabled else None
if not _CAM_MOTION:
_CAM_MOTION = json.loads(h.get("camera_keyframes", "null")) if h.get("camera_keyframes") else None
_CAM_MOTION = json.loads(config.get("camera_keyframes", "null")) if config.get("camera_keyframes") else None
if _CAM_MOTION:
_cam_center = [0.0, 0.0, 0.0]
_cam_elev = initial_camera["elevation"]