SHINE: Shear inference environment - Generative forward modelling for radio weak gravitational lensing

Introduction

The gravitational lensing effect describes how photons, emitted by distant sources, are deflected as they travel through space due to the presence of gravitational potentials between the sources and the observer, distorting the apparent shape and position of the observed objects (Bartelmann [4]). This phenomenon provides a way to probe the large-scale structure (LSS) of the Universe by mapping the distribution of matter overdensity projected along the line of sight. In the case of lensing by the LSS, the effect is very subtle, and the apparent deformation of the observed sources is minimal, with a change of less than 1%. This regime is known as weak gravitational lensing (WL) (Schneider [64]; Hoekstra & Jain [38]; Munshi et al. [51]; Bartelmann [4]; Kilbinger [43]; Mandelbaum [47]), and the distortions can be modelled locally by a linear transformation that includes an isotropic magnification component, called convergence ($\kappa$), and an anisotropic stretching component, called shear ($\gamma$). The WL signal allows constraints to be placed on cosmological parameters related to the dark matter density and the growth of the LSS (Bartelmann & Schneider [5]). In practice, the shear field is estimated statistically by measuring the shapes of large numbers of galaxies and averaging their observed ellipticities, under the assumption that their intrinsic orientations are randomly distributed. However, given that the shear is very small compared to the ellipticity of the galaxies, it is necessary to measure the galaxy shapes with extreme precision. Furthermore, these measurements must be calibrated for the instrumental response of the telescope, the point spread function (PSF), which needs to be precisely modelled (Massey et al. [49]). Finally, any additive or multiplicative bias in the shear estimate must be calibrated, requiring costly simulations, which highlights the need for unbiased methods.

The weak lensing signal has been successfully measured in past surveys using optical observations. Surveys such as CFHTLenS (Heymans et al. [35]), KiDS (de Jong et al. [23]), DES (Collaboration et al. [18]), and HSC (Miyazaki et al. [50]) have made it possible to derive cosmological constraints from WL studies thanks to their depth, sky coverage, and compact PSF. Upcoming optical WL surveys such as Euclid (Laureijs et al. [45]; Euclid Collaboration et al. [28]), LSST (Abell et al. [1]; Ivezić et al. [41]), or Nancy Grace Roman (Spergel et al. [69]; Akeson et al. [2]) will observe the sky with even greater depth and sensitivity, allowing for the detection of a higher galaxy density, and promising significant improvements in cosmological results.

Radio WL is an evolving domain, and the detection of this signal using radio observations alone has not yet been conclusive. Studies such as Chang et al. [16] report a detection in radio wavelengths, measuring the shapes of galaxies directly in Fourier space, but with a very low detection significance due to the low galaxy density detected. Other studies such as Tunbridge et al. [76]; Demetroullas & Brown [24]; Hillier et al. [36]; Liu et al. [46] perform a cross-correlation analysis, measuring galaxy shapes in both radio and optical observations. Overall, these studies highlight the major limitations of current radio telescopes in conducting WL studies, namely the low density of detected galaxies and the small survey areas compared to optical surveys [53, 31].

The next generation of radio telescopes such as the Square Kilometre Array (SKA) [25, 10] and the Deep Synoptic Array (DSA) [30] will observe the radio sky with unprecedented coverage (over 30 000 deg$^2$) and, thanks to their physical extension and sensitivity, will resolve a large number of galaxies, making WL radio studies possible [33] (Brown et al. 2015). Some of the most prominent advantages of this type of instrument is the ability to access to higher redshifts, or obtaining photometry directly from the 21 cm HI line without relying on complementary observations [7]. Furthermore, polarised radio observations could allow the intrinsic shapes of galaxies to be estimated, alleviating key systematics in shear estimation, such as intrinsic alignments [11]. Finally, another systematic that is of lesser relevance in radio is the PSF, commonly referred to as the dirty beam, which, unlike in optical telescopes, is well known and is determined by the spatial distribution of the radio antennas, the operating frequencies, and the direction of observation. Conversely, radio surveys face other challenges for WL studies, such as a lower galaxy density and poorer resolution compared to optical telescopes. Moreover, the immense volume of data that will be acquired by the new generation of interferometers poses a major technical challenge for estimating cosmic shear, requiring highly efficient and scalable algorithms.

This paper is organised as follows: in Sect. 2 we discuss state-of-the-art methods for shear estimation in radio WL, in Sect. 3 we introduce the theoretical background, particularly the WL formalism and radio interferometry imaging, in Sect. 4 we present our forward modelling framework for shear estimation, In Sect. 5 we explain the simulated dataset used in this work, in Sect. 6 we present the experimental results and finally describe future prospects and conclude in Sect. 7.

Related work

The traditional procedure for estimating cosmic shear consists of measuring the shapes of galaxies individually. There are two main ways of doing this: one approach is to measure galaxy ellipticity directly from images, for example by computing second-order moments, deriving these moments from a decomposition of the image into basis functions, or using convolutional neural networks. The other approach consists of assuming a model of the galaxy surface brightness and ellipticity, and either individually or collectively fitting them to the observations. For a review of shape measurement methods see Mandelbaum et al. [48]. The advantage of the second approach is that the PSF can be included in the forward model, whereas the first approach requires separate PSF calibration. However, the forward model approach is susceptible to model bias if it fails to represent the morphology of the observed galaxies. Model bias will require further calibration of the estimated shear, adding computational overhead. This calibration cost may be limiting for large surveys such as Euclid, Nancy Grace Roman Space Telescope, and SKA, which will observe trillions of galaxies.

