Probabilistic Modeling and Analysis
for Cosmology

Andreas Tersenov

29 January 2026

Background: Weak Lensing Analysis

An End-to-End approach: From Pixel-Level Reconstruction to Robust Cosmology

The Data/Pixels

Catalogues: shear calibration, masks, blinding, tomography...
Mass Mapping: Inverse problems, PnPMass, Denoising.
Methods: DL Priors, Convex Optimization.

The Signal & Physics

Summary Statistics: 2pt, HOS (Peaks, PDF, l1-norm)
Systematics: Baryonic Feedback, IA, Masks, Noise.
Theory: Large Deviation Theory validation, Euclid paper on theory HOS.

The Inference

Methodology: Simulation-Based Inference (SBI) & Explicit Inference.
Computational Core: JAX ecosystem, Differentiable Programming, Samplers, ML-emulators.

Differentiable pipelines (JAX/PyTorch), ML/DL, probabilistic programming (NumPyro), advanced sampling, neural density estimation, ML/DL for scientific inference, and the diagnosis of model misspecification

Impact of Mass Mapping on Cosmology Inference

Weak Lensing Mass Mapping

From convergence to shear:
From shear to convergence:

\gamma_i = \hat{P}_i \kappa

\kappa = \hat{P}_1 \gamma_1 + \hat{P}_2 \gamma_2

Mass mapping is an ill-posed inverse problem

Different algorithms have been introduced, with different reconstruction fidelities

But... does the choice of the mass-mapping algorithm have an impact on the final inferred cosmological parameters?

Or as long as you apply the same method to both observations and simulations it won't matter?

Results

Where does this improvement come from?

FoM improvement → 150%

Plug-and-Play Mass Mapping with Uncertainty Quantification

Deep Learning for Mass Mapping?

Mass mapping method	Type	Accurate	Flexible	Fast rec.	Fast UQ
Iterative Wiener	Model-driven (Gaus. prior)	✗	✓	✓	✗
MCALens	Model-driven (Gaus. + sparse)	≈	✓	✗	✗
DeepMass	Data-driven (UNet)	✓	✗*	✓	✓
DeepPosterior	Data-driven (UNet + MCMC)	✓	✓	✗	✗
MMGAN	Data-driven (GAN)	✓	✗*	≈	≈
What we'd like	Data-driven	✓	✓	✓	✓

Plug-and-Play Mass Mapping

Instead of explicitly writing down a prior, we learn what a "likely" 𝜅 looks like from simulations and enforce it through denoising.

Results

Impact of baryonic effects on cosmology inference with HOS

Testing impact of baryonic effects on WL HOS

Idea - Explore two things:

Information content of summary statistics as a function of scale cuts & survey area
Testing the impact of baryonic effects on posterior contours

This will show:

On what range of scales can the different statistics be used without explicit model for baryons
Answer the question: how much extra information beyond the PS these statistics can access in practice

cosmoGRID simulations:

Power Spectrum

Wavelet l1-norm: sum
of wavelet coefficients
within specific amplitude
ranges across different
wavelet scales

Wavelet peaks: local maxima of wavelet
coefficient maps

N-body sims, providing DMO and baryonified full-sky κ-maps from the same IC, for 2500 cosmologies. Baryonic effects are incorporated using a shell-based Baryon Correction Model (Schneider et. al 2019).

Inference method: SBI

Training objective

\mathcal{L} = - \log p_\phi ( \theta | x )

The Scaling of Baryonic Bias with Survey Area

Determining Robust Scale Cuts

Information Content at Large Scales

Results

+ Benchmarking Theoretical Wavelet l1-Norm Predictions Against Cosmological Simulations

+ Euclid Key Paper: DR1 analysis with theory-based HOS

Why Field-Level Inference?

Unifying Reconstruction, Systematics, and Cosmology

\delta_{\rm IC}

\gamma_{\rm Obs}

Explicit & Lossless

Zero Compression: Inference directly on the pixel-level data.
FLI targets the full high-dim posterior.

