PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ibrahim Elsharkawy, David Curtin, Yonatan Kahn, Neal Weiner

arXiv:2610.27068v1astro-ph.COcs.AIhep-ph

TL;DR

Systematic N-body scans across dark-matter theory space are prohibitively expensive, and existing emulators are limited by resolution or observable specificity. Astrolabe transports COLA-based particle phase space toward N-body distributions with a particle-field architecture and physical conservation constraints. It produces accurate cosmological fields and halos, generalizes to tested out-of-distribution conditions, and runs faster than conventional N-body simulations in the reported comparisons.

  • Problem

    Systematic scans of dark-matter models are computationally prohibitive, while existing emulators are limited by fixed observables or resolution.

    Source: §I · p. 2

  • Method

    Astrolabe uses physical flow matching to transport COLA-based particle phase-space states toward full N-body distributions, with a field-particle hybrid architecture.

    Source: §I · pp. 3–4

  • Results

    Astrolabe produces realistic fields and halos with accurate power spectra, halo mass functions and profiles, and generalizes to untrained dark-matter initial conditions.

    Source: §I · p. 4, §V · p. 23

  • Takeaways & Limitations

    Astrolabe provides a particle-level simulation surrogate that can be fine-tuned across box sizes and particle numbers and offers substantial inference speedups.

    Source: §I · p. 4, §V · pp. 23–24

    The presented Astrolabe version is a first iteration that can almost certainly be made more efficient.

    Source: §I · p. 4

Abstract

from arXiv · show

The microphysics of dark matter (DM) is imprinted on cosmological structure across an enormous dynamic range. Nonlinear cosmological observables are most reliably modeled with N-body and hydrodynamic simulations. However, systematic N-body scans over DM theory space are prohibitively expensive. 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 N 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 computationally tractable via a field-particle hybrid transformer in which every particle cross-attends to multiscale field tokens, so compute and memory are linear in N. Astrolabe incorporates physical flow matching by taking the prior distribution for the flow to be approximate COmoving Lagrangian Acceleration (COLA)-based coordinates rather than unphysical noise, with momentum and mass conservation enforced at all points along flow time. In the Lambda-CDM cases we test, Astrolabe produces realistic fields with accurate power spectra, halo mass functions, velocity profiles, and halo density profiles. We also observe out-of-distribution generalization, with accurate emulation of dark-acoustic oscillation initial conditions. Astrolabe is also capable of being fine-tuned to simulate very large N and we demonstrate this by emulating a TNG50-3-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 a systematic exploration of richer dark sector interactions.

I. INTRODUCTION

N-body simulations provide reliable nonlinear cosmological observables but make systematic dark-matter theory scans computationally prohibitive, while existing emulators trade observable flexibility against resolution. Astrolabe transports particle-level COLA states toward full N-body phase-space distributions, with accurate multi-metric results and generalization beyond its training conditions.

I. INTRODUCTION

  • Systematic dark-matter scans are prohibitively expensive, while existing emulators target a few observables or face resolution costs that grow as O(1/L^3).The paper identifies observable independence and resolution across scales as goals for a simulation surrogate.Source: §I · p. 2
  • Astrolabe is trained on 1,500 DMO simulations and transports COLA point clouds toward full N-body phase-space distributions at a requested redshift.COLA supplies a fast, coarse-grained approximation that reproduces large-scale evolution; the transport recovers small-scale structure.Source: §I · p. 3
  • Astrolabe produces realistic fields and halos, matching power spectra, halo mass functions and profiles, and phase-space consistency; it also generalizes to untrained dark-matter initial conditions.The reported results include matching halo mass functions down to 10^9 M⊙/h and the full truth-level comoving power spectrum at k < 10−2 h/kpc in a 10 Mpc box.Source: §I · p. 4
  • Its point-field hybrid transformer lets individual particles attend to field tokens, retaining particle-level outputs with compute and memory linear in particle number.The paper argues that point-cloud models adapt resolution to structure rather than spending voxel capacity in empty regions.Source: §I · pp. 2, 4
  • Physical flow matching transports phase-space densities while conserving total momentum and mass, and supports temporally coherent movies as the conditioning redshift advances.The flow starts from COLA states and is described as independent of the initial conditions and dark-matter parameters that produced them.Source: §I · p. 4
  • At inference, Astrolabe is 10–500 times faster than other GPU-utilizing N-body simulations in apples-to-apples comparisons.The authors describe this version as a first iteration that can likely be made more efficient.Source: §I · p. 4

