dynamics/optimization/codex.md

Dynamics 项目优化建议

本文基于对以下文件的静态检查整理：

• compute.py
• dynamics.py
• draw.py
• engines/c/main.c
• README.md

目标不是泛泛地“提速”，而是优先找出这个项目当前最可能影响性能、内存占用、可维护性和后续扩展效率的点，并给出按优先级排序的改进方向。

一、整体判断

这个项目现在的架构已经有一个很好的基础：dynamics.py 做流程编排，compute.py 做 Python 参考实现，engines/c/main.c 提供高性能引擎，draw.py 负责可视化。

从代码结构看，当前主要瓶颈不在“外层调度”，而集中在三类地方：

1. compute.py 中逐步积分时的 Python 层循环。
2. compute.py 中弹簧力/粒子间引力的双层或逐键循环。
3. display.txt / trajectory.txt 的文本格式读写。

如果后续案例规模继续增大，Python 参考引擎会很快被这三处放大；而且即便切到 C 引擎，文本 I/O 仍然会成为新的主要瓶颈。

二、最高优先级建议

1. 先把“计算核心”和“输出格式”分开优化

当前项目已经有多语言引擎，这是很正确的方向。建议把优化目标拆成两条线分别推进：

• 计算性能：优先优化 compute.py 中的力计算和积分主循环。
• I/O 性能：优先替换或补充 display.txt / trajectory.txt 的文本格式。

原因：

• 计算核心的瓶颈主要是 CPU。
• 轨迹输出的瓶颈主要是字符串格式化、磁盘写入、加载解析。
• 这两类优化手段完全不同，混在一起做容易看不出收益来源。

建议先增加一个简单的 profiling 开关，例如输出：

• 总模拟时间
• 力计算时间
• 抽帧时间
• 保存 display.txt 时间
• 保存 trajectory.txt 时间
• 加载 display.txt 时间

这样后面每做一次优化，都能知道收益来自哪一段。

2. 优先向量化 compute_force() 里的弹簧键计算

compute.py:1230 开始的 compute_force() 是 Python 引擎最关键的热点。

其中弹簧键部分目前是：

• 遍历 BOND_PAIRS
• 每根键单独算 dx/dy/dz
• 每根键单独回写两个原子的受力

这在键数增大后会变成明显瓶颈。

建议做法：

• 把 BOND_PAIRS[:, 0] 和 BOND_PAIRS[:, 1] 拆成两个索引数组 i_idx、j_idx。
• 用 NumPy 一次性计算所有键的 dx/dy/dz/dist/stretch。
• 用 np.add.at() 或等价的 scatter 累加方式，把键力一次性回写到 fx/fy/fz。

预期收益：

• 中等规模体系下，Python 参考引擎会有很可观的提速。
• 这也是最不改变项目结构、最容易验证正确性的优化。

3. 把原子间万有引力视为“可选高成本模块”，不要默认走 Python 双层循环

compute.py:1281 开始的 GRAVITY_INTERACTION 采用 for i / for j 双层循环，复杂度是 O(N^2)。

这段代码在教学上没问题，但在稍大的粒子数下会非常慢。建议：

• 明确把它标记为“仅适合小规模案例”。
• 默认推荐用户用 engine: c 跑这类 case。
• 如果未来仍想保留 Python 版本，优先考虑：
  • 小规模时保持现状。
  • 中大规模时改成 Numba 或 C 扩展。
  • 更进一步再考虑 Barnes-Hut、cell list、neighbor list 之类近似/加速算法。

如果只是当前项目阶段，我更建议：

• 不要先在 Python 里硬做复杂天体优化。
• 先把这类高复杂度场景明确导向 C 引擎。

这是投入产出比更高的路线。

4. display.txt 文本格式要补一个二进制版本

compute.py:196 的 save_display_txt() 和 compute.py:239 的 load_display_txt() 已经比 JSON 好很多，但本质仍然是大文本。

当前问题：

