从天文到星际：具有能量守恒自引力的高阶可微（磁）流体动力学
Per Astronomix ad Astra: High-Order Differentiable (Magneto)hydrodynamics with Energy-Conserving Self-Gravity
浏览论文内容
中文总结 AI 辅助
介绍用Python/JAX编写的可微（磁）流体动力学模拟器astronomix，利用自动微分实现逆建模等，在单GPU运行时表现良好且可扩展，还提出新自引力方案及相关智能技术提升性能。
中文摘要 AI 辅助
我们展示了astronomix，一个用Python/JAX编写的高性能可微（磁）流体动力学模拟器。通过与手动推导的解析函数导数和有限差分验证自动微分，它能对数百万参数进行逆建模，实现灵敏度和稳定性分析以及正确的本征模式初始化，还能在模拟器内训练机器学习模型。在给定分辨率下单GPU运行时与GPU优化代码AthenaPK相当，但因高阶在平滑问题上误差更低。它可扩展到多个GPU和节点。我们还提出了一种新的四阶自引力方案。为最大化性能，创建了一种能从JAX参考代码和测试套件生成并验证自定义Pallas GPU内核的智能技术。模拟器可通过链接获取。
英文摘要
We present astronomix, a performant differentiable (magneto)hydrodynamics simulator written in Python/JAX. We demonstrate how automatic differentiation, validated against hand-derived analytical functional derivatives and finite differences, enables inverse modeling over millions of parameters and allows for sensitivity and stability analysis as well as correct eigenmode initialization. The differentiability of astronomix furthermore enables training machine-learning models inside the simulator. On a single GPU at a given resolution, astronomix has runtimes of the same order of magnitude as the GPU-optimized code AthenaPK but reaches far lower errors on smooth problems due to its higher order. astronomix scales to multiple GPUs ($\sim 6.5$ strong scaling speedup on $8$ GPUs) and multiple nodes ($\sim 76\%$ weak scaling efficiency on $16$ GPUs over $4$ nodes). We also present a novel fourth-order self-gravity scheme which complements the fifth-order finite difference constrained transport magnetohydrodynamics scheme implemented in astronomix. To maximize performance, we created an agentic skill that generates and validates custom Pallas GPU kernels from our JAX reference code and test suite. The simulator is available at https://github.com/leo1200/astronomix.