Radio astronomy poses another challenge, which is the fact that observations are not acquired directly in the image domain but in Fourier space, where the definition of galaxy shape and ellipticity are not trivial. One option is to reconstruct the galaxy image from the Fourier samples (commonly called $uv$ samples or visibilities in radio interferometry) using a deconvolution algorithm such as CLEAN (Högbom [39]) or its more recent versions [21, 56, 19, 26] (Garsden et al. 2015) which could introduce artefacts in the shape of the reconstructed galaxy due to the complex structure of the radio PSF. In Patel et al. [53], real-space shapelets are applied to CLEAN-reconstructed images to measure the shapes of galaxies. The other option is to measure the shape directly in Fourier space. In Chang et al. [16] a method based on shapelets, introduced in Chang & Refregier [15], is used to estimate the shape of galaxies in Fourier space directly from interferometric data. Other works such as Tripathi et al. [75] use convolutional neural networks to extract features from both the galaxy image and the radio dirty beam to estimate galaxy ellipticities. However, this type of method will require retraining when applied to real data, for which there are no labels for the true shape of the observed galaxies or deconvolved images. Frameworks such as Rivi et al. [59] model the light profile of galaxies analytically in $uv$ space, which allows the likelihood to be evaluated in Fourier space directly using visibilities. However, their analytical model is limited to exponential profiles, which are overly simple for well-resolved galaxies. Lastly, the approaches mentioned above have a potential source of model bias in the shear estimation due to the simplification of galaxy shapes to the ellipticity, which is not a well defined quantity for arbitrary galaxy morphology [58].

In this paper, we present a method for estimating shear from radio galaxies that it is free from shape measurements. Our framework models shear hierarchically for multiple galaxies within the same region of the sky. This method consists of a forward galaxy model that includes both the shear distortion and the instrumental response of the radio interferometer, simulating the visibilities of each galaxy individually. We leverage two different galaxy generative models, adapting to the morphological complexity of the observed galaxies in order to quantify and address potential model biases in the case of a morphological mismatch.

Theoretical background

Weak gravitational lensing

The non-homogeneous distribution of dark matter between distant galaxies and the observer slightly deflects the emitted photons, modifying the apparent shape of the observed galaxies. In the case of lensing due to the LSS, this deformation is subtle and can be modelled locally by a linear transformation described by the following matrix:

where $\kappa$ is a scalar representing an isotropic magnification of the observed shape, and $\gamma= \gamma_{1} + i\gamma_{2}$ is a complex number that models an anisotropic stretching. These two quantities are closely related since they are given by second-order derivatives of the two-dimensional lensing potential, $A_{ij} = \delta_{ij} - \partial_i\partial_j\psi$ (see Bartelmann [4]; Kilbinger [43]). Thus, a source with intrinsic ellipticity $\epsilon^{i}$, when distorted by the WL phenomenon, will be observed with the following ellipticity

Note that we are not considering the size change by $(1-\kappa)$, as in the context of WL these distortions are very small and we can approximate the reduced shear as follows

One way to estimate shear is by averaging the observed ellipticities of multiple sources in the same region of the sky, where the shape distortion is assumed to be the same for all objects, obtaining

Then, since the intrinsic shapes of galaxies do not have a preferred orientation, the average of intrinsic ellipticities cancels out, yielding

where $n$ represents shape noise, which depends on the number of averaged galaxies. In this work, we do not model the possible intrinsic alignment (IA) of background galaxies with their hosting LSS, which can correlate galaxy shapes even in the absence of lensing (Schäfer [63]). Estimating the reduced shear from noisy ellipticity measurements yields a biased estimate, because the ellipticity estimator responds non-linearly to noise and to imperfect galaxy and PSF models. To first order this bias is described by a linear model,

with a multiplicative bias $m$ and an additive bias $c$. Calibration methods such as metacalibration ([40, 67, 66]) characterise the response to shear directly from the data and correct the estimate, at the cost of additional shape measurements for each galaxy.

Radio interferometric imaging

Radio interferometers, unlike single-dish radio telescopes, use multiple radio antennas to perform observations. The main advantage is that the angular resolution of the telescope depends on the maximum separation between antennas instead of the size of the individual dish reflectors. This allows resolutions comparable to optical telescopes (arcsec) to be achieved without the need to build extremely large, expensive antennas that are challenging to deploy and maintain. Radio interferometers correlate the incoming radio signals between pairs of antennas. This allows, through the interferometry principle, to decompose the observed sky into sinusoidal components. The complex coefficients associated with these components are called visibilities, and their relationship with the observed sky is given by the van Cittert-Zernike theorem (Thompson et al. [74]).

where $\mathcal{V}(u,v,w)$ is the complex visibility function, $I_{\mathrm{sky}}$ is the observed sky intensity as a function of the two direction cosines $(l,m)$ and $n$ is the third direction cosine

If we assume that the observed sources do not subtend large angles in the sky, the small sky approximation is fulfilled

which simplifies Eq. (7) yielding

Neglecting the $w$ term introduces a distortion associated with projecting the celestial sphere onto a plane ([73]). Nevertheless, if the imaged region is sufficiently small and centred on the direction of observation, this effect can be ignored ([22]). We can rewrite Eq. (10) in terms of the two-dimensional (continuous) Fourier transform

However, since each pair of antennas samples the complex visibility function at a specific position in the $uv$ plane, which depends on their separation, known as the baseline, only a limited number of noisy $uv$ points are measured

If we represent the set of $uv$ sampling positions by a mask in Fourier space, we can describe the measured visibilities as follows:

where $V$ are the measured visibilities, $F$ is the discrete Fourier transform, $M$ is the sampling mask (containing ones at $uv$ positions with samples and zeros elsewhere), $\odot$ is the Hadamard product which represents the element-wise matrix multiplication, and $N$ is the additive Gaussian noise. Note that the noise, which is white in $uv$ space, is correlated through the sampling mask in real space.

Methodology

The proposed approach to estimate cosmic shear consists of forward modelling the visibilities of the observed galaxies rather than performing individual shape measurements on the galaxy images. The forward model allows us to evaluate the likelihood of the observed visibilities conditioned on the galaxy morphology parameters and the shear. Using Monte Carlo sampling methods, we are able to approximate the full posterior shear distribution. This approach aims to obtain an unbiased estimate of the shear without the need to measure galaxy shapes explicitly.

Generative forward model

The forward model consists of a generative galaxy model $G_\theta(z)$, a pixel-level shear transformation, and a model of the instrumental response of the radio interferometer. Figure 1 illustrates the forward model and provides examples of its intermediate outputs. This model is coded in JAX (Bradbury et al. [9]), a library for array-oriented numerical computations, with automatic differentiation and GPU compatibility and uses JAX-GalSim, a differentiable, GPU accelerated and scalable implementation of the GalSim (Rowe et al. [61]) library, for modelling the shear operation. This facilitates the use of sampling methods such as Hamiltonian Monte Carlo (HMC) or Microcanonical Langevin Monte Carlo (MCLMC), which require calculating gradients of the likelihood with respect to the sampled parameters.