• 保存时每个数都在做字符串格式化。
• 读取时虽然用了 np.genfromtxt，但源数据仍是文本。
• 帧数和粒子数继续增大后，I/O 会成为明显瓶颈。

建议保留 display.txt 作为“人可读格式”，同时新增一个高性能格式，例如：

• display.npz
• 或 display.npy + meta.json

推荐方案：

• frames_x/y/z/vx/vy/vz 存到 np.savez_compressed
• 元信息单独存一个轻量 meta.json
• draw.py 优先读二进制，不存在时再回退到 display.txt

这样有几个好处：

• 调试和教学仍可保留文本文件。
• 大规模运行时可以直接避开文本解析开销。
• Python 和 C 引擎都可以逐步迁移，不需要一次切完。

三、中优先级建议

5. run_simulation() 里完整轨迹缓存要改成“按需流式写出”

compute.py:1481 起，如果 save_trajectory=1，会一次性申请：

• traj_x
• traj_y
• traj_z
• traj_vx
• traj_vy
• traj_vz

这意味着内存复杂度接近 O(steps * atoms)，而且还是 6 份 float64 数组。

这在教学小案例里没问题，但一旦：

• NT 很大
• 粒子数上来
• 同时还保留抽帧缓存

内存会膨胀得很快。

建议：

• 如果只是为了最终导出，改成边算边写。
• 如果还需要后续随机访问，可以改成 memmap 或分块写入 npz/hdf5/zarr。

推荐优先级：

• 先做“保存完整轨迹时改为 chunked binary writer”。
• trajectory.txt 保留为兼容输出，而不是默认主输出。

6. draw.py 的 Marker 更新应一次性切片赋值，避免逐原子 Python 循环

draw.py:517 附近的 _update_atom_positions() 在 USE_MARKER 模式下仍然逐原子循环：

• 先 for i in range(N_ATOMS)
• 再逐个写 marker_pos[i]
• 然后 balls.set_data(pos=marker_pos)

建议直接改成：

• marker_pos[:, 0] = DISP_ALL_X[f_idx]
• marker_pos[:, 1] = DISP_ALL_Y[f_idx]
• marker_pos[:, 2] = DISP_ALL_Z[f_idx]

这样更符合 Marker 模式“批量更新”的初衷。

同理，成键线 _update_bond_positions() 也可以进一步尝试批量索引构造，而不是逐键更新。

虽然这部分通常不如计算核心慢，但在大粒子数动画里会直接影响帧率。

7. apply_fixed_constraints() 每步构造 column_stack，可以改成原地掩码写回

compute.py:1408 的实现每步都会：

• np.column_stack((x, y, z))
• np.column_stack((vx, vy, vz))
• np.where(...)

这会产生额外临时数组。

建议改成：

• 预先缓存 fixed_x/fixed_y/fixed_z 三个布尔掩码
• 直接对 x/y/z/vx/vy/vz 原地赋值

例如思路上改成：

• x[fixed_x] = ATOM_POSITIONS[fixed_x, 0]
• vx[fixed_x] = 0.0

这类优化单项收益不一定最大，但由于它在每一步都执行，累计下来是有价值的。

8. apply_driving_force() 有重复小数组分配

compute.py:599 的驱动力逻辑中，每个 driver、每一步都构造：

• t_vec = np.array([t, t, t], dtype=np.float64)

这是典型的小对象重复分配。

建议：

• 直接分别计算三个轴，不需要生成 t_vec
• 或预先把 driver 参数整理成矩阵，批量更新受驱原子

如果 driver 数量很少，这不是最大瓶颈；但这类微优化很容易做，而且不会增加复杂度。

四、架构与可维护性建议

9. 减少 compute.py 的全局变量依赖，逐步收敛到状态对象

目前 compute.py 依赖大量全局变量，例如：

• ATOM_IDS
• ATOM_MASSES
• BOND_PAIRS
• METHOD
• NT
• DT

这让代码在以下场景下会越来越难维护：

