Differentiable Forward Modelling for Radio Weak Lensing

Overview

1 ARGOS and argosim A differentiable simulator for radio interferometers 3
2 Antenna optimisation Designing array layouts by gradient descent 5
3 Radio weak lensing Opportunities and challenges 8
4 SHINE Shear inference directly from visibilities 10
5 Conclusions 16

The ARGOS Project

ARGOS ARGOS European Commission

ARray for Gigahertz ObservationS: a European initiative for a low-cost, sustainable radio interferometer in Crete 🇬🇷

argos-telescope.eu

Modular, scalable and affordable design
Pathfinder: 5 × 6 m dishes installed at FORTH
Full planned array: ~1000 antennas, 1–3 GHz, baselines up to ~4 km
Science: FRBs (strongly lensed FRBs → ~1% \(H_0\)), pulsars, GWs, radio transients
Designing the array and its imaging software calls for fast, flexible simulations

ARGOS Pathfinder at FORTH: ground and aerial views

argosim: Differentiable Interferometric Simulations

argosim simulation pipeline: from antenna layout to CLEAN image

Centofanti et al. (2026a) • arXiv:2606.25573
github.com/ARGOS-telescope/argosim

Lightweight, modular Python package for radio simulations
Full chain: array → uv tracks → visibilities → dirty beam & image → CLEAN, plus PSF/image metrics
JAX backend: ~40× faster than NumPy on CPU, up to ~800× on GPU
End-to-end differentiable → gradients
Open source, Docker image and GUI
pip install argosim

Antenna Optimisation: Differentiable PSF Forward Model

Antenna positions to PSF optimisation loop

\[\hat{X} = \arg\min_X \; \frac{\lVert b(X) - b_{\rm tg} \rVert_2^2}{N_{\rm px} \max |b_{\rm tg}|^2} + \lambda_D\, g_D(X) + \lambda_R\, g_R(X) + \lambda_C\, g_C(X)\]

Proof of concept: Can we optimise the position of antennas in an array for a given science case?
Compare the PSF \(b(X)\) synthesised by argosim from antenna positions \(X\) to a given target PSF \(b_{\rm tg}\): e.g. a Gaussian PSF
Physical constraints: minimum antenna separation (\(g_D\)), available area (\(g_R\)), array centring (\(g_C\))
Any differentiable loss works, e.g. sidelobe level or FWHM directly

Antenna Optimisation Results

Initial and final beam PSF cross-section

	ISL [dB]	PSL [dB]	FWHM [px]	e
Initial	−6.3	−14.2	(40, 38)	0.32
Final	−19.7	−27.5	(32, 32)	0.16

Main lobe matches the target; sidelobes suppressed by ~13 dB
Optimum for one observing setup → next: multiple declinations, frequencies and track lengths in a single loss

Weak Lensing in One Slide

A powerful cosmological probe

Traces all matter, dark and baryonic
Sensitive to both geometry and the growth of structure: constrains \(\Omega_m\), \(S_8\) and dark energy
KiDS, DES, HSC → Euclid, Rubin, Roman, SKA

A tiny statistical signal

\[\epsilon^{\rm obs} \approx \epsilon^{\rm int} + \gamma \quad \Rightarrow \quad \langle \epsilon^{\rm obs} \rangle \approx \gamma\]

\(\lvert\gamma\rvert \sim 0.01 \ll \sigma_\epsilon \sim 0.2\) → many precise shapes

What we must control

\[\hat{\gamma}_i = (1 + m_i)\,\gamma_i + c_i\]

Multiplicative \(m\) and additive \(c\) biases must vanish
Any residual biases the cosmology

Top: lensing by large-scale structure (ESA). Bottom: cosmic-shear constraints vs Planck CMB (Wright et al. 2025)

Why Weak Lensing with the SKA?

Opportunities

AA4: 5 000 deg², ~50 M galaxies, \(z_m \sim 1.1\)
\(S_8\) to 5% alone, 3% with Euclid or LSST
With clustering and HI: ~1% on ΛCDM parameters
Different systematics → robustness test
Polarisation and kinematics separate lensing from intrinsic shapes

Challenges

Data are visibilities: galaxies are delocalised in Fourier space
Deconvolution artefacts → biased shapes
Ellipticity is ill-defined for complex morphologies → model bias

