Mapping Dark-Matter Clusters via Physics-Guided Diffusion Models

Diego Royo; Brandon Zhao; Adolfo Muñoz; Diego Gutierrez; Katherine L. Bouman

cs.CVMar 15, 2026

Mapping Dark-Matter Clusters via Physics-Guided Diffusion Models

Authors: Diego Royo, Brandon Zhao, Adolfo Muñoz, Diego Gutierrez, Katherine L. Bouman

Organizations: Universidad de Zaragoza–I3A, Spain · California Institute of Technology, USA

Abstract

Galaxy clusters are powerful probes of astrophysics and cosmology through gravitational lensing: the clusters' mass, dominated by 85% dark matter, distorts background light. Yet, mass reconstruction lacks the scalability and large-scale benchmarks to process the hundreds of thousands of clusters expected from forthcoming wide-field surveys. We introduce a fully automated method to reconstruct cluster surface mass density from photometry and gravitational lensing observables. Central to our approach is DarkClusters-15k, our new dataset of 15,000 simulated clusters with paired mass and photometry maps, the largest benchmark to date, spanning multiple redshifts and simulation frameworks. We train a plug-and-play diffusion prior on DarkClusters-15k that learns the statistical relationship between mass and light, and draw posterior samples constrained by weak- and strong-lensing observables; this yields principled reconstructions driven by explicit physics, alongside well-calibrated uncertainties. Our approach requires no expert tuning, runs in minutes rather than hours, achieves higher accuracy, and matches expertly-tuned reconstructions of the MACS 1206 cluster. We release our method and DarkClusters-15k to support development and benchmarking for upcoming wide-field cosmological surveys.

Figures & tables

Method name	Model type	# Param.	Inputs			Automation level	End-to-end time
Method name	Model type	# Param.	Photometry	SL	WL	Automation level	End-to-end time
Lenstool [ 31 ]	Parametric	$\sim$ 10–100	✓(LTM 1:1)	✓	✓	Expert-guided	Hours [ 38 ]
WSLAP+ [ 15 ]	Hybrid	$\sim$ 1000	✓	✓	✓	Expert-guided	Hours
MARS [ 10 ]	Free-form	100 $\times$ 100	✗	✓	✓	Fully automatic	Hours [ 10 ]
Napier et al. [ 40 ] *	Parametric	$\sim$ 10–100	✓(LTM 1:1)	✗	✗	Fully automatic	Minutes
Ours	Free-form	512 $\times$ 512	✓	✓	✓	Fully automatic	Minutes

Table 1: Overview of galaxy cluster mass reconstruction methods. We differentiate between tools that require expert configuration (e.g., Lenstool [ 31 ] , WSLAP+ [ 15 ] ) and fully automatic approaches (MARS [ 10 ] , the method from work by Napier et al. [ 40 ] , and ours). Our free-form approach leverages a diffusion prior trained on DarkClusters-15k , enabling 512×512 free parameters with mass reconstructions within minutes, significantly faster than traditional methods. Unlike other methods that assume that Light Traces Mass (LTM) with a 1:1 ratio, we freely model photometry along with Strong Lensing (SL) and Weak Lensing (WL). * Our reimplementation.

Figure 1: Sketch of a typical gravitational lens system. The light path of an object at angular position $\boldsymbol{\beta}$ on a source plane $S$ is deflected by an angle $\tilde{\boldsymbol{\alpha}}$ due to the presence of a galaxy cluster on the lens plane $L$ , so the object appears at angular position $\boldsymbol{\theta}$ instead.

Figure 2: Left: For the creation of our DarkClusters-15k dataset, we extract 600 simulated galaxy clusters from IllustrisTNG and SIMBA as 3D point clouds, and compute the surface mass density $\Sigma$ for $n=25$ different line of sight $\vec{r}_{i}$ projections. Right: For each projection, we then independently generate measurements for strong lensing (i.e., positions of multiple images for each source) and weak lensing (i.e., shear $\boldsymbol{\gamma}$ at the source positions).

Panel of Figure 3 — Figure 3: Left: Our diffusion-based approach produces multiple plausible reconstructions; in each row we show one representative sample of the lens convergence $\kappa$ , and the posterior mean and standard deviation over 20 samples. Right: Reliability diagram, showing close correlation between our model’s predicted uncertainty and the actual error.

Panel of Figure 4 — Figure 4: Ablation study on the different input components used by our method to estimate cluster masses. Left: Each row shows photometry $\textbf{P}_{f}$ for one cluster, and mass reconstructions using only one input as stated. Right: Adding strong (SL) and weak lensing (WL) information (SNR 0.01) improves the photometry (P) results. We report PSNR of each image w.r.t. the ground truth (GT).

Metric	KS [ 26 ]		Napier et al. [ 40 ]		UNet		Ours
Metric	TNG	SIMBA	TNG	SIMBA	TNG	SIMBA	TNG	SIMBA
PSNR $\uparrow$	20.9	21.1	26.0	22.1	31.1	29.6	33.3	30.9
SSIM $\uparrow$	0.443	0.437	0.703	0.565	0.894	0.855	0.927	0.884
PCC $\uparrow$	0.189	0.197	0.832	0.616	0.924	0.863	0.975	0.930
Held-out $\boldsymbol{\alpha}$ RMSE $\downarrow$	3.697	3.919	1.475	3.295	0.411	1.597	0.530	1.485
Held-out $\kappa$ RMSE $\downarrow$	0.0552	0.0521	0.0173	0.0169	0.0067	0.0073	0.0058	0.0067