Galaxy generative model

The generative model $G_\theta$ takes as input a set of galaxy morphology parameters, $z$, and produces an image of the associated galaxy in pixel space. We consider two types of generative models, each adapted to different morphological complexities.

Parametric generative model

This model represents the surface brightness profile of galaxies using a Spergel parametric profile (Spergel [70]). The advantage of this profile is that it has an analytical Fourier transform, making it ideal for modelling visibilities in $uv$ space. We use the implementation of this profile available in JAX-GalSim. In this model the galaxy parameters ($z$) that define the output image are: the half-light radius ($hlr$), which is related to the apparent size of the galaxy, the flux ($f$), which is related to the amplitude of the surface brightness profile, the Spergel index ($\nu$), which controls the degree of central concentration of the profile, and the intrinsic ellipticity ($e_1,e_2$). In this case, the generative model does not have any parameter $\theta$ that need to be learned or tuned other than the parameter priors, that will be discussed later. Figure 2 shows some examples of Spergel profiles in real space prior to being convolved with the radio PSF.

Vector (PDF)

Deep NN generative model

The second generative model is based on a Auto-Encoder (AE) neural network (NN) architecture, similar to that proposed

in Lanusse et al. [44]. We use a Deep Compression Auto Encoder (DC-AE) model with residual autoencoding layers, introduced in Chen et al. [17]. The AE architecture has two components: an encoder that compresses an input image into a much smaller vector known as the latent representation of the input image. And the second component is a decoder, which is trained to take the latent representation and reconstruct back the original image. This two components have a mirrored architecture and use multiple convolutional layers to compress and reconstruct back the input image. For an overview of recent autoencoder models, see Bank et al. [3]. In our case, we use an encoder with 5 hidden layers that use residual convolutional networks (ResNet, He et al. 2016) and 2D convolutional layers that reduce the size of the input image (128 × 128) to a latent variable of (4 × 4). The decoder has an equivalent structure but with pixel-shuffle up-convolution layers (Shi et al. [68]) to increase the dimensions of the latent representation to match the size of the original image. The AE is coded in equinox (Kidger & Garcia [42]), a JAX library for neural networks, which allows for direct integration into the forward model. The AE architecture and the training script are available on GitHub². Given the lack of observations of well-resolved galaxies in the radio band, we used optical observations. The AE is trained using 80 000 images of real galaxies with complex morphologies taken from the COSMOS (Scoville et al. [65]) survey, optical observations from the Hubble Space Telescope (HST). The dataset is split into 70 000 training examples and 10 000 test examples, and we incorporate random flipping of the galaxy images and their PSF to augment the data. Convolution with the HST PSF is applied at the decoder output so that the AE learns the unconvolved surface brightness profile of the COSMOS galaxies. The loss function used for training the AE is the mean absolute error (MAE) with total variation (TV) regularisation, and is written as follows

where $x$ is the real COSMOS galaxy image, $\hat{x}$ is the AE unconvolved galaxy reconstruction, $H$ is the corresponding HST PSF and $\nabla_i$ is the gradient with respect to the $i$-th coordinate, computed as the difference of adjacent pixels. The first term corresponds to the data fidelity, and the second term, weighted by $\lambda= 1\mathrm{e}{-5}$, is the total variation, which drives the denoising and smoothness of the unconvolved profiles. The classic Kullback–Leibler (KL) divergence penalty on $z$, used to enforce a Gaussian distribution in the latent space, is not included since it tends to degrade the quality of the images generated by the autoencoder. Instead, a softplus-type saturation, $z' = z/\sqrt{1+z^2/25}$, is used, matching the range of a Gaussian distribution, as suggested in Rozet et al. [62]. Once the AE is trained the decoder part can be used as a generative galaxy model, which is much more expressive than parametric profiles and can generate galaxies with well-resolved features. In this model, the galaxy parameters are replaced by the latent representation $z \in\mathbb{R}^{4 \times4}$. This allows us to generate galaxy images for different values of $z$. Figure 3 shows some examples of real COSMOS galaxies alongside the AE reconstruction, prior to convolution with the HST PSF and the respective residuals.

Vector (PDF)

Normalising flow

To generate new images of galaxies, we must sample new latent vector values. However, since no KL regularisation is used in the latent space, the distribution $p(z|X_{\mathrm{obs}})$ is unknown, making the sampling of new latent vectors a non-trivial task. In order to be able to draw samples in the latent space we use a normalising flow, which allows us to learn a bidirectional mapping between the target distribution $p(z|X_{\mathrm{obs}})$ and a base distribution, typically a standard Gaussian $u \sim\mathcal{N}(0,I)$. For a review of normalising flows, see Papamakarios et al. [52]. Since the flow is used in the context of differentiable probabilistic sampling—which requires forward modelling of several thousand galaxies and the calculation of gradients with respect to galaxy parameters—we need a flow that is fast in the forward direction and provides efficient gradient evaluation. The Real NVP flow (Dinh et al. [27]) is a type of coupling flow that exhibits computational symmetry and allows for efficient evaluation. The coupling layers split the latent vector in two (usually in half, i.e. $d = D/2$) and apply a bijection of the form

where $u$ and $u'$ are the input and output of the coupling layer, $s$ and $t$ are scaling and translation functions, and $\odot$ is the Hadamard or element-wise product. The functions $s$ and $t$ consist of multilayer perceptron (MLP) neural networks whose parameters are learned during the training of the flow. In this work, we use a Real NVP flow comprising four coupling layers, with a two-layer MLP of 128 neurons and ReLU activation. We use the FlowJAX coupling flow implementation (Ward [77]). The flow is trained by minimising the KL between the true prior distribution $p(z|X_{\mathrm{obs}})$ and the variational distribution $q(z|X_{\mathrm{obs}})$. In practice, this is achieved by minimising the negative log-likelihood over the variational distribution. We train the flow on 80 000 latent representations of the COSMOS dataset obtained with the encoder model of the trained AE. The flow is incorporated into the forward model right before the decoder, thus, instead of sampling the latent space $z$, we sample from a Gaussian distribution $u \sim\mathcal{N}(0,1)$ and obtain $z=f(u)$ using the flow.

Instrumental response

