From 757a891a43e3e628a536dcd0c3f619e832fd8b2d Mon Sep 17 00:00:00 2001
From: Ying-Li Niu <64801511@qq.com>
Date: Fri, 12 Jun 2026 07:46:28 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20display.txt=20=E4=BC=A0=E9=80=92?=
 =?UTF-8?q?=E5=8E=9F=E5=AD=90=E5=8D=8A=E5=BE=84=E6=95=B0=E7=BB=84=EF=BC=8C?=
 =?UTF-8?q?draw.py=20=E8=AF=BB=E5=8F=96=20per-atom=20=E5=8D=8A=E5=BE=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

之前所有原子使用统一的 ball_radius（来自 input.txt），
现在 display.txt header 包含 atom_radii 字段（逗号分隔
的半径数组），draw.py 据此为每个原子设置独立半径。
fallback：若没有 atom_radii 字段，仍使用 ball_radius。
---
 compute.py  | 3 ++-
 draw.py     | 7 ++++++-
 dynamics.py | 3 ++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/compute.py b/compute.py
index 6fbb319..b35bdcc 100644
--- a/compute.py
+++ b/compute.py
@@ -1460,7 +1460,8 @@ def run_simulation(save_trajectory=0):
                        "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]))}
+                       "alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])),
+                       "atom_radii": ",".join(str(r) for r in ATOM_RADII)}
     )
     print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)")
 
diff --git a/draw.py b/draw.py
index 3c48389..e8656f5 100644
--- a/draw.py
+++ b/draw.py
@@ -67,7 +67,12 @@ DISP_T    = DISP_STEP * DT
 
 # 原子信息
 ATOM_IDS  = disp_data["atom_ids"]
-ATOM_RADII = np.full(N_ATOMS, float(h.get("ball_radius", 0.5)))
+# 优先使用 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 格式不含成键信息，从原始数据加载
diff --git a/dynamics.py b/dynamics.py
index 68f5a36..1842a3f 100644
--- a/dynamics.py
+++ b/dynamics.py
@@ -253,7 +253,8 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output",
                   "gravity_strength": str(data.get("gravity_strength", 1.0)),
                   "driving_force": str(data.get("driving_force", 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_radii": _fmt_alpha(data.get("atom_radii", []))}
             
             n_frames = len(indices)
             compute.save_display_txt(