置换不变多项式的张量化高效求值用于表示势能面

中文摘要 AI 辅助

置换不变多项式（PIPs）连同相关多项式基不变描述符（如基本不变量（FIs））被广泛用于构建含有相同原子的分子的高保真全局势能面（PESs）。尽管单项式对称化方法（MSA）通过递归因式分解实现了PIPs的快速求值，但它导致深度嵌套的计算图，可能占用大量内存，并且在现代自动微分框架中效率低下。在本工作中，我们引入了JaxPIP，一种基于JAX的实现，将PIP/FI求值重新表述为张量化线性代数运算。通过用稠密矩阵运算结合对数-指数变换和分段求和替代递归因式分解，求值变得规则且GPU友好。这使得通过即时编译实现高效执行，并支持能量和力（以及高阶导数）的大规模批量求值。所得到的架构以完全向量化的方式支持系综模拟，如准经典轨迹（QCT）、路径积分分子动力学（PIMD）和扩散蒙特卡洛（DMC）。如示例所示，JaxPIP为具有可扩展GPU执行的分子系统高效模拟提供了一条实用途径。

英文摘要

Permutation invariant polynomials (PIPs), together with related polynomial-based invariant descriptors such as fundamental invariants (FIs), are widely used in constructing high-fidelity global potential energy surfaces (PESs) of molecules that contain identical atoms. Although the monomial symmetrization approach (MSA) enables fast evaluation of PIPs through recursive factorization, it leads to deeply nested computational graphs that may require large memory and are inefficient in modern automatic differentiation frameworks. In this work, we introduce JaxPIP, a JAX-based implementation that reformulates PIP/FI evaluation into tensorized linear algebra operations. By replacing recursive factorization with dense matrix operations combined with log-exp transformation and segmented summation, the evaluation becomes regular and GPU-friendly. This allows efficient execution with just-in-time compilation and enables large-scale batch evaluation of energies and forces (as well as higher-order derivatives). The resulting architecture supports ensemble simulations such as quasi-classical trajectory (QCT), path-integral molecular dynamics (PIMD), and diffusion Monte Carlo (DMC) in a fully vectorized manner. As demonstrated in examples, JaxPIP provides a practical route for efficient simulations of molecular systems with scalable GPU execution.

发表机构

• Chongqing University（重庆大学）
• University of New Mexico（新墨西哥大学）

机构由 AI 辅助整理，请以论文原文为准。