Once the galaxy image has been generated and shearing is applied using JAX-GalSim, the instrumental response is applied in Fourier space. The radio PSF is modelled as described in Eq. (13), using a sampling mask ($M$) in $uv$ space. This set of samples constitutes the simulated visibilities. We simulate the $uv$ sampling mask with argosim (Centofanti et al. [14])³, a modular, lightweight, and JAX-based package for radio interferometric simulations. We use SKA1-Mid AA4 as our reference interferometer. We calculate the $uv$ tracks for an 8-hour observation, sampling every 5 minutes. We simulate an observation at zenith, which corresponds to a declination of approximately $-30$ degrees. We simulate the $uv$ tracks for a single frequency band centred at 1.4 GHz. The resulting $uv$ coverage is shown in Figure 4. In this work, we simulate a single galaxy per FOV and assume that all the simulated galaxies are located at the same pointing centre, using the same radio PSF for all simulated galaxies. This should not affect the results, as the $uv$ mask is deterministic, but it does greatly simplify data generation. In reality, effects such as radio frequency interference (RFI) or calibration errors can alter the radio PSF, making it not fully deterministic. However, these details are not addressed in this work. The $uv$ samples are gridded to 128 $\times$ 128 pixels to match the size of the simulated galaxy images, using a pixel scale of 0.15 arcsec pixel$^{-1}$. At this pixel scale, the $uv$ samples associated with the longest baseline fit just within the 128 $\times$ 128 image.

Posterior inference via differentiable sampling

Forward modelling the visibilities allows us to model the likelihood of the observed data. The model for the observation of a single galaxy is written as

where $M$ is the $uv$ sampling mask, $F$ is the Fourier transform, $\gamma\otimes\cdot$ represents the shear operation and $N$ is noise added to the sampled visibilities. The noise in $uv$ space is Gaussian and independent per visibility. Then, for the observation of $N_{\mathrm{gal}}$ galaxies, the likelihood is written as follows

where $V_{\mathrm{obs},i}[j]$ is the $j$-th visibility sample of the $i$-th observed galaxy, and $\sigma_j$ is the noise associated with the $j$-th baseline. Finally, we can write the joint posterior for the galaxy parameters and the shear in terms of the likelihood using Bayes’ theorem, resulting in

Since the galaxy morphology parameters and shear are uncorrelated, we can factorise the prior as the product of two independent priors, $p(z,\gamma)=p(z)p(\gamma)$.

Monte Carlo sampling

To estimate the posterior distribution, we use Monte Carlo sampling methods. Given that we are dealing with multiple galaxies simultaneously, the number of morphological parameters is considerably large (over 700 parameters for $N_{\mathrm{gal}}=100$), making the use of conventional sampling methods such as Metropolis-Hastings unfeasible. We therefore use Microcanonical Langevin Monte Carlo (MCLMC), a gradient-based sampling method that is highly efficient for sampling in high-dimensional spaces [60]. MCLMC automatically tunes its hyperparameters and is orders of magnitude faster than other gradient-based methods such as Hamiltonian Monte Carlo (HMC, [6]). These algorithms model the trajectory of the sampled particles using dynamics inspired by physical systems. In particular, they use Hamiltonian dynamics, which utilise the gradients of the target distribution to simulate the trajectories of the sampled particles. We exploit the fact that our forward model is built in JAX meaning that the likelihood is fully differentiable with respect to the sampled parameters. We use the MCLMC implementation provided by the blackjax Bayesian inference library [13]. To increase the number of independent samples, we run multiple chains in parallel. All chains are initialised randomly from points drawn from the prior distributions of the tracked parameters. Then, before running the MCLMC adaptation and sampling, we perform a maximum a posteriori (MAP) estimation for each chain using gradient descent. After sampling, we verified the convergence of the chains using the Effective Sample Size (ESS) statistic, which measures the number of uncorrelated samples for each of the tracked parameters.

Kernel density estimation

Once the Monte Carlo chains have taken a sufficient number of independent samples from the posterior distribution, we proceed to estimate the probability density. We used a simple Gaussian mixture model (GMM) after verifying that the distribution of the posterior samples is fairly round and unimodal. The posterior density is then written as follows

with $K$ Gaussian components. We use the scikit-learn ML library [55] to fit the GMM parameters, fixing the number of components at $K=5$. To combine multiple posterior distributions, in cases where the posterior has a single mode and a Gaussian-like shape, we perform the analytical multiplication of the likelihoods, dividing $N-1$ times with the prior $(\Sigma_p)$, yielding

Radio simulations

Given the absence of radio observations with sufficient resolution and galaxy density to perform a WL-scale analysis, in this study we apply the proposed methodology to simulated data. Furthermore, this allows us to evaluate the performance of our method in a more controlled setting, given that we know the ground truth value of the applied shear. We simulated two datasets of galaxy observations with varying degrees of morphological complexity. In the first dataset each galaxy is represented by a single component parametric profile, and in the second dataset we turn real optical images of galaxies from the COSMOS survey into radio observations [65].

Parametric radio galaxies

The parametric radio galaxies are modelled by Spergel profiles using JAX-GalSim. Spergel profiles can be described analytically both in real space and Fourier space, and offer greater flexibility than exponential or Gaussian profiles. These profiles are characterized by two main parameters: their size, represented by the half-light radius (hlr), and the Spergel index $\nu$. The form of the Spergel profile in the 2D-projected Fourier space is as follows

where $c_{\nu}$ is a tabulated constant that depends on $\nu>-1$. In real space, this profile is described as follows

where $f_{\nu}(u)=(u/2)^{\nu}K_{\nu}(u)/\Gamma(\nu+1)$ and $K_{\nu}(u)$ is a modified spherical Bessel function of the third kind. The analytical expressions of this profile are described in more detail in Spergel [70]. The Spergel profile can approximate an exponential profile for $\nu=0.5$, equivalent to a Sérsic index $n=1$, and a de Vaucouleurs-type profile for $\nu=-0.6$, equivalent to a Sérsic index $n=4$. In Tan et al. [72] the authors present a comparison between the Sérsic and Spergel profiles and provide an analysis of the ranges of values for the Sérsic index $n$ and the Spergel index $\nu$ that result in similar surface brightness. We believe that, for this study, the type of profile used is not a determining factor, as they share a similar morphological complexity. An important factor in our preference for the Spergel profile over the Sérsic profile is that the former has a differentiable JAX-based implementation in JAX-GalSim, which is essential for constructing a fully differentiable forward model.

