AI 中文总结
研究旨在构建高效泊松求解器，提出基于 JAX 的几何多重网格框架，用于粒子-网格引力。在固定网格上可替代 FFT，在移动网格上是可行求解器，通过热启动切比雪夫多重网格减少计算量，提升性能，还嵌入可微移动网格方法，成为两者间实用桥梁。
AI 中文摘要
高效、可微的泊松求解器是现代粒子-网格模拟和场级推理管道的关键组成部分。基于 FFT 的求解器在固定笛卡尔网格上非常有效，但存在全局全对全通信且依赖于自适应或非笛卡尔坐标中会丢失的对称性。本文提出了一种用于粒子-网格引力的基于 JAX 的几何多重网格框架，指出多重网格有两个互补作用：在固定网格上可替代 FFT，在移动网格上是可行的求解器。对于静态 FastPM 演化，热启动切比雪夫多重网格作为缺陷校正方法，利用时间步间的时间相干性减少达到场级精度所需的 V 循环次数。在大网格尺寸下，减少了内存压力，与分布式 FFT 相比有相当或更快的时钟性能，在固定最终网格尺寸下总 GPU 时间最多减少一半。然后将相同求解器嵌入可微移动网格粒子-网格方法中。结果表明几何多重网格可成为快速固定网格 PM 方法和可微自适应力宇宙学模拟之间的实用桥梁。
英文摘要
Efficient, differentiable Poisson solvers are a key component of modern particle--mesh simulations and field-level inference pipelines. FFT-based solvers are extremely effective on fixed Cartesian meshes, but they impose global all-to-all communication and rely on symmetries that are lost in adaptive or non-Cartesian coordinates. In this work, we present a JAX-native geometric multigrid framework for particle--mesh gravity and argue that multigrid plays two complementary roles: on fixed meshes it can be a competitive, communication-avoiding alternative to FFTs, while on moving meshes it becomes the enabling solver. For static FastPM evolution, warm-started Chebyshev multigrid acts as a defect-correction method, exploiting temporal coherence between time steps to reduce the number of V-cycles required for field-level accuracy. At large mesh sizes this reduces memory pressure and yields comparable or faster wall-clock performance than distributed FFTs, with up to a factor of two reduction in total GPU time at fixed final mesh size. We then embed the same solver in a differentiable moving-mesh particle--mesh method, where adaptive coordinate deformation produces a variable-coefficient curvilinear Poisson equation that cannot be solved by ordinary FFT diagonalization. The resulting method concentrates force resolution in nonlinear structures while retaining a regular, JAX-compilable, automatically differentiable array workflow. These results suggest geometric multigrid can be a practical bridge between fast fixed-grid PM methods and differentiable adaptive-force cosmological simulations.
Comments19 pages, 16 figures
