加速大规模常微分方程系综
Accelerating massive ensembles of ordinary differential equations
- University of Cambridge（剑桥大学）
- Kavli Institute for Cosmology University of Cambridge（卡弗里宇宙学研究所 剑桥大学）
- University of Vienna（维也纳大学）
- Ludwig-Maximilians-Universität München（慕尼黑大学）
机构由 AI 辅助整理，请以论文原文为准。
中文总结 AI 辅助
针对大规模独立低维ODE系综，提出GPU加速求解器库modax，通过线程级并行和稀疏技术，在Lorenz和Robertson系统上分别实现30倍和700倍加速，并成功应用于宇宙学模型。
中文摘要 AI 辅助
求解大量独立的小型常微分方程（ODE）是计算科学与工程中的常见任务，例如在贝叶斯参数估计、蒙特卡洛不确定性量化以及非耦合物理系统（如宇宙学的线性化爱因斯坦-玻尔兹曼方程）的积分中。现代图形处理单元（GPU）具有数千个核心，非常适合此类并行工作负载，但实现这一潜力需要仔细关注GPU特定的架构约束。我们提出了modax，一个GPU加速的ODE求解器Python库，专门针对求解大量低维独立问题进行了优化。我们的求解器与JAX生态系统兼容，但采用低级、基于线程的编程范式编写，比纯JAX对应物更适合高度发散轨迹的系综。在Lorenz、Robertson和van der Pol格点系统上的基准测试将modax与diffrax、this http URL和torchdiffeq在系综大小、维度和轨迹发散性方面进行了比较。我们的显式Tsit5求解器在Lorenz系统上比diffrax实现了30倍加速，而我们的隐式Rodas5P求解器在Robertson系统上实现了超过700倍的加速，证明了线性隐式Rosenbrock-Wanner方法在GPU上的优越性。由于使用了利用稀疏性的技术和数据结构，我们的隐式求解器能很好地扩展到更高维度。最后，我们将modax应用于宇宙学领域的三个示例——从原始丰度进行贝叶斯参数估计、全局21cm信号的蒙特卡洛不确定性量化以及原始功率谱中非耦合傅里叶模式的计算——并展示了在建模由低维（<200D）ODE描述的物理现象方面显著提高了效率。
英文摘要
Solving large ensembles of small, independent ordinary differential equations (ODEs) is a common task in computational science and engineering, arising for example in Bayesian parameter estimation, Monte Carlo uncertainty quantification, and the integration of uncoupled physical systems such as the linearised Einstein-Boltzmann equations of cosmology. Modern graphics processing units (GPUs) with thousands of cores are well suited to such parallel workloads, but realising this potential requires careful attention to GPU-specific architectural constraints. We present modax, a Python library of GPU-accelerated ODE solvers specifically optimised for solving large ensembles of low-dimensional, independent problems. Our solvers are compatible with the JAX ecosystem but are written in a low-level, thread-based programming paradigm better suited to ensembles of highly divergent trajectories than their pure-JAX counterparts. Benchmarks on the Lorenz, Robertson and van der Pol lattice systems compare modax against diffrax, DiffEqGPU$.$jl and torchdiffeq across ensemble size, dimensionality, and trajectory divergence. Our explicit Tsit5 solver achieves 30x speed-up over diffrax on Lorenz systems, and our implicit Rodas5P solver achieves over 700x speed-up on Robertson systems, demonstrating the superiority of linearly implicit Rosenbrock-Wanner methods on GPUs. Our implicit solver scales well to higher dimensions thanks to the use of sparsity-exploiting techniques and data structures. We conclude by applying modax to three examples from the field of cosmology - Bayesian parameter estimation from primordial abundances, Monte Carlo uncertainty quantification of the global 21cm signal, and the computation of uncoupled Fourier modes in the primordial power spectrum - and demonstrate significant efficiency improvements in our ability to model physical phenomena described by low-dimensional (<200D) ODEs.