Once the Spergel profiles have been generated, we assign an intrinsic ellipticity and flux to each individual galaxy, and apply the same fiducial shear to all the galaxies. In this study, we tested various shear values, which we discuss later in Sect. 6. For the population of the parametric dataset, we used the Tiered-Radio Extragalactic Continuum Simulation (T-RECS) presented in Bonaldi et al. [8]. These simulations provide a catalogue of star-forming galaxies that matches the flux, size, shape, redshift, among other sources parameters, to recent observations. As in Tripathi et al. [75] we used a selection cut in the integrated flux, ranging from 50 µJy to 200 µJy. Figure 5 shows the joint distribution of flux and size for the selected galaxies. We randomly sample independent flux and hlr values from exponential and log-normal distributions fitted to the T-RECS catalogue. The exponential and log-normal fits are shown in red and blue, respectively, over the marginal histograms in Figure 5. With regard to the Spergel index $\nu$, it is sampled from a logistic-normal⁴ distribution (not to be confused with log-normal) scaled to $\nu_{\max}=1$ and $\nu_{\min}=-0.7$, limiting the Spergel index to its valid range of application. Finally, the intrinsic galaxy ellipticities are sampled form independent distributions $e_i \sim\mathcal{N}(0,\sigma_e)$ with $\sigma_e=0.2$. To ensure the numerical stability of the GalSim shear operation, we limit the ellipticity to a unit modulus.

The simulated galaxies are drawn in Fourier space, where the $uv$ mask described in Sect. 4.1.2 is applied to obtain the visibilities. Galaxies are simulated separately, each yielding its own $uv$ set. Next, we add independent noise to each visibility to obtain the noisy observations. The signal-to-noise ratio (S/N) is defined as follows

Vector (PDF)

The noise level is adjusted to achieve a S/N ranging from 10 to 80. In this study, we use the same noise level for all the visibilities. The noise level used for the parametric galaxies is $\sigma_{uv}=0.02$.

Realistic radio-like galaxies

We constructed a second dataset containing more complex galaxy morphologies using observations of real galaxies from the COSMOS survey. The COSMOS galaxies consist of observations taken by the Hubble Space Telescope (HST) in the optical band. Given that optical telescopes have higher resolution and sensitivity than radio telescopes, they are able to detect a greater number of well-resolved galaxies. In our efforts to consider complex galaxy morphologies, we artificially adapted these galaxies to radio-like observations. In Patel et al. [54], the authors analyse the correlation between the shapes of galaxies observed using both optical telescopes (HST) and radio interferometers (VLA). They conclude that, on a galaxy-by-galaxy basis, optical and radio shapes are weakly correlated; however, the shape distributions in the optical and radio sky were very similar. We use the COSMOS 25.2 catalogue available on GalSim, which contains approximately 87 000 galaxies with the corresponding original HST PSF. We applied a magnitude cut at 24.0, excluding large galaxies with highly resolved features, resulting in a sample of approximately 40 000 galaxies. The COSMOS galaxies are co-added images, that is, they are a combination or stack of multiple observations, which results in a non-trivial noise distribution. Furthermore, as we need to construct a mock dataset of sheared galaxies and it is not preferable to shear the noisy images, we whitened⁵ the COSMOS galaxies. We used the trained AE described in Sect. 4.1.1 to reconstruct the unconvolved surface brightness profile of the COSMOS galaxies. Then, we applied a fiducial shear to the whitened galaxies and drew them in Fourier space with a size of 128×128 pixels, using the pixel scale of the COSMOS images (0.03 arcsec pixel$^{-1}$). Finally, we obtain the complex visibilities by applying the radio PSF in Fourier space as described in Sect. 4.1.2. In this dataset, we use two different pixel scales: one for generating the COSMOS galaxies and another for the radio PSF (0.15 arcsec pixel$^{-1}$). This artificially shrinks the radio PSF by a factor of 5, so that the radio adapted COSMOS galaxies are well resolved. This is because the simulated radio PSF is much larger than the HST PSF. To model observational noise, we add independent Gaussian noise to the $uv$ samples as described in Sect. 5.1. The noise level used for the realistic radio–like galaxies is $\sigma_{uv} = 0.01$ (which is not directly comparable to the parametric dataset as the parametric and autoencoded galaxies have different fluxes), resulting in an S/N range of 20 to 60. The S/N distribution is shown in Figure 3 for both the parametric and the realistic galaxies. The latter contain no galaxies of very low S/N (between 10 and 20), however, the distribution is consistent with that of the parametric galaxies. An example of a COSMOS galaxy adapted to a radio-like format is shown in Figure 7, which displays the absolute values of the gridded visibilities and its Fourier counterpart, the galaxy dirty image.

Results

In this section, we present the results of the shear posterior estimation. First, we show the results obtained using the parametric model and compare them with other state-of-the-art methods for shear measurement. We then demonstrate the model bias that can occur when considering galaxies with more complex morphologies than those represented by the adopted generative model. Finally, we address this bias by using a more complex generative model, which is capable of representing the morphology of the observed galaxies.

Parametric galaxies

We estimate the shear posterior distribution using the forward model and the MCLMC sampling algorithm for different fiducial shear values. We selected nine ground-truth points on the $(\gamma_1,\gamma_2)$ plane to obtain equidistant values between $-0.05$ and $0.05$ for each shear component. For each fiducial shear, we simulated the visibilities of 10 000 parametric galaxies, as described in Sect. 5.1. The number of galaxies used corresponds to the galaxy density expected for the first phase of SKA (SKA1), $n_{\mathrm{gal}} \sim2.7$ arcmin$^{-2}$ ([32]), over an area of 1 deg$^2$, approximately the primary beam of the MeerKAT antennas, which gives roughly 10 000 galaxies. In this test, we use the GalSim-based Spergel generative model, whose morphology is consistent with the simulated observations. The sampled parameters and their prior distributions are listed in Table 1.

Parameter	Prior distribution	Size	Units
$h_{\mathrm{lr}}$	LogisticNormal(0.1, 3.0)	$(N_{\mathrm{gal}},)$	arcsec
$\nu$	LogisticNormal($-0.7$, 1.0)	$(N_{\mathrm{gal}},)$	-
flux	LogisticNormal(0.03, 0.25)	$(N_{\mathrm{gal}},)$	mJy
$(e_1,e_2)$	$\mathcal{N}(0, 0.2)$	$(N_{\mathrm{gal}},2)$	-
$(\gamma_1,\gamma_2)$	$\mathcal{N}(0, 0.1)$	$(1,2)$	-