II. ASTROLABE

Astrolabe models the particle phase-space density conditioned on redshift and cosmological parameters. The section introduces flow matching as a way to transport samples from a prior toward that target distribution.

II. ASTROLABE

  • Astrolabe's task is to learn and sample the six-dimensional phase-space density of particles conditioned on redshift and cosmological parameters.Each particle's state comprises co-moving position and peculiar velocity.Source: §II · p. 4

A. Review of standard flow matching

  • Flow matching learns a velocity field that transports samples from a prior distribution toward the desired data distribution.Integrating the learned field's ordinary differential equation from a prior sample yields target-distribution samples at flow time 1.Source: A. Review of standard flow matching · p. 5
  • Training constructs conditional paths between prior and training samples, then fits the velocity field using the path's time derivative.The passages describe linear paths with endpoints at the prior sample at time 0 and the training sample at time 1.Source: A. Review of standard flow matching · pp. 5–6

B. Physical flow matching with COLA

Physical flow matching transports COLA phase-space samples to the full N-body distribution while enforcing momentum and mass conservation along the flow. COLA supplies a fast large-scale approximation, while the particle-level transport corrects its small-scale and velocity limitations.

B. Physical flow matching with COLA

  • Physical flow matching transports samples from the approximate COLA phase-space density to the full N-body distribution using a conditional path between them.COLA reproduces large-scale structure but fails to reproduce substructure, small-scale power, and velocity distributions.Source: B. Physical flow matching with COLA · p. 6
  • Using COLA as the prior avoids training two separate generative flows and provides a computationally efficient large-scale transport.The authors report difficulty achieving desired visual fidelity when using N-body initial conditions as the prior.Source: B. Physical flow matching with COLA · p. 6
  • Unlike mesh-based corrections such as COCA [68] and PGD [69], Astrolabe operates directly on particles, allowing sub-mesh corrections and stochastic realizations from one low-mesh COLA prior.The particle-based approach could also accommodate distinct baryonic particle types with minimal method changes beyond training data.Source: B. Physical flow matching with COLA · p. 6
  • Projecting the flow velocity into the zero-net-momentum tangent space conserves total momentum at every integration step, while keeping particle count fixed conserves mass.The method focuses on momentum conservation; for equal-mass particles, the projection removes the mean change in phase-space velocity.Source: B. Physical flow matching with COLA · p. 7

C. Field-Particle Hybrid Architecture

Astrolabe separates particle-local interactions from long-range field context, combining them through a hybrid transformer with linear complexity in particle number. Its particle and multiscale field tokens are conditioned on flow time and cosmological variables to produce per-particle phase-space flow velocities.

C. Field-Particle Hybrid Architecture

  • Splitting particle and field interactions gives Astrolabe linear complexity in particle number, avoiding the prohibitive O(N^2) cost of all-to-all particle attention.The architecture uses particle and field channels, each with linear complexity in N.Source: C. Field-Particle Hybrid Architecture · p. 8
  • Each particle token combines an embedding of its COLA state and initial coordinates with a pooled encoding of its 256 nearest neighbors’ physically motivated pair features.Radial basis features provide a flexible encoding intended to build in the 1/r^2 central-force nature of gravity.Source: C. Field-Particle Hybrid Architecture · p. 8
  • Multiscale field tokens combine convolutional features, cell summary statistics, and cell location and size embeddings across grids with g = 2, 4, 8, 16, and 32.Particles are rasterized into a six-channel field before cells are encoded.Source: C. Field-Particle Hybrid Architecture · pp. 8–10
  • Distance-gated cross-attention keeps coarse field scales global while restricting fine-scale attention to nearby cells, limiting attended tokens per particle.A learned bias based on particle-to-cell geometry supplies an additional positional locality prior.Source: C. Field-Particle Hybrid Architecture · p. 10, Figure 7 · p. 10
  • The generator head maps conditioned particle tokens to six output channels representing each particle’s phase-space flow velocity.Conditioning includes flow time, cosmological variables, and the current-state power spectrum.Source: C. Field-Particle Hybrid Architecture · p. 10

