PXAID

PXAID

PYXAID是由Oleg V. Prezhdo和Alexey V. Akimov共同开发的优秀非绝热分子动力学软件，在凝聚态体系的激发态动力学模拟中得到了广泛应用。随后，Alexey V. Akimov开发的Libra进一步拓展了相关方法体系。此外，Hefei-NAMD、NEXMD、SHARC和Newton-X等软件也在激发态动力学研究中得到了广泛应用，并表现出良好的性能。

在PYXAID基础上，本项目开发了PXAID，一个基于Python语言的轻量化重构版本。该实现基于原程序的代码逻辑进行重新实现，并采用更加现代化的Python数值计算框架。在保留从头算电子-核运动耦合核心理论框架的基础上，充分结合NumPy、SciPy以及Numba等高性能科学计算库，实现了非绝热分子动力学过程的高效模拟。其中，核心传播、表面跳跃与退相干等关键计算模块采用Numba JIT加速，并结合FFT优化自相关函数与退相干速率计算，从而显著提升长时间轨迹下的计算效率。同时，PXAID引入了改进的表面跳跃算法（如SDM方法）以及更加灵活的模块化设计，从而提升了程序在复杂体系中的适用性与可扩展性。此外，PXAID支持Linux与Windows等多平台环境运行，适用于大规模计算以及后续与机器学习方法的集成。

使用方法

julia namd.jl

性能对比

在以下两种硬件平台上，对PYXAID、MAXAID和JUXAID不同版本进行了性能测试：

• 服务器平台：英特尔® 至强® 金牌 6444Y，256 GB 内存，CentOS

• 桌面平台：Intel(R) Core(TM) Ultra 9 285H，32 GB 内存，Windows 11

相对速度对比表

软件版本	相对速度（PYXAID=1）	性能表现说明
PYXAID	1×	基准参考
MAXAID	1/15×	电子态数增加时性能显著下降
JUXAID (v10)	1×	与 PYXAID 基本一致
JUXAID (v12)	6×	优化
JUXAID (v14)	13×	优化
JUXAID (v15)	19×	关闭部分非必要输出
PXAID (v10)	1/36×	初始 Python 重构版本，性能较低
PXAID (v12)	1/15×	优化
PXAID (v13)	11×	优化
PXAID (v14)	34×	优化
PXAID (v15)	154×	关闭部分非必要输出

更新日志

v15 相比 v14 的主要改进

1. 增加输出文件保存开关

• v14 默认保存全部输出文件

• v15 增加四类输出控制参数：

  • SAVE_DECOHERENCE_RATES

  • SAVE_ICOND_FILES

  • SAVE_ME_ENERGIES

  • SAVE_ME_POP

• 用户可根据需求灵活选择是否保存对应文件

2. 降低文件IO开销

• v14 会持续写入多类文本文件

• v15 可关闭部分非必要输出

• 在大规模 icond 和长时间轨迹计算中可显著减少写文件耗时

v14 相比 v13 的主要改进

1. 引入多线程并行计算（Numba parallel）

• v13 主要采用单线程轨迹传播

• v14 在核心传播函数中引入：

  • parallel=True

  • prange

• 实现多条 surface hopping trajectories 并行计算

• 显著提升大规模轨迹模拟效率

2. 引入线程局域统计数组（TLS）

• v13 所有轨迹共享：

  • sh_pops

  • se_pops

• v14 为每个线程建立：

  • sh_tls

  • se_tls

• 降低多线程写冲突与同步开销

3. 多轨迹任务静态分块

• v13 顺序遍历：

  • for itraj in range(num_sh_traj)

• v14 按线程静态划分轨迹区间：

  • lo

  • hi

• 提高线程负载均衡与缓存局部性

4. 多线程随机数流独立化

• v13 所有轨迹共享随机数序列

• v14 使用：

  • seed + 1000003 * tid

• 为不同线程建立独立随机数流

• 降低并行随机数相关性

5. 核心NAMD循环完全并行化

• v13 仅核心数值函数JIT化

• v14 将：

  • surface hopping

  • electronic propagation

  • decoherence

  • population accumulation

  整体集成到并行JIT核心中

• 进一步降低Python层调度开销

6. 减少共享数组竞争

• v13 所有轨迹直接更新全局统计数组

• v14 使用线程局域累加后统一归约

• 提高多核扩展效率

7. 并行传播与JIT进一步融合

• v13 使用：

  • @njit(cache=True, fastmath=True)

• v14 核心传播函数改为：

  • @njit(fastmath=True, parallel=True)

• 进一步提升热点循环执行效率

8. 多线程下保持数组扁平化结构

• v13 已采用：

  • np.ndarray[:,:, :]

  批量存储 Hamiltonian

• v14 在并行环境下继续保持连续内存布局