Reconstructing ALL Latent Variables

Beyond θcosmo reconstruct the entire formation history.
Initial Conditions (ICs)
Evolved DM Distribution
Peculiar Velocities

FLI Advantages:

Self-Consistent Systematics: Modeling noise/masks + forward model other effects physically in 3D (cosmology + systematics become part of one model, which we can sample).
Interpretability: analyzing stat errors directly in data space
Signals mix less → model misspec. easier to detect
Validation: Posterior Predictive Tests & Cross-Correlations (e.g., kSZ, X-ray) without CV.
Discovery: A platform to test physics at the field level.

FLI for Cosmic Shear

FLI for Primordial Non-Gaussianity

\theta_{\rm cosmo}

Gravity solver

Baryonification

Lensing+IA

Survey effects & systematics

Initial conditions

Evolved density

DM+baryons

Lensing maps

Predicted observations

Part 1: Gravity Solver

Goal: a differentiable, scalable gravity backbone that unlocks nonlinear information without bias.

Challenges:

Computational efficiency
- Scale to Euclid/LSST volume (Gpc)
- Resolve smaller (non-linear) scales (kpc)
Differentiability at scale: stable gradients for HMC, memory/communication

Big dynamic range

Scalability + resolution → use fully non-linear, distributed (multi-GPU) PM simulations

Solution: A Hybrid Physical-Neural Approach

JaxPM + jaxdecomp + jax-healpy

3D domain decomposition and distributed FFTs

Use FNO instead

Performs global convolutions in Fourier space, learning the underlying continuous operator → learn effective Green’s function of the Poisson eqn from the data
Captures the non-local tidal deformations
Can distinguish between different cosmic environments
Resolution Invariant: learned on small, high-res boxes → Zero-shot deployment on big lightcone mesh without retraining.

Computational efficiency

Resolution independence: Neural operators can generalize across resolutions after single-
scale training, facilitating multi-resolution analysis strategies

For field-level inference, we need correct velocities and displacement histories, not just final density fields. Therefore, the advanced version of my proposal is to implement Active Neural Force Correction. We inject the neural network inside the time-integration loop to correct the PM forces at every step.

Run Coarse PM, but at every time step, call an FNO to predict the "Residual Force" (Exact Force - PM Force) and kick the particles.
FNO learns the mapping
between function spaces (the density field 𝜌 to the displacement potential 𝜓 )

F_\theta(\mathbf{x}, a)=\frac{3 \Omega_m}{2} \nabla\left[\phi_{P M}(\mathbf{x}) * \mathcal{F}^{-1}\left(1+f_\theta(a,|\mathbf{k}|)\right)\right]

"Hybrid Physical Neural ODEs for Fast N-body simulations" Lanzieri, Lanusse, Starck

PM solvers still face a resolution bottleneck

(mesh grid → smooths out gravity)

Strategy: inject neural correction into time-integration loop to recover small scales

Learned Isotropic Fourier Filter

predict the complex 3D force field from density, instead of just modifying potential

\theta_{\rm cosmo}

Gravity solver

Baryonification

Lensing+IA

Survey effects & systematics

Initial conditions

Evolved density

DM+baryons

Lensing maps

Predicted observations

Part 2: Baryon Modeling

Challenge:

More accurate modeling of baryons beyond simple thermodynamics
Scalable/Fast enough for FLI + differentiable

In BORG-WL → EDG

Displaces particles to minimize enthalpy assuming a simple polytropic eqn of state.
It cannot create features that don't exist in the potential (e.g., AGN bubbles, non-local feedback effects).

Conditional Optimal Transport Flow-Matching (Kerrigan et al., 2024)

The Bridge: Evolve directly from DM → Gas.
Conditioned on Cosmology & the IC (paired datasets)
Learns the optimal transport map from gravity solver output to hydro density fields.
Captures non-linear feedback morphology that simple EOS models may miss.
Reproduces the correct HOS of the gas distribution.

Baryon Emulator

PM Density

Velocity

Density