Table 2: Evaluation results for galaxy clusters at

z_{L}=0.5

in the TNG and SIMBA test sets of DarkClusters-15k . We report the Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), Pearson Correlation Coefficient (PCC), and the

\kappa

\boldsymbol{\alpha}

and

\boldsymbol{\gamma}

RMSE at held-out points (i.e., points without any photometry or lensing signals).

Panel of Figure 5 — Figure 5: Comparison of our method with Kaiser-Squires (KS) [ 26 ] , Napier et al. [ 40 ] and our UNet baselines. We show lens convergence $\kappa$ for five test samples of DarkClusters-15k at different redshifts $z_{L}$ and simulations. Napier et al. produce an overly smooth mass prediction, and yield inaccurate estimates in certain regions due to its strict 1:1 mapping of mass and photometry. Also, contrary to both Napier et al. and our UNet baselines, our method includes gravitational lensing observables and generates samples from a posterior distribution with calibrated uncertainty estimation; we show one sample of this posterior along with the mean and std. dev. computed for 20 generated samples.

Panel of Figure 6 — Figure 6: Our mass estimate for the MACS 1206 cluster, compared with the one shown in work by Caminha et al. [ 8 ] . We show our dimensionless surface mass density $\kappa$ estimate along with its isolines and critical curve at $z_{S}\approx 1.44$ ; the reference work shows the diffuse cluster halo which does not account for light sources in the lens plane and thus is naturally smoother.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Distribution of the galaxy clusters selected from IllustrisTNG and SIMBA to construct DarkClusters-15k . We show their virial masses (i.e., their total gravitationally-bound masses) and their virial radii (i.e., the radius that encloses that mass).

Explore similar work

Generative Diffusion Priors for 3D Mapping of the Dark Universe
May 30, 2026Brandon Zhao, Diana Scognamiglio, Olivier Doré +13D Reconstruction Bayesian Inverse Problems
Department of Computing and Mathematical Sciences, California Institute of Technology · Jet Propulsion Laboratory, California Institute of Technology · Department of Physics, Duke University +1
Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines
Jul 21, 2026Guillaume Payeur, Laurence Perreault-Levasseur, Gabriel Missael Barco +1Pixel-Space Diffusion Models Diffusion Posterior Sampling
Department of Physics, Universit´e de Montr´eal, Montreal QC, Canada · 2Mila - Quebec AI Institute, Montreal QC, Canada · 3Ciela Institute, Institute for Astrophysics and Machine Learning, Montreal QC, Canada +3
A plug-and-play approach with fast uncertainty quantification for weak lensing mass mapping
Mar 23, 2026Hubert Leterme, Andreas Tersenov, Jalal Fadili +1Uncertainty Quantification
Université Caen Normandie, ENSICAEN, CNRS, Normandie Univ, GREYC UMR 6072, F-14000 Caen, France · Université Paris-Saclay, Université Paris Cité, CEA, CNRS, AIM, 91191, Gif-sur-Yvette, France · Institutes of Computer Science and Astrophysics, Foundation for Research and Technology Hellas (FORTH), Heraklion, Greece

May 30, 2026J/K move · Enter open · S save

Generative Diffusion Priors for 3D Mapping of the Dark Universe

Brandon Zhao, Diana Scognamiglio, Olivier Doré, Katherine L. Bouman

Department of Computing and Mathematical Sciences, California Institute of Technology · Jet Propulsion Laboratory, California Institute of Technology · Department of Physics, Duke University+1

3D Reconstruction Bayesian Inverse Problems Diffusion Posterior Sampling Diffusion Models

Open paper AskPDF arXiv

Reconstructing the three-dimensional distribution of dark matter from weak-lensing observations is a central but highly ill-posed inverse problem in cosmology. Unlike standard 3D reconstruction with multiple viewpoints, we observe the universe from a single line of sight, through noisy shape distortions of galaxies with uncertain distances, so meaningful recovery of the 3D matter field requires strong prior assumptions. Existing methods either produce point estimates with handcrafted priors or use neural ensembles for approximate Bayesian uncertainty, and struggle to capture the non-Gaussian, filamentary structure of the cosmic web. With the advent of new high-resolution cosmological simulations, we now have an alternative source of prior knowledge that captures the nonlinear statistics of structure formation with far greater fidelity than analytic prescriptions. We leverage these simulations to build a new dataset

\texttt{Conicus3D}

, which enables us to learn a data-driven diffusion-model prior capturing the full 3D distribution of dark matter structure across cosmic time. Building on recent plug-and-play approaches, we modify a diffusion-based posterior sampling scheme to the 3D weak-lensing setting, combining the learned prior with a differentiable physical forward model. On realistic simulations targeting a modern weak lensing survey, our approach yields substantially improved 2D and 3D reconstruction accuracy over baseline methods. Moreover, it produces posterior samples whose statistics closely track the underlying simulations, while remaining robust to moderate shifts in cosmology.

Abstract

Explore similar work

Generative Diffusion Priors for 3D Mapping of the Dark Universe

Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

A plug-and-play approach with fast uncertainty quantification for weak lensing mass mapping