• 提高缓存命中率与线程数据访问效率

9. 并行版本进一步降低Python解释器参与度

• v13 部分轨迹控制仍存在Python层调度

• v14 将多轨迹传播完全下沉至JIT并行核心

• 降低解释器与函数调度开销

10. 显著提升整体运行性能

• v13 相对 PYXAID 速度约为：

  • 11×

• v14 在多线程环境下进一步提升

• 在当前测试体系中性能超过 JUXAID v14

v13 相比 v12 的主要改进

1. 引入Numba JIT加速核心数值计算

• v12 主要依赖 Python 解释执行

• v13 对核心热点函数引入：

  • @njit(cache=True, fastmath=True)

• 显著降低 Python 解释器开销

2. 核心电子动力学完全JIT化

• v12 使用普通 Python 函数执行传播

• v13 将：

  • 电子传播

  • 表面跳跃

  • 退相干

  • population更新

  等核心模块全部编译为机器码执行

• 大幅提升长时间动力学模拟效率

3. FFT实现进一步优化

• v12 使用：

  • np.fft.fft

  • np.fft.ifft

• v13 在实值退相干相关函数计算中改为：

  • scipy.fft.rfft

  • scipy.fft.irfft

• 利用实数能量涨落序列减少FFT计算量与内存占用

4. 引入next_fast_len优化FFT长度

• v12 使用固定2幂长度补零

• v13 使用：

  • next_fast_len

• 自动选择更高效FFT长度

5. 三维数组批量存储Hamiltonian

• v12 使用 list[np.ndarray]

• v13 改为：

  • np.ndarray[:,:, :]

• 提高内存连续性与缓存命中率

6. 核心传播函数“去对象化”

• v12 使用 ElectronicStructure 对象传播

• v13 核心JIT部分直接操作：

  • C

  • A

  • g

  • tau_m

  • t_m

  等底层数组

• 降低对象访问与动态调度开销

7. 减少临时数组与Python层循环

• v12 仍存在部分 Python 层循环

• v13 在JIT函数中大量采用显式数组原位更新

• 显著降低内存分配与GC压力

8. Hamiltonian构造进一步JIT化

• v12 使用普通 Python 构造 Hme_batch

• v13 引入：

  • _build_spin_orbital_H_jit

  • _build_Hme_batch_jit

• 加速多电子态哈密顿量构建过程

9. 随机跳跃过程JIT化

• v12 使用 Python random

• v13 在核心传播中使用：

  • np.random

  并集成到JIT函数内部

• 提高表面跳跃模拟效率

10. 显著提升整体运行性能

• v12 相对 PYXAID 速度约为：

  • 0.067×

• v13 相对 PYXAID 速度达到：

  • 11×

• 实现数量级性能提升

v12 相比 v10 的主要改进

1. 自相关函数FFT加速

• v10 使用双重循环计算自相关函数（O(N²)）

• v12 使用 FFT 实现严格等价的自相关计算（O(N log N)）

• 在长时间轨迹下显著提升decoherence计算速度

2. 参数结构化（移除大量全局变量）

• v10 使用大量全局变量（hbar、kb、dt等）

• v12 引入 NAMDParams 数据类统一管理参数

• 提高代码可维护性与函数接口一致性

3. 多电子态映射预计算

• v10 在每个时间步重复计算电子态映射

• v12 预计算：

  • transition_map

  • diag_orbital_map

• 显著减少重复计算开销

4. Hamiltonian构造模块化

• v10 在主程序中集中构造 Hamiltonian

• v12 拆分为多个独立函数

• 提高代码可读性与后续扩展能力

5. 避免频繁对象复制

• v10 使用多个 ElectronicStructure 对象列表传播

• v12 改为单对象逐步更新传播

• 降低内存占用与对象管理开销

6. 数组操作优化

• v10 存在较多 Python 层循环与临时数组生成

• v12 更多采用 NumPy 向量化计算

• 提高矩阵运算效率

7. hop算法优化

• v10 使用：

  • copy

  • cumsum

  • where

• v12 改为在线累加概率

• 减少临时数组分配

8. 更适合后续高性能优化

• v12 数据结构进一步规范化

• 为后续：

  • Numba JIT

  • 并行计算

  • FFT优化

  提供基础

引用

这是PYXAID的原始文献：

• Akimov A V, Prezhdo O V. The PYXAID program for non-adiabatic molecular dynamics in condensed matter systems. Journal of Chemical Theory and Computation, 2013, 9(11): 4959–4972.

• Akimov A V, Prezhdo O V. Advanced capabilities of the PYXAID program: integration schemes, decoherence effects, multiexcitonic states, and field–matter interaction. Journal of Chemical Theory and Computation, 2014, 10(2): 789–804.

本工作：

论文正式发表后补充。