III. TRAINING AND INFERENCE

The section covers how the training set is built, how the model is trained, and how inference is performed at scale.

III. TRAINING AND INFERENCE

  • The section first describes the composition and generation of the training set.This is covered in Section III A.Source: §III · p. 11
  • It then presents the model’s training procedure.This is covered in Section III B.Source: §III · p. 11
  • The final part explains how inference is performed at scale.This is covered in Section III C.Source: §III · p. 11

A. The N-body and COLA dataset

The training dataset pairs GIZMO N-body snapshots with COLA runs sharing the same initial conditions, across 1 and 10 Mpc/h boxes. Its model-size presets specify token and attention dimensions.

A. The N-body and COLA dataset

  • Matched GIZMO and COLA data provide shared-initial-condition training examples across 1 and 10 Mpc/h boxes.The N-body set contains 1,500 runs with 100 snapshots each, totaling roughly 200 billion particles.Source: A. The N-body and COLA dataset · p. 11
  • COLA runs use 40 timesteps and a 256^3 particle mesh, with separate runs for each initial-condition and GIZMO-snapshot pair.This accounts for differing COLA convergence across redshift snapshots.Source: A. The N-body and COLA dataset · p. 11
  • Table I defines model-size presets through token dimension D, body and head block counts, attention heads, and MLP hidden dimension.The supplied table caption identifies these preset fields but gives no preset values.Source: Table I · p. 12

B. Training procedure

Training uses particle subsampling and halo-weighted sampling to make large particle counts practical and give small halos more influence. Inference integrates the learned flow from COLA states, and the results section evaluates ΛCDM, DAO, larger simulations, and compute savings.

B. Training procedure

  • Uniformly subsampling particles gives an unbiased estimator of the full loss gradient and reduces per-step cost from O(N) to O(n).Grid tokens and neighbor candidates are still built from all N particles, while only n ≪ N particles are embedded and supervised.Source: B. Training procedure · p. 12
  • Small-halo weighting counteracts the per-particle mean loss’s dominance by bulk motion by increasing draws for particles ending in low-mass halos.The weighting is training-only and biases gradients toward spending more capacity on small halos.Source: B. Training procedure · p. 12

C. Inference

  • Inference solves the flow-matching ODE from a held-out COLA phase-space distribution, using Euler steps and momentum projection at every step.The desired generated state is at flow time t = 1.Source: C. Inference · p. 12
  • Inference can be split across devices because particle evaluations communicate through field tokens and local neighbors, with per-device compute and memory scaling as O(N/ndevice).This supports probing generalization beyond the particle number used in training, assuming the model generalizes to different particle numbers and resolutions.Source: C. Inference · p. 13

IV. RESULTS

  • The results evaluate unseen ΛCDM conditions, DAO initial conditions, fine-tuning to larger particle counts and box sizes, and compute savings; a figure also examines NFW halo density profiles.The supplied passages describe these evaluation areas but do not provide their quantitative outcomes.Source: §IV · pp. 13–14, Figure 10 · p. 14

A. ΛCDM (in-distribution) evaluation

On ΛCDM test simulations, Astrolabe improves several measures of small-scale structure over COLA, while performance varies by observable and box size.

A. ΛCDM (in-distribution) evaluation

  • The evaluation uses eight initial conditions of 2M-particle ΛCDM simulations at 1 and 10 Mpc/h, testing ten redshifts against COLA and Astrolabe variants.The variants include small and medium models and fast-inference modes.Source: A. ΛCDM (in-distribution) evaluation · p. 14
  • Halo checks quantify NFW concentration, inner density slope, and Jeans consistency, alongside power-spectrum and velocity-dispersion errors.The concentration fit uses halos above the specified particle-mass cutoff; the inner slope is measured at r < 0.05 · R200.Source: A. ΛCDM (in-distribution) evaluation · pp. 15, 17
  • Astrolabe recovers much of the small-scale power lost in COLA, with better power-spectrum performance in the 10 Mpc/h box than in the 1 Mpc/h box.The authors suggest stronger nonlinear dynamics in the smaller box may make COLA harder for Astrolabe to correct.Source: A. ΛCDM (in-distribution) evaluation · p. 16, Figure 12 · p. 16
  • Astrolabe overproduces small halos but recovers a more accurate halo velocity-dispersion distribution than COLA.Figure 12 reports these statistics at z = 0 for both box sizes.Source: Figure 12 · p. 16, A. ΛCDM (in-distribution) evaluation · p. 17
  • Astrolabe improves halo consistency over COLA, bringing cored profiles toward—but not fully to—cuspy profiles and improving NFW concentrations and sub-virial-radius Jeans consistency.The Jeans improvement is reported for virialized dynamics below the virial radius.Source: A. ΛCDM (in-distribution) evaluation · pp. 15–17

