🔭 astrophysics

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

This paper introduces a JAX-native geometric multigrid framework that serves as a memory-efficient, communication-avoiding alternative to FFTs for fixed-mesh particle simulations and acts as the essential differentiable solver for moving-mesh cosmological simulations, thereby bridging the gap between fast fixed-grid methods and adaptive-force field-level inference.

Original authors: Benjamin Horowitz

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin Horowitz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the universe as a giant, invisible web of gravity, stretching across billions of light-years. To understand how this web forms, scientists run massive computer simulations. These simulations are like a game of "connect the dots," but instead of dots, they are billions of particles representing dark matter, and the "dots" are connected by the force of gravity.

The biggest headache in these simulations is solving the Poisson equation. Think of this as the rulebook that tells every particle how to pull on every other particle. It's the math that turns a map of where matter is (the density) into a map of how hard it's pulling (the gravity).

For decades, the standard way to solve this was using a tool called FFT (Fast Fourier Transform). Imagine the FFT as a super-fast, global translator. It takes the whole map, translates it into a secret code, does a quick math trick, and translates it back. It's incredibly fast on a flat, unchanging grid. But there's a catch: to do its magic, the FFT needs every single computer in the cluster to shout its data to every other computer at the same time. It's like a classroom where every student has to pass a note to every other student simultaneously. On a small scale, it's fine. But as the simulation gets bigger and the computers get more numerous, this "all-to-all" shouting match slows everything down and eats up a ton of memory.

The Paper's Main Discovery: A Smarter, Local Approach

Benjamin Horowitz's paper suggests a different way to play the game, using a method called Geometric Multigrid. Instead of a global translator, imagine a team of local editors working on a newspaper.

The "Warm-Start" Trick: In a time-stepping simulation, the universe doesn't change wildly from one second to the next. The gravity map from yesterday is almost the same as today. The paper shows that if you use yesterday's map as a "warm start" (a head start) for today's calculation, you don't need to do the heavy lifting from scratch.
The Multigrid Editor: The multigrid method works like a team of editors with different magnifying glasses. One editor looks at the whole page to fix the big, blurry smudges (long-range gravity). Another zooms in to fix the tiny, jagged edges (short-range gravity). They pass the corrections back and forth.
The Result: On fixed grids, this local approach is faster and uses less memory than the global FFT translator. In tests on NVIDIA A100 GPUs, the new method was up to 2.4 times faster than the old FFT method for large grids (like a 1024³ mesh). It also allowed simulations to run on half the number of computers (nodes) because it didn't need to store massive temporary copies of the data.

The Big Leap: The Moving Mesh

The real magic happens when the paper introduces a Moving Mesh.

Imagine the standard simulation grid as a rigid, square fishing net. If you try to catch a school of tiny fish (dense clusters of matter) in a vast ocean (empty space), most of the net is wasted on empty water. You either need a huge net to catch the fish, or you miss the details.

The paper proposes a quasi-Lagrangian moving mesh. This is a net that can stretch and shrink.

Where the fish are dense (galaxies and clusters), the net contracts, squeezing the cells together to get a super-sharp, high-resolution view.
Where the water is empty (voids), the net expands, wasting no effort on empty space.

This is like having a camera that automatically zooms in on the action and zooms out on the background. The paper shows that this method recovers much more detail in the dense regions than a static net of the same size. For example, in a 256³ particle simulation, the moving mesh captured small-scale structures that a static mesh of the same size missed.

The "Differentiable" Superpower

Here is the most playful part: The entire system is built using JAX, a programming framework that allows the computer to not just calculate the answer, but to learn from it.

Usually, if you want to figure out what the universe looked like in the past based on what it looks like now, you have to guess and check. But because this new solver is "differentiable," you can run the simulation backward. You can tell the computer, "Make the past look like this," and it calculates exactly how to tweak the initial conditions to get that result.

The paper demonstrates this by reconstructing a cosmic web. The moving mesh successfully "learned" to concentrate its resolution exactly where the gravity was strongest, reproducing the complex, web-like structures of the universe with high fidelity.

What the Paper Rules Out and Limits

It's important to know what this method doesn't do, or what it argues against:

It's not a magic bullet for everything: The paper explicitly states that for very small grids, the old FFT method is still faster. The new method only wins when the grid gets large (like 512³ or bigger) and the "shouting" of the FFT becomes too expensive.
It's not a fully adaptive mesh in the traditional sense: Unlike some other methods that can break a grid into tiny, irregular shards (like Voronoi cells), this method keeps the grid as a regular, stretchable rectangle. It doesn't tear the net; it just stretches it. This means it can't handle extreme distortions where the net would fold over itself, so the authors had to add "limiters" to stop the cells from getting too squished.
It's a simulation, not a physical discovery: The results are based on computer simulations (using the CAMELS suite and AREPO code as references). The paper suggests this is a practical bridge for future simulations, but it doesn't claim to have solved the mysteries of the universe itself.

How Sure Are They?

The authors are very confident in the performance numbers. They measured the time it took to run steps on real hardware (Perlmutter supercomputer with A100 GPUs) and found consistent speedups of 1.3 to 2.4 times depending on the grid size and number of computers.

They are also confident in the accuracy. The moving mesh results matched the high-resolution reference simulations very closely, with cross-correlation coefficients around 0.97 (where 1.0 is a perfect match).

