PyBird-JAX: Accelerated inference in large-scale structure with model-independent emulation of one-loop galaxy power spectra

The paper introduces $\texttt{PyBird-JAX}$, a differentiable, JAX-based framework that utilizes neural network emulators to accelerate one-loop EFTofLSS galaxy power spectrum calculations by 3–4 orders of magnitude while maintaining high accuracy, thereby enabling rapid and efficient cosmological inference for upcoming large-scale structure surveys.

The Gist

Imagine you are trying to understand the universe by looking at a giant, cosmic map of galaxies. Astronomers call this the "Large-Scale Structure" (LSS). It's like looking at a city from a plane: you can see the streets (galaxy clusters) and the open spaces, but the details are fuzzy because gravity has been pulling things together for billions of years.

To understand how this city formed, scientists use a complex mathematical recipe called the Effective Field Theory of Large-Scale Structure (EFTofLSS). This recipe helps them predict what the galaxy map should look like based on different theories of how the universe works (like how much dark matter there is or how fast it's expanding).

However, there's a big problem: The recipe is incredibly slow.

The Problem: The "Slow Cooker"

Imagine trying to bake a cake, but the oven takes 10 minutes to preheat and 5 hours to bake a single slice. If you want to test 1,000 different recipes to find the perfect one, you'd be waiting for years.

In cosmology, the "oven" is the computer code (called PyBird) that calculates these predictions. To get precise answers for upcoming massive surveys (like DESI or Euclid), scientists need to run this code millions of times. The old version of PyBird was accurate but so slow that it was like trying to run a marathon in concrete boots. It took about half a second to calculate one prediction. Multiply that by millions of attempts, and you're stuck.

The Solution: PyBird-JAX and the "Smart Assistant"

The authors of this paper, Alexander Reeves, Pierre Zhang, and Henry Zheng, built a new tool called PyBird-JAX. Think of it as upgrading from a slow, manual oven to a high-speed, AI-powered microwave.

They did two main things to make it fast:

1. The "JAX" Upgrade (The New Engine):
   They rewrote the code using a special programming language called JAX. Imagine JAX as a super-efficient delivery service. Instead of a single truck (a standard computer processor) delivering one package at a time, JAX sends out a fleet of drones (GPUs) that can deliver thousands of packages simultaneously. It also "pre-cooks" the instructions so the computer doesn't have to think about how to do the math every time; it just does it instantly.

2. The "Emulator" (The Smart Assistant):
   This is the real magic. The slowest part of the recipe involves calculating complex loops (like trying to predict how a crowd of people will move). Instead of calculating every single step from scratch, PyBird-JAX uses a Neural Network Emulator.

   • The Analogy: Imagine a master chef who has cooked a million different meals. Instead of measuring every spice and timing every second for a new dish, the chef looks at the ingredients and instantly guesses how long it will take and how it will taste, based on their massive experience.
   • How it works: The team trained a computer brain (a Neural Network) on millions of examples. They taught it to recognize the "shape" of the universe's data. Now, when the code needs a prediction, the emulator doesn't do the heavy lifting; it just looks up the pattern and gives the answer in a flash.
   • The "Model-Independent" Trick: Most smart assistants need to be retrained if you change the recipe slightly (e.g., if you switch from a chocolate cake to a carrot cake). This new assistant is special. It doesn't care about the specific ingredients (the cosmological model); it only cares about the shape of the data. So, it works for any theory of the universe without needing to be retrained.

The Results: From Hours to Milliseconds

The speed difference is mind-blowing:

• Old PyBird: Takes about 560 milliseconds (over half a second) to calculate one prediction.
• PyBird-JAX (Standard): Takes about 100 milliseconds.
• PyBird-JAX (with Emulator): Takes about 1.2 milliseconds on a normal computer and 0.2 milliseconds on a powerful graphics card.

That is a 3,000 to 4,000 times speed-up. It's like going from walking to the store to teleporting.

Why Does This Matter?

Because it's so fast, scientists can now:

1. Test More Theories: They can try thousands of different versions of the universe's history to see which one matches the real data best.
2. Handle More Data: Upcoming surveys will map millions of galaxies. The old code would choke on this data; the new code handles it like a breeze.
3. Be More Precise: Because the code is so fast, they can use "gradient-based" methods. Imagine trying to find the bottom of a valley in the dark. The old way was to feel around randomly. The new way uses a map that tells you exactly which way is "down" (using math called automatic differentiation), so you find the answer much faster and more accurately.

The Bottom Line

PyBird-JAX is a revolutionary tool that turns a slow, difficult calculation into a lightning-fast, routine task. It allows astronomers to use the most powerful computers (GPUs) to decode the secrets of the universe with unprecedented speed and accuracy, ensuring we are ready for the massive flood of data coming from the next generation of telescopes.

In short: They took a cosmic puzzle that used to take a lifetime to solve and figured out how to solve it in the time it takes to brew a cup of coffee.

Technical Summary

Here is a detailed technical summary of the paper "PyBird-JAX: Accelerated inference in large-scale structure with model-independent emulation of one-loop galaxy power spectra."

1. Problem Statement

The field of cosmology has entered an era of precision, driven by large-scale structure (LSS) surveys like DESI and Euclid. Analyzing these datasets requires modeling the galaxy power spectrum at the one-loop level within the Effective Field Theory of Large-Scale Structure (EFTofLSS).

• Computational Bottleneck: Traditional EFTofLSS calculations involve expensive loop integrals and IR-resummation (infrared resummation) to account for non-linearities and redshift-space distortions. These operations are computationally intensive, taking hundreds of milliseconds per evaluation on a CPU.
• Parameter Space Complexity: Full-shape analyses introduce $\mathcal{O}(10-40)$ nuisance parameters (bias, counterterms, stochastic terms) per redshift bin. Combined with cosmological parameters, this creates high-dimensional parameter spaces ($\sim 100+$ dimensions) that are prohibitively slow to explore with traditional Markov Chain Monte Carlo (MCMC) methods.
• Limitations of Existing Emulators: Previous attempts to speed up inference using neural network (NN) emulators often required retraining for every specific cosmological model or were limited to specific cosmological parameter spaces, lacking true model independence.

2. Methodology

The authors present PyBird-JAX, a complete rewrite of the PyBird codebase using the JAX library, combined with a novel model-independent neural network emulator (PyBird-Emu).

A. JAX Implementation

• Differentiability: The code is fully differentiable, enabling automatic differentiation (AD) for gradient-based optimization and Fisher matrix calculations.
• Just-In-Time (JIT) Compilation: Compiles Python/NumPy code to optimized machine code for CPUs and GPUs.
• Vectorization: Uses vmap to process batches of cosmological parameters in parallel, significantly speeding up ensemble sampling.

B. Model-Independent Emulation Strategy

Instead of emulating the output based on cosmological parameters (e.g., $\Omega_m, h$), the authors emulate the computationally expensive loop integrals and IR-resummation terms based on the input linear matter power spectrum ($P_{lin}$).

1. Input Representation: The linear power spectrum is decomposed into a compact set of 80 spline coefficients on a fixed grid of $k$-knots (log-spaced from $0.0001$to$0.7 , h \text{ Mpc}^{-1}$).
2. Training Data Generation:
   • A "CosmoRef" bank of $10^6$linear power spectra is generated using the Boltzmann solver CLASS, sampling a wide range of cosmological parameters (including$\Lambda$CDM,$w_0w_a$CDM, neutrino mass, etc.).
   • A Gaussian Copula resampling strategy is used to generate additional training data. This decouples 1D marginals from inter-knot correlations, allowing the training set to explore regions slightly outside the original cosmological bank while maintaining physical smoothness.
3. Neural Network Architecture:
   • Separate NNs are trained for the monopole, quadrupole, and hexadecapole of the one-loop power spectrum.
   • Architecture: 3 hidden layers with 512 neurons each, using a specific activation function with trainable hyperparameters.
   • Output compression via Principal Component Analysis (PCA) reduces the output dimensionality to 256 modes, improving training stability and speed.
4. Scope: The emulator handles the loop integrals and IR-resummation. The linear power spectrum and counterterms are computed analytically (as they are fast and trivially dependent on $P_{lin}$), ensuring the emulator only approximates the most expensive $\sim 10\%$ of the signal calculation.

3. Key Contributions

• Speed: Achieves a 3–4 orders of magnitude speed-up over the original PyBird code.
  • CPU: $\sim 1.2$ ms per evaluation (vs. $\sim 560$ ms).
  • GPU: $\sim 0.2$ ms per evaluation.
  • Vectorized GPU: $\sim 0.023$ ms for a batch of 128.
• Model Independence: The emulator is trained on the shape of the power spectrum, not specific cosmological parameters. It works for any cosmology (including non-standard models like Early Dark Energy or Self-Interacting Neutrinos) without retraining, provided the underlying EFTofLSS assumptions hold.
• Differentiable Pipeline: Enables gradient-based sampling (HMC-NUTS), Fisher forecasting, and Taylor expansion of observables, which are impossible with standard non-differentiable codes.
• Integration: Compatible with major Boltzmann solvers (CLASS, CAMB via finite differences) and samplers (emcee, BlackJAX, Nautilus).

4. Results and Validation

The authors validated PyBird-JAX and PyBird-Emu across three distinct scenarios:

A. PT Challenge Simulations (Accuracy)

• Test: Fitted high-resolution N-body simulations (PT Challenge) with a total volume of $566 , (\text{Gpc}/h)^3$.
• Outcome: PyBird-Emu recovered cosmological parameters ($\Omega_m, \sigma_8, h$) with posteriors indistinguishable from the full PyBird calculation (differences $< 0.3\%$).
• Range: Validated up to $k_{max} = 0.3 \, h \text{ Mpc}^{-1}$ for specific "wedge" combinations, demonstrating robustness at scales where EFTofLSS is typically pushed to its limits.

B. BOSS Data (Coverage & Robustness)

• Test: Analyzed SDSS BOSS DR12 Luminous Red Galaxies.
• Scenarios: Tested standard $\Lambda$CDM, $w_0w_a$CDM, Early Dark Energy (EDE), and Self-Interacting Neutrinos (SI$\nu$).
• Outcome: The posteriors from PyBird-Emu were virtually identical to the full PyBird results for all models, including EDE and SI$\nu$ which were not explicitly represented in the training set. This confirms the emulator's ability to generalize to "out-of-distribution" cosmologies.

C. Stage-4 Survey Forecast (Efficiency)

• Test: A mock DESI-like survey with 7 redshift bins and 84 nuisance parameters (total $\sim 100$ parameters).
• Sampler Performance:
  • Vectorized Ensemble Sampling (emcee): Achieved the highest Effective Sample Size per second (ESS/s = 3.4) by leveraging GPU parallelism.
  • Gradient-Based Sampling (HMC-NUTS): Achieved a 90% acceptance rate (vs. 19% for emcee) but was limited by the cost of gradient computation ($\sim 3\times$ likelihood cost).
  • Taylor Expansion: Using AD to Taylor expand observables around a fiducial point provided a fast alternative, running in comparable time to the emulator.
• Convergence: A typical Stage-4 MCMC analysis converges in minutes on a single GPU.

5. Significance

• Enabling Next-Gen Surveys: PyBird-JAX makes it feasible to perform full-shape analyses on upcoming Stage-4 surveys (DESI, Euclid, LSST) which require scanning massive parameter spaces with high precision.
• New Inference Capabilities: The differentiability of the code allows for advanced statistical techniques, such as:
  • Non-flat Volume Measures: Correcting for "volume projection effects" in marginal posteriors (detailed in a companion paper).
  • Gradient-Based Optimization: Rapidly finding Maximum A Posteriori (MAP) estimates.
• Paradigm Shift: It moves LSS analysis from "slow, brute-force sampling" to "fast, differentiable, and vectorized inference," reducing the time for a full cosmological analysis from days/weeks to minutes.
• Open Science: The code is publicly released and designed to be easily integrated into existing pipelines (MontePython, DESI's desilike, Euclid's cloe).

In conclusion, PyBird-JAX represents a major leap forward in computational cosmology, combining the physical rigor of EFTofLSS with the computational efficiency of modern machine learning and differentiable programming to unlock the full potential of upcoming galaxy clustering data.