B. Out-of-distribution generalization: case study with dark acoustic oscillations

The authors test whether an Astrolabe model trained on ΛCDM generalizes to dark-acoustic-oscillation initial conditions outside its training distribution. It generally improves on COLA, with a reported velocity exception in the 1 Mpc/h box.

B. Out-of-distribution generalization: case study with dark acoustic oscillations

  • The OOD test perturbs initial conditions with dark acoustic oscillations from a dark-radiation bath, using four variants defined through the ETHOS parametrization.The variants vary the dark-radiation-to-CMB temperature ratio and the leading momentum-transfer-rate coefficient.Source: B. Out-of-distribution generalization:… · pp. 17–18
  • Astrolabe significantly improves on COLA across the DAO tests, except for velocity distribution in the 1 Mpc/h box, the study’s only obvious out-of-distribution failure.The tests include four DAO configurations and three initial conditions per configuration.Source: B. Out-of-distribution generalization:… · p. 18, Table III · p. 19
  • Astrolabe better captures the evolution of the DAO-to-ΛCDM power-spectrum ratio across redshifts and improves small-scale halo concentration, inner-slope, and Jeans metrics over COLA.The DAO-to-ΛCDM power-spectrum error is compared with GIZMO in Figure 16.Source: B. Out-of-distribution generalization:… · p. 18, Figure 16 · p. 19
  • A model trained only on ΛCDM produces z = 0 fields for substantially different DAO initial conditions without additional training.The authors describe this as evidence of generalizable features of dark-matter gravitational evolution.Source: B. Out-of-distribution generalization:… · p. 19

C. Fine-tuning for Larger Particle Count

The authors fine-tune Astrolabe-Small to test larger particle counts and different box sizes, then evaluate it on TNG simulations. Fine-tuning improves results, including on TNG100, though resolution and held-out-simulation limits qualify the findings.

C. Fine-tuning for Larger Particle Count

  • The larger-particle-count experiment fine-tunes Astrolabe-Small on early-redshift snapshots, evaluates at z = 0, and compares against COLA and the unfine-tuned model.The 20 Mpc/h case preserves the 2M 10 Mpc/h interparticle spacing, while 1 Mpc/h tests increased resolution.Source: C. Fine-tuning for Larger Particle Count · pp. 19–20
  • For 16M-particle tests, fine-tuning improves held-out z = 0 performance at both box sizes, but benefits unseen initial conditions for the 20 Mpc/h box and only marginally helps at 1 Mpc/h.The authors suggest adapting to box size may be easier than increasing resolution, which may require full retraining.Source: Table IV · p. 20, C. Fine-tuning for Larger Particle Count · p. 22
  • The TNG evaluation fine-tunes on TNG50-3-Dark snapshots at z = 2.0, 1.0, 0.5, and 0.24, then evaluates held-out z = 0 snapshots.TNG100-3-Dark serves as a proxy held-out simulation despite its different particle count.Source: C. Fine-tuning for Larger Particle Count · p. 22
  • Fine-tuned Astrolabe improves significantly on COLA for TNG50-3-Dark and the technically out-of-distribution TNG100-3-Dark; without fine-tuning, Astrolabe performs worse than COLA.Emulating TNG50-3-Dark required evaluation on 150M particles with 10 Euler steps.Source: Table V · p. 21, C. Fine-tuning for Larger Particle Count · p. 22

D. Comparison of compute costs

The compute study compares Astrolabe with GPU-based PKDGRAV3 for generating z = 0 snapshots at matched nonlinear-physics resolution. Fast inference reduces compute further, with reported speedups increasing with particle count.