Since the MCLMC sampling algorithm requires the gradients of the forward model to be calculated for each simulated galaxy, we must process the 10 000 galaxies in smaller batches due to GPU memory constraints. We use a single Nvidia Tesla V100 GPU with 32 GB of memory, and the forward modelling of the galaxy images, together with the gradient computation graphs, can easily exceed the available memory for a few hundred galaxies. We therefore sample the posterior distribution using batches of 100 galaxies at a time (giving a total of 100 batches). Sampling takes around 10 minutes, using 10 chains in parallel, taking 200 000 samples and yielding a total ESS of around 4 000 samples for the shear parameters. As we divide the dataset into 100 batches of 100 galaxies each, we can parallelise the runs across multiple GPUs, achieving a performance of approximately 60 ms per galaxy. In this paper, we do not seek to optimise the performance of the algorithm, but simply to demonstrate its ability to produce an unbiased estimate of cosmic shear. After obtaining the samples for a given batch, we estimate the posterior density by fitting a GMM. Finally, we combine the GMMs from every batch to estimate the posterior density for the entire dataset. The posterior distributions have a single mode and a Gaussian-like appearance, so they are combined by analytically multiplying their densities. Since the posteriors are likelihood-dominated, the precision of the prior accounting for approximately 1% of the posterior precision, we directly multiply the posteriors without subtracting the prior.

Figure 8 shows the results of the shear estimation. In each panel, we present the estimation error for each shear component as a function of the different fiducial values. For each shear value, we calculate the mean and standard deviation of the GMM derived from the 10 000 observations. The error bars represent the $1\sigma$ and $3\sigma$ intervals. For reference, we show the shear estimation error of the DeepShape method ([75]), which uses convolutional neural networks to estimate the individual ellipticity of each galaxy and then estimates the shear as the average of these ellipticities. The results of the galaxy-by-galaxy ellipticity estimates for various fiducial shear values were provided by the authors of DeepShape. We have reproduced the Figure 9 of their paper, simplifying the data points to a linear interpolation to improve readability. In [75], the authors also use 10 000 images of parametric galaxies with an equivalent radio PSF, size and flux distribution. The difference between their dataset and ours lies in the type of profile used, as DeepShape simulates galaxies using Sérsic profiles. However, we consider that this difference has no significant impact on the results, as both models produce similar galaxy morphologies for the chosen parameter ranges, and use the same $uv$ coverage, and number of galaxies. Comparing the two methods, we note that our shear estimate is more accurate, with an uncertainty approximately three times smaller for the same number of galaxies. Furthermore, our method shows no apparent signs of significant multiplicative bias.

Vector (PDF)

Posterior calibration

Sampling the posterior probability naturally enables uncertainty quantification over the shear estimates. However, the process does not guarantee that the estimated density is well calibrated. To assess the quality of the estimated shear posterior, we perform a probability coverage test for a single reference shear value. The selected shear is $\gamma_{\mathrm{GT}} = (-0.05, 0.05)$. Since we are dealing with simulations, we know the expected value of the posterior probability; we therefore carry out an analysis based on the distribution of the high posterior density (HPD) credible regions [20, 71]. For the chosen ground-truth shear value, we estimate the posterior density using our framework a total of $N$ times. Then, for each estimated posterior $\hat{p}_n(\gamma\mid D)$, we find the interval of highest density that contains the ground-truth shear value,

We then compute the credible level $C_n$ by integrating the posterior density over this interval, as follows

The theoretical distribution of $C_n$ for a well-calibrated posterior is uniform between 0 and 1. Therefore, for a given nominal confidence level $\alpha\in[0,1]$, the fraction $K/N$ of posteriors for which the ground-truth shear falls within the $\alpha$-credible region, that is, $C_n \le\alpha$, should be equal to $\alpha$. The number of posteriors for which the ground truth value falls inside the $\alpha$-credible region, $K$, is distributed as a Binomial$(N,\alpha)$, with mean $N\alpha$ and variance $N\alpha(1-\alpha)$. Therefore, the expected value of the ratio $K/N$ is $\alpha$, and its standard deviation is

We performed the coverage test for $N = 125$ posterior estimates and present the results in Figure 9, which shows the empirical coverage ratio for different values of the credibility level $\alpha$. The diagonal line $\alpha:\alpha$ represents the ideal coverage, and

we have outlined the $2\sigma$ band in grey. We note that the empirical coverage, shown in blue, is consistent with the theoretical coverage, indicating that the estimated posteriors are properly calibrated. However, it should also be noted that it systematically lies below the ideal coverage, which we will analyse in the future with a larger number of simulations.

Realistic galaxies

In this section, we test the forward model on galaxies with more complex morphologies, evaluating shear estimation in a scenario where the generative model is unable to capture the morphological complexity of the observations. We use the COSMOS real galaxy dataset adapted to radio interferometric observations, which is described in Sect. 5.2. We estimate the shear posterior for a total of 10 000 whitened COSMOS galaxies. As we are solely seeking to demonstrate the existence of model bias in the case of morphological mismatch, we perform this test for a single fiducial ground-truth shear value of $\gamma_{\mathrm{GT}} = (-0.05, 0.05)$.

Spergel generative model

First, we evaluate the forward model using the parametric generative model, which models galaxies using a Spergel profile. We subdivide the dataset into batches of 100 galaxies to ensure they fit in the GPU memory, and we sample the posterior using MCLMC, as described in Sect. 4.2.1. Given that the size and range of the COSMOS profiles do not follow the same distribution as the T-RECS dataset, we use different priors, which are shown in Table 2. Once the posterior samples have been obtained, we estimate the density using a GMM. Finally, we analytically combine the posteriors across the multiple runs. The posterior shear estimate from 10 000 galaxies using the Spergel generative model is shown in violet in Figure 10. The black cross indicates the ground truth shear value. The constraining power of the parametric model is similar to the results shown in Sect. 6.1, however, in this case we observe a considerable bias, larger than $10\sigma$. This bias is to be expected, as Spergel profiles are unable to capture the complex galaxy morphologies present in the COSMOS dataset.