~1 min. Forecasts from the AASKAII chapters (2026). SKA-Mid AA4 Band 2 survey: 5 000 deg², 2.7 galaxies per square arcminute (about 50 million in total, our estimate from these two numbers), median redshift 1.1. The 5% and 3% on S8 assume high-fidelity shape reconstruction, which is the problem SHINE addresses. Figure: 1 and 2 sigma forecasts for SKA-Mid AA4 alone (dashed blue), a Stage-IV optical survey (red), their cross-correlation (green), which removes additive systematics by construction, and the full combination (filled), which can remove multiplicative ones too. The ~1% figure combines lensing, clustering and HI (Harrison et al. 2026, arXiv:2607.11254). If asked about AA: all of this needs AA4. At least 1 resolved galaxy per square arcminute is needed to detect the shear signal: AA4 (0.55 arcsec PSF) clears this, AA (2.06 arcsec) never does at any depth, and a 75 km layout just reaches it. Current radio shear methods reach m of a few 10^-2 and c of a few 10^-4 on AA4 simulations (Tripathi et al. 2026). Polarisation and kinematics can raise statistical power by a factor of about 5-10. Bottom figure: (a) gridded visibilities, (b) dirty image at S/N 30.</aside>

Top: forecast constraints (after Harrison et al. 2026). Bottom: radio-adapted COSMOS galaxy (Centofanti et al. 2026b)

SHINE: SHear INference Environment

SHINE: forward-modelling framework for probabilistic shear inference

Centofanti et al. (2026b) • Coming soon to arXiv!
github.com/CosmoStat/SHINE

Forward-modelling framework for probabilistic shear estimation
No shape measurement: model the data directly, not galaxy ellipticities
JAX-powered: autodiff and GPU acceleration (JAX-GalSim, BlackJAX)
Versatile: designed for both optical and radio surveys
Radio forward model built on argosim

Generative Forward Model of Radio Visibilities

Generative forward model

\[V^{\rm obs} = M \odot \mathcal{F}\big(\gamma \otimes \mathcal{G}_\theta(z)\big) + N\]

\(\mathcal{G}_\theta(z)\): galaxy generator (simulates galaxy images)
Shear \(\gamma\) applied at pixel level with JAX-GalSim
uv mask \(M\) from argosim: SKA-Mid, 8 h track, 1.4 GHz
Gaussian likelihood directly in visibility space: no imaging, no deconvolution, no shape measurement

Sampling the Shear Posterior

Hierarchical model: galaxies in a patch share one \(\gamma\); morphology parameters \(z_i\) are marginalised
100 galaxies per batch → >700 parameters → gradient-based MCLMC sampling
Posterior per batch fitted with a GMM, then batches combined analytically
~60 ms per galaxy on V100 GPUs (parametric model)

Results: Parametric Galaxies

Shear error vs fiducial shear compared to DeepShape

10 000 T-RECS-like galaxies (Spergel profiles), SKA-Mid, S/N 10–80
~3× smaller uncertainty than DeepShape (Tripathi et al. 2025), with no apparent multiplicative bias
Coverage test (125 posteriors): uncertainties are well calibrated

A Learned Galaxy Model: Autoencoder

Training data: 80 000 HST COSMOS galaxies (no well-resolved radio sample exists)
Architecture: Deep Compression AE (ResNet blocks), 128×128 image → 4×4 latent \(z\)
Loss: the HST PSF is applied to the decoder output, so the AE learns the unconvolved profile

\[\mathcal{L}(\hat{x}) = \lVert x - H \ast \hat{x} \rVert_1 + \lambda\, \mathrm{TV}(\hat{x})\]

Prior: Real NVP normalising flow, \(z = f(u)\) with \(u \sim \mathcal{N}(0, I)\)
Speed: distilled decoder (18 M → 0.8 M parameters, float16), ~10× faster

Top: AE architecture. Bottom: COSMOS galaxies and their unconvolved AE reconstructions (Centofanti et al. 2026b)

Realistic Galaxies: Model Bias and How to Fix It

10 000 HST COSMOS galaxies adapted to radio, with \(\gamma_{\rm true} = (-0.05, 0.05)\)

Spergel model: morphology mismatch → bias > 10σ
AE generative model with flow prior on the latent space: bias reduced to < 2σ
MAP point estimate: up to 30× faster and consistent with the posterior, but without uncertainty

Conclusions

argosim • Centofanti et al. (2026a) • arXiv:2606.25573

Lightweight, JAX-based, differentiable interferometric simulator (pip install argosim)
Demonstration of gradient-based array design: ISL −6 → −20 dB
- Many possible evaluation metrics for science cases
- Can incorporate physical constraints on antenna positions

SHINE for radio weak lensing • Centofanti et al. (2026b)

Shear inferred directly from visibilities: no imaging, no deconvolution, no shape measurement
Parametric galaxies: ~3× smaller uncertainty than DeepShape, no apparent \(m\)-bias, well calibrated
Realistic galaxies: a learned generative achieves a model bias from of < 2σ on real COSMOS images

Towards cosmology with SKA

AA4 forecasts of 3–5% on \(S_8\) need unbiased shear: this is the problem we are tackling
Still to add: detection, full-field modelling, varying PSFs, SEDs, intrinsic alignments
Then: scale to ~50 M galaxies and apply to real radio data