Temperature

"BaryonBridge: Interpolants models for fast hydrodynamical simulations" Horowitz et al. 2025

Part 3: Sampling methods

Scalable posterior exploration in high-dim spaces

\mathcal{L}(\theta)=\mathbb{E}\left[\left\|u_\theta\left(X_t, t\right)-\left(\frac{X_1-X_t}{1-t}-\nabla E\left(X_1\right)\right)\right\|^2\right]

(unadjusted) MCLMC for scaling to ≥10^6 dim

"Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching", Havens et al.

Augment sampling with learned proposal distribution via NFs
Flow Matching to learn probability path between prior and posterior integrated into a MCLMC/Annealed HMC loop
provide high-fidelity proposal distributions, significantly accelerating convergence

Diffusion samplers >1 energy evaluations
per gradient update

Amortizes the cost of the forward model can perform hundreds of cheap network updates for every single expensive simulation

Adjoint Sampling learns directly from the energy function and its gradient, without requiring many data samples (many grad updates/sample

Part 4: Pixel-Level Forward Modeling

Foundation models

To reuse & share prior models, build flexible model, capable of modeling different types of data
Diff. models from corrupted data using an Expectation-Maximization approach
Multimodal generative models to learn both from imaging & spectroscopy
More info in latent space to improve the galaxy characteristics measurement
Foundational data-driven prior models → adapt to different desired populations at inference time, and can be used for many types of observations.

End-to-end diff/ble shear forward model

Integrate detection, deblending, shear estimation, robustness to image defects.
Denoising diffusion model priors over galaxy morphologies → optimize the posterior + sample
Inferred population-level distributions of galaxy properties, and individual galaxy shapes jointly from the pixel data.

Combine multi-band, multi-instrument data to produce consistent UQ and posterior samples for each object.

t=T

t=0

\theta_{\rm cosmo}

Gravity solver

Baryonification

Lensing+IA

Survey effects & systematics

Initial conditions

Evolved density

DM+baryons

Lensing maps

Predicted observations

Part 6: Field Level Systematics

Exploiting Phase Information to Distinguish Physics from Artifacts → Euclid realism

In power spectra, systematics often mimic cosmological signals.
In the 3D field, they possess distinct spatial signatures.
By differentiating through the forward model, the inference engine uses phase information to separate signal from noise.

Inject systematics directly into the pipeline:
- IA → TATT
- Spatially varying shear bias and survey depth
Survey Geometry
Photo-z:
- BHM to sample n(z) and marginalise over redshift uncertainties, using forward model of the photometric fluxes

Probabilistic Modeling and Analysis for Cosmology

Background: Weak Lensing Analysis

An End-to-End approach: From Pixel-Level Reconstruction to Robust Cosmology

Impact of Mass Mapping on Cosmology Inference

Weak Lensing Mass Mapping

Results

Where does this improvement come from?

Plug-and-Play Mass Mapping with Uncertainty Quantification

Deep Learning for Mass Mapping?

Plug-and-Play Mass Mapping

Results

Impact of baryonic effects on cosmology inference with HOS

Testing impact of baryonic effects on WL HOS

cosmoGRID simulations:

Inference method: SBI

Results

+ Benchmarking Theoretical Wavelet l1-Norm Predictions Against Cosmological Simulations

+ Euclid Key Paper: DR1 analysis with theory-based HOS

Why Field-Level Inference?

Unifying Reconstruction, Systematics, and Cosmology

FLI for Cosmic Shear

FLI for Primordial Non-Gaussianity

Part 1: Gravity Solver

Goal: a differentiable, scalable gravity backbone that unlocks nonlinear information without bias.

Solution: A Hybrid Physical-Neural Approach

Part 2: Baryon Modeling

Baryon Emulator

Part 3: Sampling methods

Part 4: Pixel-Level Forward Modeling

Part 6: Field Level Systematics

Exploiting Phase Information to Distinguish Physics from Artifacts → Euclid realism

Probabilistic Modeling and Analysis
for Cosmology