Parameter	Prior distribution	Size	Units
h$\mathrm{lr}$	LogisticNormal(0.1, 5.0)	$(N_{\mathrm{gal}},)$	arcsec
$\nu$	LogisticNormal($-0.7$, 1.0)	$(N_{\mathrm{gal}},)$	-
flux	LogisticNormal(0.005, 1.0)	$(N_{\mathrm{gal}},)$	mJy
$(e_1,e_2)$	$\mathcal{N}(0, 0.2)$	$(N_{\mathrm{gal}},2)$	-
$(\gamma_1,\gamma_2)$	$\mathcal{N}(0, 0.1)$	$(1,2)$	-

Vector (PDF)

NN-based generative model

To address this bias, we consider the AE-based generative model, which does capture complex galaxy morphologies. We repeated the shear inference using 10 000 radio-adapted COSMOS galaxies. The galaxy parameters in this model correspond to the latent vector $z \in\mathbb{R}^{4 \times4}$ of the AE. In order to sample in the AE latent space, we use the normalising flow as described in Sect. 4.1.1. The sampled parameters and their priors are listed in Table 3. Sampling with MCLMC requires calculating the gradients of the forward model with respect to the sampled parameters. Given that the forward model comprises a NN-based flow and generative model, such gradients require considerable GPU memory and are slower to compute than in the case of a simple parametric model. In order to run the MCLMC sampling of 100 galaxies on a single Nvidia Tesla V100 GPU with this model, we carried out various performance optimisation efforts.

Parameter	Prior distribution	Size	Units
$u$	$\mathcal{N}(0,I)$	$(N_{\mathrm{gal}},4,4)$	-
$(\gamma_1,\gamma_2)$	$\mathcal{N}(0,0.1)$	$(1,2)$	-

The simplest solution is to use a smaller generative model. However, it is difficult to successfully train a small autoencoder model with few parameters that is able to capture the complex structure of the target images. Then, the chosen approach is to use a technique known as knowledge distillation (Buciluă et al. [12]; Hinton et al. [37]). This technique involves using a pre-trained model (the decoder part of the AE) as a teacher for a considerably smaller model, the student. This task is simpler than training the entire AE model, as the student model only needs to learn the mapping from the latent space to the surface brightness profile of the galaxies and not the full encoding and decoding task. We used a student model with the same number of layers but a depth (number of channels) that is a quarter of that of the teacher model, reducing the number of parameters from 18 million to 790 thousand. This significantly reduces the memory footprint of the gradients and also speeds up sampling, as the forward pass of the distilled model is 10 times faster. To speed up the forward pass even further, we recompiled the distilled model using half-precision, with float16 instead of float32. In the tests we carried out, the reduction in numerical precision did not significantly impact simulated visibilities ($< 0.02\%$ relative error) or the computed gradients ($< 1\%$ relative error). Finally, we limit the padding size used by JAX-GalSim to interpolate galaxy images and plot them in Fourier space. These optimisations allow us to use MCLMC sampling with batches of 100 galaxies, limiting the number of parallel chains to four. We took a total of 5 000 samples, resulting in an ESS of around 200 samples for the shear parameters. Sampling takes 50 minutes, approximately 300 ms per galaxy, by running the 100 runs in parallel across multiple GPUs.

Once we have sampled the posterior for each batch of 100 galaxies, the posteriors from each run are combined by analytically multiplying the estimated densities. The results of the shear posterior inference are shown in orange in Figure 10. As can be seen in this case, the shear estimate is not biased, with the ground-truth shear value lying close the edge of the $2\sigma$ region. The ability of the NN-based generative model to represent the observations makes it possible to address the model bias observed in the parametric case. We also note that the posterior estimate is slightly larger than in the parametric case. This is because the NN-based generative model has greater flexibility, and there are more combinations of latent parameters and shear that adequately represent the observed data.

Scalability and performance

Surveys such as SKA2 promise to observe a high galaxy density, over $10\,\mathrm{arcmin}^{-2}$ ([32]). Given their extensive sky coverage, it is expected that several trillion galaxies will be detected. Therefore, producing shear maps of large regions of the sky requires processing a very large number of galaxies. Estimating the full shear posterior using probabilistic sampling techniques allows us to quantify the uncertainty of the estimate and marginalise over the latent parameters; however, it is slow compared to point estimate methods. Moreover, estimating the full posterior becomes considerably slower the more complex the forward model is. As an example, the proposed parametric model takes around 10 minutes to get 2 500 effective shear samples for a batch of 100 galaxies, whilst the NN-based model takes 50 minutes for a much smaller number of effective samples. This performance limitation is due not only to the time taken by the model to simulate the observations, but also to the time taken to compute the gradients of the model with respect to the sampled parameters. With the aim of assessing the scalability of our method for future surveys, we investigate an alternative to full posterior estimation. In this section, we carry out a test in which we use the forward model to predict a point estimate of the shear using the same observations as those considered in the previous case. The approach chosen as the point estimation method consists of posterior density maximization (MAP).

Unlike the MAP optimisation used to initialise the sampling algorithm, in this case we want to ensure convergence to a maximum of the posterior density, where in the previous case where it was sufficient to be close to a likely region in order to carry out the sampling. We have observed experimentally that initialising the latent parameters to zero yields better results. We believe this is due to the flexibility of the AE model in representing the data; consequently, initialising at a random point in the latent space can lead to convergence at local maxima. We initialise the shear randomly, sampling it from the prior. We then perform four MAP optimisations in parallel to verify that they converge to the same shear value. We use a learning rate $lr = 0.05$ to ensure convergence, and applied a 2× factor to the $lr$ related to the shear parameters. The aim is for the optimisation to adjust the shear parameters rather than searching for latent parameters that explain the observations for a given initial shear. We run the MAP optimisation for a total of 1 500 gradient descent steps using the Adam optimiser from the optax library. As in the previous tests, due to GPU memory limitations, the dataset is divided into batches of 100 galaxies, obtaining 100 shear point estimates.

The 100 runs are executed in parallel, taking a total of 6 minutes, compared to almost 1 hour for the full posterior. The MAP estimation performance can be further optimised by running just one MAP per run, reducing the estimation time by a factor of four, resulting in a processing time of 9 ms per galaxy.

