Fast(er)PM and Moving Mesh: JAX-native Geometric Multigrid Methods
📄 arXiv:2607.10983 · 📥 PDF · 2026-07-13 · astro-ph.CO
Authors: Benjamin Horowitz [arXiv · scholar]
🕰 Orloj analysis
Tento článek představuje JAX-nativní geometrické multigridové metody pro efektivní a diferencovatelné Poissonovy řešiče v simulacích částic-mřížka, zejména pro gravitaci. Metoda slouží jako konkurenční alternativa k FFT na pevných mřížkách a jako nezbytný řešič pro adaptivní pohyblivé mřížky, zlepšující výkon a rozlišení v nelineárních strukturách.
💡 Práce nabízí významné vylepšení výpočetní efektivity a adaptability pro kosmologické simulace díky JAX-nativnímu multigridu a jeho aplikaci na pohyblivé mřížky.
✓ code_available, falsifiable, modest_claims
📄 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.