Fast(er)PM and Moving Mesh: JAX-native Geometric Multigrid Methods

Abstract

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.