• 并行运行多个 case
• 写单元测试
• 替换不同 force model
• 将来做 GUI 或服务化封装

建议中期做一个 SimulationState / SimulationConfig / SystemData 分层：

• 配置类：步长、方法、开关、输出选项
• 系统类：原子、键、边界、驱动参数
• 状态类：当前 x/y/z/vx/vy/vz

不需要一次性重构完，先从最核心的 compute_force()、run_simulation() 入手即可。

10. Python 参考引擎和 C 引擎的“功能等价层”建议更明确

README 里已经强调了多语言对比，这是项目亮点。为了后续更稳，建议把“等价层”标准化：

• 相同输入
• 相同输出语义
• 相同采样规则
• 相同边界行为
• 相同 driver 行为

然后做一个最小一致性测试集：

• case01-case06 都能跑
• Python 与 C 输出在容差内一致
• 不同 method 的结果有基准对照

这样后续做任何优化时，都更容易大胆改，不怕悄悄改坏物理行为。

五、需要尽快处理的正确性/维护风险

这些不一定直接是“性能问题”，但会影响后续优化效率，建议优先修一下。

11. dynamics.py 的绘图分支存在明显变量依赖不完整的风险

在 dynamics.py:331 一带，绘图代码直接使用：

• all_x
• all_y
• all_z
• all_vx
• all_vy
• all_vz
• data

但从当前文件上下文看，这些变量只在某些分支里才会存在，尤其 Python 引擎路径下很可能未定义。

这会带来两个问题：

• 某些 case 可能直接在绘图阶段报错。
• 优化时很难判断问题来自性能还是流程分支。

建议：

• 在进入绘图逻辑前统一构造标准数据对象。
• 不要依赖分支里“顺便留下来的局部变量”。

12. README.md 描述与当前实现已经有部分不一致

例如 README 仍强调：

• trajectory.txt (JSON, 统一格式)
• sample.py -> display.txt

但当前实现里：

• Python 路径已经直接写 display.txt
• save_trajectory 变成可选
• sample.py 在主流程中的角色已经下降

文档不一致本身不会拖慢程序，但会拖慢后续协作和排障效率，尤其在你继续演进输出格式时会更明显。

建议在做 I/O 优化时顺手更新 README，避免认知分叉。

六、推荐实施顺序

如果按“最少改动获得最大收益”的原则，我建议这样排：

1. 给 dynamics.py / compute.py 增加基础计时统计。
2. 向量化 compute_force() 的弹簧键计算。
3. 把 draw.py 的 Marker 更新改成切片赋值。
4. 新增 display.npz，draw.py 优先读取二进制。
5. 把 save_trajectory=1 改成分块二进制输出，而不是全量内存缓存。
6. 修复 dynamics.py 绘图分支的数据来源问题。
7. 再考虑 compute.py 全局变量收敛和更深层的结构重构。

七、如果只做三件事，最值得做什么

如果你现在只想投入一小轮精力，我建议只做这三项：

1. 向量化 compute_force() 的弹簧键部分。
2. 增加 display.npz 二进制输出与读取。
3. 修正 dynamics.py 绘图阶段的数据流一致性。

原因：

• 第 1 项直接优化 Python 引擎核心热点。
• 第 2 项直接优化所有引擎共享的 I/O 瓶颈。
• 第 3 项能降低后续改动时的维护风险。

这三项一起做，收益通常比零散微优化更明显。

八、结论

这个项目最值得肯定的地方，是已经天然分成了：

• 参考实现
• 高性能引擎
• 统一可视化管线

这意味着它非常适合做“分层优化”，不需要推倒重来。

从当前代码看，后续最有效的路线不是继续在外围加流程判断，而是：

• 把 Python 热点循环尽量向量化
• 把文本轨迹格式逐步替换为二进制主格式
• 把数据流和状态管理收紧

如果按这个方向推进，这个项目会同时得到：

• 更好的性能
• 更低的内存占用
• 更稳定的多引擎一致性
• 更容易继续扩展新的物理项和可视化形式