Field-level inference of primordial non-Gaussianity from galaxy redshift surveys
Hugo SIMON,
PhD student supervised by
Arnaud DE MATTIA and François LANUSSE
CSI, 2025/09/02
Cosmological inference
\(\Omega := \{ \Omega_m, \Omega_\Lambda, H_0, \sigma_8, f_\mathrm{NL},...\}\)
inference
\(P\)
\(\Omega\)
\(\delta_L\)
\(\delta_g\)
Cosmological inference
\(\Omega := \{ \Omega_m, \Omega_\Lambda, H_0, \sigma_8, f_\mathrm{NL},...\}\)
\(\Omega\)
\(\delta_L\)
\(\delta_g\)
inference
Please define model
Model: an observation generator, link between latent and observed variables
\(F\)
\(C\)
\(D\)
\(E\)
\(A\)
\(B\)
\(G\)
Frequentist:
fixed but unknown parameters
Bayesian:
random because unknown
\(A\)
\(B\)
\(C\)
\(F\)
\(D\)
\(E\)
\(G\)
Represented as a DAG which expresses a joint probability factorization, i.e. conditional dependencies between variables
latent param
latent variable
observed variable
Definitions
evidence/prior predictive: $$\mathrm{p}(y) := \int \mathrm{p}(y \mid x) \mathrm{p}(x)\mathrm d x$$
posterior predictive: $$\mathrm{p}(\tilde y \mid y) := \int \mathrm{p}(\tilde y \mid x) \mathrm{p}(x \mid y)\mathrm d x$$
\(y\)
\(x\)
Bayes theorem
\(\sim\)
\(\underbrace{\mathrm{p}(x \mid y)}_{\text{posterior}}\underbrace{\mathrm{p}(y)}_{\text{evidence}} = \underbrace{\mathrm{p}(y \mid x)}_{\text{likelihood}} \,\underbrace{\mathrm{p}(x)}_{\text{prior}}\)
\(y\)
\(x\)
Marginalization:
\(y\)
\(x\)
\(z\)
\(z\)
\(x\)
\(y\)
\(x\)
\(y\)
\(x\)
Conditioning:
Probabilistic Programming: Wish List
\(\begin{cases}x,y,z \text{ samples}\\\mathrm{p}(x,y,z)\\ U= -\log \mathrm{p}\\\nabla U,\nabla^2 U,\dots\end{cases}\)
1. Modeling
e.g. in JAX with \(\texttt{NumPyro}\), \(\texttt{JaxCosmo}\), \(\texttt{JaxPM}\), \(\texttt{DISCO-DJ}\)...
2. DAG compilation
PPLs e.g. \(\texttt{NumPyro}\), \(\texttt{Stan}\), \(\texttt{PyMC}\)...
3. Inference
e.g. MCMC or VI with \(\texttt{NumPyro}\), \(\texttt{BlackJax}\), \(\texttt{Stan}\), \(\texttt{PyMC}\), \(\texttt{emcee}\)...
4. Extract and viz
e.g. KDE with \(\texttt{GetDist}\), \(\texttt{GMM-MI}\), \(\texttt{mean()}\), \(\texttt{std()}\)...
\(x\)
\(z\)
\(y\)
\(x\)
\(y\)
1
2
3
4
Field-level inference
Summary stat inference
\(\Omega\)
\(s\)
\(\delta_g\)
\(\Omega\)
\(\delta_L\)
\(s\)
marginalize
condition
marginalize
\(\Omega\)
\(s\)
\(\delta_g\)
\(\Omega\)
\(\delta_L\)
condition
Two approaches to cosmological inference
Cosmo model
\(\mathrm{p}(\Omega,s)\)
\(\mathrm{p}(\Omega \mid s)\)
\(\Omega\)
\(\delta_g\)
\(\mathrm{p}(\Omega,\delta_L,\delta_g, s)= \mathrm{p}(s \mid \delta_g) \, \mathrm{p}(\delta_g \mid \Omega,\delta_L)\, \mathrm{p}(\delta_L \mid \Omega)\, \mathrm{p}(\Omega)\)
\(\mathrm{p}(\Omega,\delta_L \mid \delta_g)\)
\(\mathrm{p}(\Omega \mid \delta_g)\)
\(\delta_g\)
\(\Omega\)
\(\delta_L\)
\(s\)
Two approaches to cosmological inference
Cosmo model
Problem:
- \(s\) is too simple \(\implies\) lossy compression
- \(s\) is too complex \(\implies\) intractable marginalization
The Problem:
- high-dimensional integral $$\mathrm{p}(\Omega \mid \delta_g) = \int \mathrm{p}(\Omega, \delta_L \mid \delta_g) \;\mathrm d \delta_L$$
- To probe scales of \(15\ \mathrm{Mpc}/h\) in DESI volume, \(\operatorname{dim}(\delta_L) \simeq 1024^3\)
The Promise:
- "lossless" explicit inference
Field-level inference
Summary stat inference
We gotta pump this information up
- Field-level
- CNN...
- WST...
- Halo, Peak, Void, Split, Hole...
- 3PCF, Bispectrum
- 2PCF, Power spectrum
- 1D-PDF
$$0-$$
$$\mathrm H(\delta_g)-$$
- At large scales, matter density field almost Gaussian so power spectrum is almost lossless compression, and is relatively tractable
- To prospect smaller non-Gaussian scales, let's use:
We gotta pump this information up
- Field-level
- CNN, GNN...
- WST, 1D-PDFs, Holes...
- Peak, Void, Split, Cluster...
- 3PCF, Bispectrum
- 2PCF, Power spectrum
$$0-$$
$$\mathrm H(\delta_g)-$$
- At large scales, matter density field is Gaussian so power spectrum is lossless compression, and is relatively tractable
- To prospect smaller non-Gaussian scales, let's add:
- all the data
- learn the stat
- multiscale count
- object correlations
- more correlations
- standard analysis
-
Prior on
- Cosmology \(\Omega\)
- Initial field \(\delta_L\)
- EFT Lagrangian parameters \(b\)
(Dark matter-galaxy connection)
-
LSS formation: 2LPT or PM
(BullFrog or FastPM) - Apply galaxy bias
- Redshift-Space Distortions
- Observational noise
Field-Level Modeling
Fast and differentiable model thanks to (\(\texttt{NumPyro}\) and \(\texttt{JaxPM}\))
How to N-body-differentiate?
\((\boldsymbol q, \boldsymbol p)\)
\(\delta(\boldsymbol x)\)
\(\delta(\boldsymbol k)\)
paint*
read*
fft*
ifft*
fft*
*: differentiable, e.g. with via \(\texttt{JaxPM}\), in \(\mathcal O(n \log n)\)
apply forces
to move particles
solve Vlasov-Poisson
to compute forces
\(\begin{cases}\dot {\boldsymbol q} \propto \boldsymbol p\\ \dot{\boldsymbol p} = \boldsymbol f \end{cases}\)
\(\begin{cases}\nabla^2 \phi \propto \delta\\ \boldsymbol f = -\nabla \phi \end{cases} \implies \boldsymbol f \propto \frac{i\boldsymbol k}{k^2} \delta\)
MCMC sampling
Canonical MCMC samplers
Recipe😋 to sample from \(\mathrm p \propto e^{-U}\)
- take particle with position \(\boldsymbol q\), momentum \(\boldsymbol p\), mass matrix \(M\), and Hamiltonian $$\mathcal H(\boldsymbol q, \boldsymbol p) = U(\boldsymbol q) + \frac 1 2 \boldsymbol p^\top M^{-1} \boldsymbol p$$
- follow Hamiltonian dynamics during time \(L\)
$$\begin{cases} \dot {{\boldsymbol q}} = \partial_{\boldsymbol p}\mathcal H = M^{-1}{{\boldsymbol p}}\\ \dot {{\boldsymbol p}} = -\partial_{\boldsymbol q}\mathcal H = - \nabla U(\boldsymbol q) \end{cases}$$and refresh momentum \(\boldsymbol p \sim \mathcal N(\boldsymbol 0,M)\)
- usually, perform Metropolis adjustment
- this samples canonical ensemble $$\mathrm p_\text{C}(\boldsymbol q, \boldsymbol p) \propto e^{-\mathcal H(\boldsymbol q, \boldsymbol p)} \propto \mathrm p(\boldsymbol q)\,\mathcal N(\boldsymbol 0, M)$$
gradient guides particle toward high density sets
scales poorly with dimension
must average over all energy levels
Hamiltonian Monte Carlo (e.g. Neal2011)
Microcanonical MCMC samplers
Recipe😋 to sample from \(\mathrm p \propto e^{-U}\)
-
take particle with position \(\boldsymbol q\), momentum \(\boldsymbol p\), mass matrix \(M\), and Hamiltonian $$\mathcal H(\boldsymbol q, \boldsymbol p) = \frac {\boldsymbol p^\top M^{-1} \boldsymbol p} {2 m(\boldsymbol q)} - \frac{m(\boldsymbol q)}{2} \quad ; \quad m=e^{-U/(d-1)}$$
-
follow Hamiltonian dynamics during time \(L\)
$$\begin{cases} \dot{\boldsymbol q} = M^{-1/2} \boldsymbol u\\ \dot{\boldsymbol u} = -(I - \boldsymbol u \boldsymbol u^\top) M^{-1/2} \nabla U(\boldsymbol q) / (d-1) \end{cases}$$ and refresh \(\boldsymbol u \leftarrow \boldsymbol z/ \lvert \boldsymbol z \rvert \quad ; \quad \boldsymbol z \sim \mathcal N(\boldsymbol 0,I)\)
-
usually, perform Metropolis adjustment
- this samples microcanonical/isokinetic ensemble $$\mathrm p_\text{MC}(\boldsymbol q, \boldsymbol u) \propto \delta(H(\boldsymbol q, \boldsymbol u)) \propto \mathrm p (\boldsymbol q) \delta(|\boldsymbol u|^2 - 1)$$
single energy/speed level
let's try avoiding that
gradient guides particle toward high density sets
MicroCanonical HMC (Robnik+2022)
Canonical/Microcanonical MCMC samplers
- \(\mathcal H(\boldsymbol q, \boldsymbol p) = \frac {\boldsymbol p^\top M^{-1} \boldsymbol p} {2 m(\boldsymbol q)} - \frac{m(\boldsymbol q)}{2} \quad ; \quad m:=e^{-U/(d-1)}\)
- samples microcanonical/isokinetic ensemble $$\mathrm p_\text{MC}(\boldsymbol q, \boldsymbol p) \propto \delta(\mathcal H(\boldsymbol q, \boldsymbol p)) \propto \mathrm p (\boldsymbol q) \delta(\dot{\boldsymbol q}^\top M \dot{\boldsymbol q} - 1)$$
Hamiltonian Monte Carlo (e.g. Neal2011)
MicroCanonical HMC (Robnik+2022)
- \(\mathcal H(\boldsymbol q, \boldsymbol p) = U(\boldsymbol q) + \frac 1 2 \boldsymbol p^\top M^{-1} \boldsymbol p\)
- samples canonical ensemble $$\mathrm p_\text{C}(\boldsymbol q, \boldsymbol p) \propto e^{-\mathcal H(\boldsymbol q, \boldsymbol p)} \propto \mathrm p(\boldsymbol q)\,\mathcal N(\boldsymbol 0, M)$$
What it looks like
Inferring jointly cosmology, bias parameters, and initial matter field allows full universe history reconstruction
- Different samplers and strategies used for field-level (e.g. Lavaux+2018, Bayer+2023). Additional comparisons required.
- We provide a consistent benchmark for field-level from galaxy surveys. Build upon \(\texttt{NumPyro}\) and \(\texttt{BlackJAX}\).
Samplers comparison
= NUTS within Gibbs
= auto-tuned HMC
= adjusted MCHMC
= unadjusted Langevin MCHMC
10 times less evaluations required
Unadjusted microcanonical sampler outperforms any adjusted sampler
Model preconditioning
- Sampling is easier when target density is isotropic Gaussian
- The model is reparametrized assuming a tractable Kaiser model:
linear growth + linear Eulerian bias + flat sky RSD + Gaussian noise
10 times less evaluations required
Benchmark results
- Promising for future inferences, going multi-GPU using JaxDecomp
- Code readily available at github.com/hsimonfroy/benchmark-field-level
\(128^3\) PM on 8GPU:
4h MCLMC vs. \(\geq\)80h NUTS
Mildly dependent with respect to formation model and volume
Probing smaller scales could be harder
Handling bias in unadjusted MCLMC
- Microcanonical dynamics \(\implies\) energy should not vary
- Numerical integration yields quantifiable errors that can be linked to bias
- Stepsize can be tuned to ensure controlled bias, see Robnik+2024
reducing stepsize rapidly brings bias under Monte Carlo error
Primordial Non-Gaussianity from galaxies
Local-type PNG is constrained by the induced scale-dependent bias
\(\phi_{\mathrm{NL}}=\phi+{\color{purple}f_{\mathrm{NL}}}\phi^{2}\)
\(\delta(\boldsymbol k)\simeq\left(b_{1}+ b_\phi {\color{purple}f_\mathrm{NL}}k^{-2} \right) \delta_L(\boldsymbol k)\)
Field-level modeling of PNG
$$\begin{align*}w_g&=1+{\color{purple}b_{1}}\,\delta_{\rm L}+{\color{purple}b_{2}}\delta_{\rm L}^{2}+{\color{purple}b_{s^2}}s^{2}+ {\color{purple}b_{\nabla^2}} \nabla^2 \delta _{\rm L}\\&\quad\quad\! + {\color{purple}b_\phi f_{\rm NL}} \phi + {\color{purple} b_{\phi\delta} f_{\rm NL}} \phi \delta_{\rm L}\\\Delta \boldsymbol q_\parallel &= H^{-1} \dot{\boldsymbol q}_\parallel + {\color{purple}b_{\nabla_\parallel}} \nabla_\parallel \delta_\mathrm{L}\end{align*}$$
\(\phi_{\mathrm{NL}}=\phi+{\color{purple}f_{\mathrm{NL}}}\phi^{2}\)
\(\boldsymbol q_\mathrm{LPT} \simeq \boldsymbol q_\mathrm{in} + \Psi_\mathrm{LPT}(\boldsymbol q_\mathrm{in}, z(\boldsymbol q_\mathrm{in}))\)
one-shot 2LPT light-cone
\(n_g^\mathrm{obs}(\boldsymbol q) = (1+\delta_g(\boldsymbol q))\, {\color{purple}\bar n_g(\,r)}\, {\color{blue}W(\boldsymbol q)}\, {\color{purple}\beta_i} {\color{green}T^i(\theta)}\)
RIC relax + selection + imag. templates
\(\delta_g \sim \mathcal N(\delta_g^\mathrm{det}, \sigma^2)\) with
\(\sigma(k) = {\color{purple}\sigma_0}(1+{\color{purple}\sigma_2} k^2 + {\color{purple}\sigma_{\mu2}}(k\mu)^2)\)
EFT-based modeling, many scale cuts alleviating discretization effects (see Stadler+2024)
A word on integral constraints
Radial Integral Constraint
\(\delta_g \propto n_g - \braket{n_g}\approx n_g - \bar n_g(r)\)
i.e. impose \(\bar \delta_g(r) = 0\)
Global Integral Constraint
\(\delta_g \propto n_g - \braket{n_g} \approx n_g - \bar n_g\)
i.e. impose \(\bar \delta_g = 0\)
PRELIMINARYPreliminary results
on LRG SGC footprint
\(k_\mathrm{max} \approx 0.04\, [h/\mathrm{Mpc}]\)
\(\sigma[f_\mathrm{NL}] \approx 20\), consistent with power spectrum analysis (Chaussidon+2024)
Next steps:
- EFT model validation (w/o PNG) on AbacusSummit
- EFT model validation (w/ PNG) on PNG-Unitsims
- Systematics model validation on contaminated mocks
- Application to DESI DR1 LRG and QSO
Where we are
Tally
- Currently visiting Montréal in Laurence Perreault-Levasseur team
- Talks:
- Optimal cosmo information extraction at Sesto (for Euclid people)
- DESI meeting at Berkeley (for colab)
- CoBALt at Institut Pascal (for inflation theorists)
- Bayesian Deep Learning 3 at APC (for deep learners)
- ED Festival (for particle physicists)
- GDR Cophy 2h tutorial
- Papers:
- stat paper at NeurIPS2024 (from master internship)
- Benchmarking Field-Level in review on JCAP
- PNG measurement at the field level in prep
- Teaching: Bachelor 2 Biostats (20h) and Master 1 Maths (15) courses at UPsaclay
- Formation:
- VSS, Science ethics, Sustainable dev. (Open Science left)
- Euclid summer school
Next steps
- Scientific:
- Validation for PNG inference and application to DESI
- Alternative sampling method for field-level inference
- Going Multi-GPU
- Manuscript:
- Detailed plan at the end of December, first chapter in January
- Looking and applying for postdocs this Autumn,
hence the meetings and visits...
Thank you!