The results obtained from the 100 runs are combined, and the mean is shown in Figure 11 in green, alongside the full posterior estimate from the previous test. We note that the MAP estimate agrees with the estimated posterior and is as close to the ground truth as the maximum of the estimated posterior. The MAP estimate does not coincide with the maximum of the estimated posterior distribution, as the posterior shown here is marginalised with respect to the galaxy parameters, whereas the MAP is optimised over the joint posterior distribution. One of the main limitations of this method is that we do not have a direct quantification of the uncertainty in the estimated shear. Furthermore, as we are not marginalising over the forward model's galaxy parameters, any degeneracy between this parameters and the shear will potentially introduce a bias. Finally, in Refregier et al. [57], the authors show that maximum likelihood estimation (MLE) methods for template fitting galaxy shapes from optical images suffer from noise bias, leading to biases in cosmic shear estimation. Our scenario is different, as the likelihood is modelled in Fourier space and therefore the noise does not interact in the same way with the shape of the galaxy; however, further tests must be carried out to guarantee an unbiased point estimate of the shear.

Vector (PDF)

Conclusions and perspectives

Conclusions

In this paper we propose a shape-measurement-free method for estimating cosmic shear, based on a forward model for simulating galaxy observations. The model consists of three components: a generative model of galaxy images, the distortion introduced by cosmic shear, and the instrumental response of the radio interferometer. For a given combination of galaxy parameters, cosmic shear, and radio PSF, the model simulates the radio visibilities in the $uv$ plane, which are then compared with the observations. Using a probabilistic sampling algorithm, the posterior density of the shear parameters can be estimated conditioned on a set of observed galaxies. Since galaxy shape and cosmic shear are inherently degenerate, observations of multiple galaxies from the same sky region are used jointly, under the assumption that they all share the same shear value. We employ two types of generative models: one based on the parametric Spergel surface brightness profile, whose morphological parameters are the half-light radius (hlr), flux, Spergel index $\nu$, and the intrinsic ellipticity of the galaxy; and a NN-based generative model, consisting of the decoder part of an autoencoder architecture trained on real galaxy images (COSMOS) from the Hubble Space Telescope, which encodes galaxy morphology directly in the latent space of the autoencoder. Our forward model is implemented in JAX and is fully differentiable, which allows the use of gradient-based sampling methods such as MCLMC, which are highly efficient when sampling a large number of parameters.

Since real radio observations with sufficient resolution and galaxy density for WL studies are not yet available, we test our method on simulated radio galaxy visibilities. In a first evaluation, we use parametric profiles in real space with a parameter distribution based on the T-RECS catalogue of radio star-forming galaxies, as employed in other state-of-the-art works on galaxy shape measurement in radio interferometry. We use a radio PSF based on the SKA-Mid antenna layout for an 8-hour observation at zenith. We demonstrate that our framework is capable of estimating the shear for a total of 10 000 simulated galaxy visibilities, and compare the results with other state-of-the-art methods such as DeepShape. Our forward model provides an estimate with lower uncertainty and no appreciable bias. Furthermore, we assess the quality of the estimated posterior by performing a coverage test, which shows that the posterior is correctly calibrated.

We then test the parametric forward model on simulations of galaxies with complex morphologies. For these simulations, we use an autoencoder trained on real galaxy images from the HST COSMOS survey to generate realistic surface brightness profiles. We simulate the visibilities of 10 000 galaxies sheared by the same amount, and estimate the shear using the parametric forward model. We show that the resulting shear estimate is biased, due to the inability of the parametric model to capture complex galaxy morphologies. To address this, we repeat the same test replacing the parametric forward model with the NN-based generative model, and show that accurately representing the morphology of the target galaxies leads to an unbiased cosmic shear estimate.

Finally, we used the NN-based forward model to perform a point estimate of the shear using a MAP optimisation. Obtaining a point estimate is considerably faster than a full posterior estimate but lacks uncertainty quantification. Furthermore, it may potentially be biased due to degeneracies between the galaxy parameters and the shear, in which case an additional calibration step will be required. We show that we were able to use the proposed forward model to obtain a point estimate of the shear that is consistent with the sampled posterior distribution. This method is up to 30 times faster than sampling the full posterior. Further tests are needed to properly assess the bias of the point estimate.

Future work

In this work we demonstrate the ability to obtain a posterior estimate of cosmic shear from radio data by forward modelling the observed galaxies. To this end, we consider several simplifications. We do not address galaxy detection, which may introduce a bias due to the preferential selection of galaxies with certain characteristics. Galaxies are simulated individually, centred on postage stamps, rather than simulating the $uv$ samples for all the galaxies distributed across a single FOV. We use the same radio PSF for all observations. We do not account for the spectral energy distribution (SED) of the observed galaxies when simulating visibilities across different frequency channels. Furthermore, we do not take into account the possible intrinsic alignment of background galaxies with their hosting LSS, which may correlate the shapes of galaxies even in the absence of shear. These considerations remain to be explored in future work. This work is a demonstration of the method using radio simulations, and represents an initial step towards a shear inference environment (SHINE⁶) that we are currently developing. Ultimately, SHINE will forward model galaxy observations across the entire FOV, jointly fitting the total number of galaxies and their positions, using a spatially and spectrally varying PSF, and accounting for the spectra of the observed galaxies, while simultaneously estimating the shear and the morphological and physical parameters of the galaxies, such as their metallicity or star formation rate, which are related to their spectra. SHINE will make use of both radio and optical observations, with the ability to jointly estimate the posterior of the target parameters using data from multiple surveys.

Acknowledgements. This work was granted access to the HPC resources of IDRIS under the allocation AD01041410R3 and AD011012983R4 made by GENCI. This work was supported by the European Community through the grant ARGOS (contract no. 101094354). This work was supported by the TITAN ERA Chair project (contract no. 101086741) within the Horizon Europe Framework Programme of the European Commission, and the Agence Nationale de la Recherche (ANR-22-CE31-0014-01 TOSCA). We would like to thank Priyam-vad Tripathi for sharing the data used to reproduce the DeepShape results plot.

Data Availability. The scripts and notebooks for reproducing the results and figures in this paper are available at: https://github.com/CentofantiEze/Radio_WL_Generative_Forward_Model. The trained models and metrics are available at https://zenodo.org/records/21131669.

References