D. Comparison of compute costs

  • The comparison measures node-hours for z = 0 snapshots while increasing box size with particle number to hold simulation resolution constant.The setup therefore probes the same level of nonlinear physics across particle counts.Source: D. Comparison of compute costs · p. 22, Figure 20 · p. 23
  • Fast inference cuts costs by using fewer Euler steps and evaluating only a subset of particles, copying flow velocities to neighboring particles.The reported settings are five Euler steps and subsampling factor four.Source: D. Comparison of compute costs · p. 23, Figure 20 · p. 23
  • Astrolabe-Small achieves approximately 10–250× speedup over PKDGRAV3, and fast inference achieves approximately 60–500× speedup, depending on particle number.Fast inference uses five Euler steps and a particle-subsampling factor of four.Source: D. Comparison of compute costs · p. 23, Figure 20 · p. 23

V. CONCLUSION AND OUTLOOK

Astrolabe combines particle-level emulation with physical flow matching to produce realistic fields and halos across tested settings at substantially lower compute cost than traditional N-body simulations. The conclusion identifies baryonic emulation and broader dark-matter scans as future directions.

V. CONCLUSION AND OUTLOOK

  • Astrolabe acts on particle phase space rather than voxel summaries, using a field-particle hybrid architecture that scales to large N, as demonstrated on the 150M-particle TNG50-3-dark simulation.The conclusion identifies the approach as supporting scans through dark-matter theory space and other cosmological parameters that traditional simulations make computationally intractable.Source: §V · p. 23
  • Astrolabe produces realistic fields and halos with accurate spectra, velocity dispersions, and halo profiles, while achieving a compute advantage of one or more orders of magnitude over PKDGRAV3.The advantage varies with particle number and inference mode; fine-tuning also generalizes across box sizes and boundary conditions not seen in training.Source: §V · pp. 23–24
  • Future work proposes training an Astrolabe-style emulator for baryonic physics and exploring whether the architecture can support larger dark-sector scans.The authors say comparable costs for baryonic emulation could increase compute savings; they frame improved understanding of dark-matter effects as a possibility.Source: §V · p. 24

INTERACTIVE VISUALIZATION, CODE, AND DATA AVAILABILITY

  • The study's code and dataset are to be released soon, and an online post provides interactive visualizations and films of Astrolabe and N-body fields.The passage provides the blog URL: https://ibrahimelsharkawy.com/blog/astrolabe-flow-matching.html.Source: INTERACTIVE VISUALIZATION, CODE, AND… · p. 24

AI DISCLOSURE

  • Claude assisted with implementation, compute-job management, plot optimization, interactive figures, and manuscript work; these uses were human-verified.The passage also specifies literature searches and draft feedback among the manuscript uses.Source: AI DISCLOSURE · p. 24

Appendix A: Inference Ablation

The appendices describe an inference-speed tradeoff, multiresolution position encoding, model hyperparameter tables, and additional halo measurements across simulation cases.

Appendix A: Inference Ablation

  • Choosing n = 5 and c = 4 increases inference speed with controlled performance loss; full six-dimensional neighbor selection was less successful than position-space selection.The scan varies Euler steps and particle subsampling, assigning unevaluated particles the velocity of their nearest evaluated neighbor in three-dimensional position space; its error is normalized to the n = 10, no-subsampling baseline.Source: Appendix A · p. 25

Appendix B: Hash Encoding of Particle Positions

  • The position encoder concatenates learned features from 12 hash grids spanning resolutions from 32^3 to 2767^3 cells, then maps the 24-dimensional vector through an MLP.The grids are finer than the particle mesh used for training; particle coordinates are encoded across multiple scales because the network can struggle with small-scale correlations.Source: Appendix B · p. 26, Figure 22 · p. 26

Appendix C: Model and Training Hyperparameters

  • Table VI lists hyperparameters that vary between Astrolabe-Small and Astrolabe-Medium, while Table VII lists shared hyperparameters.The supplied passages identify the tables' scopes but do not provide their parameter values.Source: Table VI · p. 26, Appendix C · p. 26, Table VII · p. 27

Appendix D: More Halo Results

  • Appendix D tabulates concentration, inner slope, and Jeans-test values for ΛCDM, DAO, and TNG simulation cases, including different box sizes and DAO strengths.The tables report means and variances over halos meeting the stated particle-count thresholds; the passages do not provide the table cell values.Source: Appendix D · p. 26, Table VIII · p. 27, Table IX · p. 27
Loading 2610.27068v1…