However, they are more cautious about the future applications. They suggest that this method could be a game-changer for field-level inference (figuring out the universe's history) and for running massive simulations that are currently too big for standard computers. But they admit that for a real-world scientific study, more work is needed to tune the "limiters" and add noise models.

In a Nutshell

This paper suggests that by switching from a "global shout" (FFT) to a "local whisper" (Multigrid) and letting the grid stretch like a living net, we can run bigger, faster, and more detailed simulations of the universe. It's a way to get high-definition gravity in the places that matter most, without needing a supercomputer the size of a city. And the best part? The computer can now learn from the simulation, potentially helping us reverse-engineer the history of the cosmos.

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

Problem Statement
Modern particle–mesh (PM) simulations and field-level inference pipelines rely heavily on efficient, differentiable Poisson solvers. The current standard, Fast Fourier Transform (FFT)-based solvers, excels on fixed Cartesian meshes but imposes global all-to-all communication patterns that are inefficient on modern GPU architectures and incompatible with adaptive or non-Cartesian coordinates. Furthermore, FFTs cannot solve the variable-coefficient Poisson equations required for moving-mesh methods, which offer adaptive force resolution but break the spectral diagonalization property. There is a need for solvers that are not only fast and scalable on distributed GPUs but also flexible enough to support differentiable, adaptive-force cosmological simulations.

Methodology
The authors present a JAX-native geometric multigrid framework designed to address these limitations in two distinct regimes:

Static Mesh (Fast(er)PM):
- Warm-Start Strategy: Instead of solving the Poisson equation from a "cold" start (zero initial guess) at every time step, the solver utilizes temporal coherence. The potential from the previous time step serves as an initial guess (warm start), reducing the residual error significantly.
- Chebyshev Smoothing: The authors employ a Chebyshev-accelerated Jacobi smoother within the multigrid V-cycle. This polynomial smoother targets the high-frequency error spectrum more effectively than standard damped-Jacobi iterations without altering the local communication pattern.
- Defect Correction: The multigrid solver acts as a defect-correction method, rapidly converging the small residual left by the warm-started guess.
- Implementation Optimizations: The code utilizes jaxdecomp for distributed halo exchanges, whole-program compilation to fuse kernels, and communication-avoiding strategies (e.g., fusing restriction/prolongation with smoothing steps) to minimize memory footprint and latency.
Moving Mesh (Differentiable PM):
- Coordinate Deformation: The method introduces a scalar deformation potential $\psi$ that maps a regular computational grid to a physical mesh that deforms with the matter distribution (quasi-Lagrangian). Cells contract in overdense regions and expand in voids.
- Curvilinear Poisson Solver: This deformation induces a spatially varying metric, transforming the Poisson equation into a variable-coefficient Laplace–Beltrami equation. FFTs cannot diagonalize this operator; thus, the geometric multigrid method is extended to handle the curvilinear operator on the deformed grid.
- Differentiability: The entire workflow, including the mesh deformation, curvilinear solves, and particle updates, is constructed using JAX primitives. This allows for reverse-mode automatic differentiation, enabling gradient-based field-level inference even with the adaptive mesh.
- Limiters: To prevent mesh folding during strong collapse, the authors implement limiters (clips and threshold switches) that are treated via a "straight-through" estimator in the backward pass to maintain differentiability while ensuring numerical stability.

Key Results

Static Mesh Performance: On NVIDIA A100 GPUs, the warm-started Chebyshev multigrid solver outperforms distributed FFTs for large mesh sizes (e.g., $1024^3$ $102 4^{3}$ and $2048^3$ $204 8^{3}$ ).
- Speed: For a $1024^3$ mesh, the one-cycle warm multigrid solve is 1.5–1.6 $\times$ faster than FFT at 8–32 GPUs and up to 2.4 $\times$ faster at 128 GPUs, where FFTs become latency-bound due to all-to-all communication.
- Memory: The multigrid approach avoids the complex transposed scratch arrays required by FFTs. Consequently, it fits on roughly half the number of GPUs required by the spectral path for the same mesh size (e.g., a $2048^3$ mesh fits on 32 GPUs with multigrid but requires 64 GPUs for FFT).
- Accuracy: With two warm cycles, the multigrid solution achieves field-level accuracy comparable to the FFT reference (power spectrum transfer function $P/P_{FFT} \approx 1.001$ ), with errors localized to high wavenumbers in collapsed regions.
Moving Mesh Performance:
- Accuracy: Compared to static PM simulations with the same nominal cell count, the moving-mesh method recovers significantly more small-scale correlation and improves the high-density tail of the mass distribution, better matching high-resolution $N$ -body references (e.g., CAMELS/CV0).
- Differentiability: The authors demonstrate a successful end-to-end differentiable reconstruction experiment. The method can reconstruct initial conditions from a target $z=0$ density field, with the mesh adapting to concentrate force resolution in nonlinear structures. The moving-mesh reconstruction retains higher transfer function values ( $T \approx 0.95$ ) at the Nyquist frequency compared to static meshes ( $T \approx 0.78$ ).

Significance and Claims
The paper argues that geometric multigrid serves as a practical bridge between fast, fixed-grid PM methods and differentiable, adaptive-force cosmological simulations.

For Fixed Grids: The authors claim that multigrid is a competitive, communication-avoiding alternative to FFTs on modern GPU architectures, particularly for large meshes where global transposes become a bottleneck.
For Moving Grids: The framework enables adaptive force resolution without abandoning the regular array operations and accelerator-friendly programming model of PM methods. By embedding the solver in a differentiable workflow, it allows for the optimization of coordinate deformations for specific scientific observables.
Scope: The authors position the moving-mesh method as an intermediate solution between cheap fixed-grid PM and more expensive TreePM or fully unstructured AMR methods. They acknowledge that the current implementation relies on numerical limiters that require systematic calibration and that the reconstruction experiment is a proof-of-concept rather than a full cosmological inference pipeline.

The work suggests that GPU-native geometric multigrid can effectively complement or replace FFT-based solvers in regimes where memory constraints, communication overhead, or the need for adaptive, differentiable force resolution are paramount.

More like this