Menu

Blog astrolabe: A 3D Particle-Level Surrogate for Cosmology Simulation with Physical Flow Matching

October 2026 Paper Code Coming soon! Cosmology Flow Matching Generative Models Simulation Emulation World Models

Paper in collaboration with: Yonatan Kahn, Neal Weiner, David Curtin

by Ibrahim Elsharkawy, last edited October 2026

Loading the simulation…
Drag to orbit · scroll to zoom · shift-drag to pan · drag the bar to compare

A 3D interactive figure to inspect N-body Simulation vs astrolabe. The N-body simulation is on the left of the draggable bar, and astrolabe is on the right. Play runs a pre-designed camera path while z evolves. You can also orbit, zoom and pan yourself at any redshift, and drag the bar across to compare.

The microphysics of dark matter is thought to be imprinted on cosmological structure across an enormous dynamic range. The most reliable forward model of dark sector physics to nonlinear cosmological observables is the N-body or hydrodynamic simulation. At a cost of O(106) CPU-hours per full-resolution run, systematic scans over dark-matter theory space are computationally intractable with N-body simulations. Machine learning emulators provide a possible solution, but are currently limited by either low resolution or only the ability to emulate specific observables. We introduce astrolabe, a flow-matching-based generative emulator designed to be free of both constraints. astrolabe acts directly on the 6N phase-space coordinates of individual particles. With no voxel grid in the output, there is no voxel-resolution ceiling, and in principle any observable can be measured from inference. astrolabe is made computationally tractable by employing a field-particle hybrid transformer in which every particle cross-attends to multiscale field tokens, so compute and memory are linear in particle count N. We further introduce so-called physical flow matching to improve inductive bias. The flow transports approximate COmoving Lagrangian Acceleration (COLA)1 based positions rather than noise, and momentum conservation is exactly enforced at all points along flow time.

Three rows of density maps at z = 2.13, 0.92, 0.32 and 0: the initial condition at the bottom, COLA fields above it, Astrolabe fields on top, with arrows for the COLA evolution and the momentum-conserving flow
Illustration of physical flow matching, with flow time decoupled from redshift. Bottom: Simulation initial conditions. Middle row: COLA fields, evolved under cosmic time from the initial conditions. Top row: astrolabe density field, evolved under flow time from the COLA distributions in the middle row. PDF

In the cases we test, astrolabe produces realistic fields with accurate power spectra, halo mass functions, and velocity profiles for the CDM fields it was trained on. We also observe out-of-distribution generalization, with accurate emulation of dark-acoustic oscillation initial conditions. We further find astrolabe capable of being fine-tuned to simulate large particle numbers, and we demonstrate this by emulating a TNG-50-dark simulation at a fraction of the cost. This method represents a first step toward a general and powerful cosmology simulation surrogate, which may eventually permit the inclusion of baryonic physics and richer dark sector interactions.

Four density maps of the same 10 Mpc/h box at z = 0, each with a zoom on the most massive halo: N-body, COLA, Astrolabe-Small with ten steps, Astrolabe-Small fast inference Enclosed mass against radius for stacked halos in the 1 Mpc/h box: N-body, COLA, Astrolabe-Small and Astrolabe-Medium
Left: Example CDM field for 1283 ≈ 2M particles in a 10 Mpc/h box. Σ is a surface density in units of particles per pixel. Top left: A GIZMO2 N-body simulation, which serves as the truth-level phase space density for this example. Top right: Coarse approximation to the truth-level density given by COLA with a mesh of 2563 and 40 timesteps1. Bottom left: astrolabe-Small evaluated with ten integration steps and no subsampling. Bottom right: astrolabe-Small Fast Inference evaluated with 2 integration steps and c = 8 subsampling. Right: Consistency test for halo density profiles in the 1 Mpc/h box. Halo mass profile averaged over all halos with more than 5000 particles, with concentration c from an NFW fit and inner slope γ displayed. PDF (left) PDF (right)

Movies of N-body Simulation vs astrolabe for various test cases.

2M CDM

  • CDM, 10 Mpc/h: N-body simulation vs Astrolabe-Small

  • CDM, 1 Mpc/h: N-body simulation vs Astrolabe-Small

2M DAO

  • DAO weak, 10 Mpc/h: kpeak = 31 h/Mpc; N-body simulation vs Astrolabe-Small

  • DAO strong, 10 Mpc/h: kpeak = 18 h/Mpc; N-body simulation vs Astrolabe-Small

  • DAO weak, 1 Mpc/h: kpeak = 95 h/Mpc; N-body simulation vs Astrolabe-Small

  • DAO strong, 1 Mpc/h: kpeak = 32 h/Mpc; N-body simulation vs Astrolabe-Small

TNG-50/100-Dark-33

  • TNG50-3-Dark, z = 0: N-body simulation vs Astrolabe-Small, fine-tuned on TNG50

  • TNG100-3-Dark, z = 0: N-body simulation vs Astrolabe-Small, fine-tuned on TNG50

References

  1. S. Tassev, M. Zaldarriaga and D. Eisenstein, Solving Large Scale Structure in Ten Easy Steps with COLA, JCAP 06 (2013) 036. arXiv:1301.0322
  2. P. F. Hopkins, A new class of accurate, mesh-free hydrodynamic simulation methods, Mon. Not. Roy. Astron. Soc. 450 (2015) 53–110. arXiv:1409.7395
  3. D. Nelson et al., The IllustrisTNG simulations: public data release, Comput. Astrophys. Cosmol. 6 (2019) 2. arXiv:1812.05609
← All posts