© The Authors 2026

1 Introduction

Gravitational lensing is the deflection of electromagnetic radiation from distant galaxies by intervening large-scale structure. In the weak-lensing regime, this effect manifests as a small but coherent distortion of galaxy shapes across the sky, commonly referred to as cosmic shear. Weak gravitational lensing provides a powerful probe of the matter distribution and the growth of cosmic structure, and plays a central role in contemporary cosmology (Bartelmann & Schneider 2001; Copeland et al. 2006; Clifton et al. 2012). The shear signal is encoded in the observed shapes of faint galaxies, making its extraction fundamentally dependent on both highly accurate shape measurements and large galaxy samples in order to suppress intrinsic shape noise.

To date, the majority of cosmic shear measurements have been carried out in the optical waveband (Tyson et al. 1990; Bacon et al. 2000; Van Waerbeke et al. 2000; Wittman et al. 2000), driven by the high angular resolution, sensitivity, and source densities achieved by optical surveys. In this context, a mature ecosystem of shape measurement techniques has been developed and extensively validated. Weak-lensing has also been detected at radio wavelengths (Chang et al. 2004), and through radio-optical cross-correlation analyses (Demetroullas & Brown 2015, 2017), thus demonstrating the viability of radio data for cosmic shear studies. However, despite these successes, shape measurement methods in the radio domain remain significantly less mature than their optical counterparts, and have so far seen more limited adoption in precision cosmological analyses.

The upcoming generation of radio surveys, particularly those enabled by SKA-Mid, will deliver the sensitivity, angular resolution, and galaxy number densities required for competitive weak-lensing analyses (Brown et al. 2015). Radio observations also offer several complementary advantages, including a deterministic point spread function and access to higher-redshift galaxy populations. Together, these developments place increasing demands on radio-specific shape measurement methodologies and motivate renewed efforts to develop robust and scalable approaches that are explicitly designed for radio interferometric data.

Radio interferometric imaging, however, differs fundamentally from optical imaging and introduces additional challenges for weak-lensing shape measurement. In radio interferometry, images are reconstructed from incomplete Fourier measurements, leading to correlated noise and reconstruction artefacts that depend on the sampling pattern and deconvolution strategy. As a result, radio weak-lensing is particularly sensitive to errors introduced during the imaging stage, and residual reconstruction errors can propagate directly into biased estimates of galaxy morphology. Accurate shape measurement in the radio therefore depends critically on the fidelity and stability of the reconstruction process.

To mitigate these effects, early radio weak-lensing studies focused on fitting parametric models directly in the visibility domain, thereby avoiding biases introduced by image reconstruction. While conceptually appealing, visibility space approaches are computationally demanding and rely on simplified parametric descriptions of the sky brightness distribution (Chang et al. 2004; Patel et al. 2010; Rivi et al. 2016). Image plane methods offer a more computationally efficient alternative and are better suited to large survey volumes, but their performance is inherently limited by the quality of the reconstructed images, particularly for faint and extended sources (Patel et al. 2015). These limitations become increasingly severe as observations move towards wider fields and higher source densities.

In a companion paper Tripathi et al. (2025) (hereafter Paper I), we introduced DeepShape, a fully supervised deep learning framework for radio weak-lensing shape measurement in the image domain. DeepShape combines a plug-and-play (PnP) image reconstruction stage with an equivariant convolutional neural network (CNN) for shape estimation, and was shown to achieve accuracy comparable to the state-of-the-art visibility domain methods while retaining the computational efficiency of image-based approaches. However, this analysis was limited to isolated sources, and did not address the effects of source entanglement present in realistic wide field observations.

In practice, radio interferometers observe wide fields containing large numbers of galaxies simultaneously. A single SKA-Mid pointing is expected to cover an area of order 1 deg² and include up to 10⁴ galaxies with a sufficient signal-to-noise ratio for shape measurement, together with a substantially larger population of fainter sources (Bacon et al. 2020). All of these objects contribute to the measured visibilities, making the isolation of individual galaxy signals inherently non-trivial. As a consequence, shape measurement of individual galaxies requires either explicit source separation based on prior information about source locations or the use of computationally intensive joint parametric modelling approaches.

For this paper we extended the DeepShape framework to wide field radio interferometric observations by explicitly addressing the challenge of source separation. Our approach combines a faceting-based reconstruction strategy that isolates the contribution of individual sources from the full visibility dataset with the existing PnP deconvolution algorithm. To further suppress contamination from neighbouring emission, we introduced a variational autoencoder (VAE)-inspired network enhanced with multi-head attention (MHA) layers, trained specifically to deblend overlapping sources in the image plane. The final stage applies the previously developed equivariant shape measurement network with an encoder to correct residual point spread function (PSF) artefacts. By integrating faceting and learned deblending within a single end-to-end pipeline, this framework enables robust and scalable shape estimation in crowded radio fields, and is well-suited to the data volumes and complexity expected for next-generation radio weak-lensing surveys.

The paper is organized as follows. In Sect. 2 we provide a brief overview of the radio interferometric imaging process and the challenges associated with source separation. Sect. 3 details our simulations and the data processing steps required to create the training and testing datasets. In Sect. 4 we describe the three key components of our framework: the facetingbased image reconstruction, the VAE-inspired deblender, and the shape measurement network. Our results are presented in Sect. 5, where we demonstrate the performance of the proposed pipeline in realistic wide field scenarios. Finally, Sect. 6 summarises our conclusions and outlines potential avenues for future work.

2 Theory

2.1 Radio interferometric imaging

Radio interferometers do not form direct images of the sky. Instead, they sample its complex Fourier transform at discrete spatial frequencies using arrays of antennas. These samples, known as visibilities, correspond to the correlated response of an antenna pair (baseline) to the incident electric field, integrated over time and frequency. The measurement process is described by the radio interferometric measurement equation (Smirnov 2011), $$\begin{array}{rl}& V(u,v,w)=\iint \frac{I_T(l,m)}{n}{e}^{-2\pi j(ul+vm+w(n-1))}\text{d}l\text{d}m\\ & n=\sqrt{1-l^2-m^2},\end{array}$$(1)

where (u, v, w) denote the baseline coordinates, (l, m) are direction cosines on the sky, and IT (l, m) is the true sky brightness distribution. Under the small field of view approximation, this relation reduces to a linear Fourier measurement model, $$\mathbf{v}\approx G^{†}F{\mathbf{i}}_T+\eta ,$$(2)

where υ are the measured visibilities, i_T denotes the discretized sky image, F is the Fourier transform operator, G^† accounts for sampling on an irregular Fourier grid, and η ~ N(0, σ_V) represents instrumental noise with variance σ_V. Applying the adjoint operator yields the dirty image: $${\mathbf{i}}_D=F^{†}G\mathbf{v}.$$(3)

Assuming GG^† is approximately diagonal, i_D is related to the true sky image through a convolution, $${\mathbf{i}}_D=H{\mathbf{i}}_T+\mathbf{n},$$(4)

where H ≡ F^† Diag(GG^†)F is the convolution operator and n is the correlated image-space noise. The related approximation error is typically remedied through an iterative procedure referred to as major loops.

This measurement process differs fundamentally from optical imaging, where observations correspond to a direct convolution with a compact and spatially smooth PSF. In radio interferometry, incomplete sampling of the Fourier plane leads to PSFs with extended sidelobes, while the non-coplanarity of the antenna array introduces additional errors through the w term. These effects, combined with correlated noise, make radio image reconstruction a highly non-trivial inverse problem.

2.2 Source separation

Source blending further complicates radio weak-lensing analyses. One form of blending arises in the image plane, analogous to the optical domain, where the emission from multiple galaxies overlaps spatially. This effect becomes increasingly important in deep surveys, where high source densities substantially increase the probability of overlapping emission (Melchior et al. 2021). Image plane blending has been extensively studied in optical weak-lensing, motivating the development of both modelbased deblending techniques such as SEXTRACTOR (Bertin & Arnouts, S. 1996) and SCARLET (Melchior et al. 2018), as well as data-driven approaches based on deep learning (Hemmati et al. 2022; Burke et al. 2019; Wang et al. 2022). While historically less severe in radio surveys, this form of blending is expected to become increasingly relevant for SKA-Mid, where the source density for galaxies with S/N > 10 may reach approximately 2.7 gal arcmin⁻² (Brown et al. 2015).

Radio interferometry also introduces an additional and fundamentally different source of blending. The measured visibilities encode contributions from all sources within the primary beam, meaning that information about individual objects is intrinsically mixed in the Fourier domain. No single visibility corresponds to a unique sky location. As a result, bright sources can introduce structured, noise-like artefacts through extended sidelobes, even when well separated from the target in the image plane. This form of visibility domain blending poses a major challenge for accurate reconstruction and shape measurement of faint galaxies.

Existing radio weak-lensing pipelines address source separation using either image plane or visibility plane strategies. Image plane approaches operate on reconstructed interferometric images and typically rely on source finding and decomposition algorithms such as PYBDSF (Mohan & Rafferty 2015) to identify and isolate individual galaxies prior to shape measurement. These methods are computationally efficient and easy to integrate into existing imaging pipelines, but their performance is inherently limited by the fidelity of the reconstructed images and by residual sidelobe contamination from bright sources. Visibility plane methods instead operate directly on the measured interferometric data, either by jointly modelling all sources in the field simultaneously (Chang et al. 2004; Rivi et al. 2019) or by applying faceting techniques to isolate the visibilities associated with individual sources or small sky regions (Rivi & Miller 2018). While these approaches are more faithful to the underlying measurement process and can better account for source coupling, they become computationally demanding for wide field observations with large numbers of sources. In contrast, deep learning based source separation methods for radio data remain relatively limited, with most existing approaches focusing on sources with simple morphologies (see Sortino et al. 2023, and references therein).

3 Simulations

Our simulation pipeline builds upon the framework introduced in Paper I, with key modifications to incorporate wide field observations. The main simulation and observational parameters are summarised in Table 1.

3.1 WideSet: Training and validation

The WideSet dataset was constructed from the wide area T-RECS catalogue, which spans approximately 400 deg², and is used for both training and validation. Since the exact pointing strategy of future SKA Mid surveys has not yet been defined, we placed the simulated field centre at (α₀, δ₀) = (56°, −30°). The celestial coordinates (α, δ) of individual sources were obtained by converting the flat sky coordinates provided in the T-RECS catalogue into equatorial coordinates relative to this central position.

From the parent catalogue, we generated 150 distinct sky patches of area 1 deg², each approximately matching the field of view of a single SKA Mid primary beam at 1.4GHz. These patches were constructed to be disjoint, ensuring that no sources are shared between training and validation subsets and thereby preventing data leakage (Sasse et al. 2025). While real surveys often employ overlapping pointings to improve signal-tonoise, as in Harrison et al. (2020), this simplification allows for a clean separation between subsets. Given our chosen pixel scale, each patch corresponds to an image of size 25 200 × 25 200 pixels, rounded to optimise fast Fourier transform (FFT) performance.

Within each patch, galaxies are simulated using Sérsic surface brightness profiles (Sérsic 1963). We restrict the simulated population to a flux density range of 1 μJy ≤ S₀ < 200 μJy. Sources fainter than this limit are unlikely to be detected above the receiver thermal noise, while brighter galaxies are rare in the T-RECS catalogue, with a surface density of approximately 0.06 gal arcmin⁻², compared to ~32.31 gal arcmin⁻² within our selected flux range.

Although such bright sources are not included in the present simulations, they can affect real interferometric observations. Even when located outside the main lobe of the primary beam, they may contribute residual power to the measured visibilities through secondary sidelobes (Mort et al. 2017). If present within the primary beam, they can also introduce ringing artefacts in the dirty image through PSF sidelobes, potentially obscuring nearby faint galaxies. These effects are generally more pronounced at lower observing frequencies, but residual contamination can be mitigated through preprocessing techniques such as source peeling (Noordam 2004) or direction-dependent corrections applied during image gridding (Repetti et al. 2017). We envisage a similar preprocessing step within our framework to remove contributions from such sources when present.

Galaxy images are generated using the GALSIM¹ package (Rowe et al. 2015), and the simulation of individual galaxies is parallelised across multiple cores using DASK² (Dask Development Team 2016) to reduce computation time. A representative random patch from the WideSet dataset is shown in Figure 1.

Fig. 1 Example sky patch from WideSet. For visualisation, only the central 0.1 deg2 region is shown. The colourmap is scaled as log10(S + 10−9), where S denotes the flux density of each pixel in units of Jy per pixel, to enhance contrast and improve visual clarity.

3.2 DeepSet: Testing

The testing dataset, DeepSet, was derived from the 1 deg² deep T-RECS catalogue. We again assumed that the catalogue is centered at (α₀,δ₀) = (56°, −30°), consistent with the WideSet configuration. This patch already matches the desired field size and was therefore treated as a single pointing without further subdivision. The same selection criteria applied to WideSet were enforced here.

In contrast to the wide catalogue used for WideSet, which assigns galaxies random positions, the deep catalogue preserves realistic large-scale structure by linking galaxy positions to dark matter haloes from a cosmological simulation. Halo positions in the deep catalogue are drawn from the high-resolution P-Millennium simulation (Baugh et al. 2019), based on the Planck best-fitting cosmology (Planck Collaboration XVI 2014). The haloes are then populated with radio-emitting star-forming galaxies using an abundance matching procedure (Moster et al. 2013), matching the halo mass function to the observed radio luminosity function. This ensures that both the clustering of haloes and the distribution of galaxies within them are physically motivated, producing a survey-like galaxy spatial distribution. Since DeepSet is based on this deep catalogue, it provides a realistic test bed for our pipeline.

3.3 Visibility simulation

Visibilities for both WideSet and DeepSet patches are simulated using the Radio Astronomy Simulation, Calibration, and Imaging Library (RASCIL)³. For each patch, a separate Measurement Set (MS) is generated, containing the full visibility data corresponding to a single SKA-Mid pointing. Non-coplanar baseline effects are accounted for using the improved W-stacking algorithm (Ye et al. 2021). Each patch is treated as an independent observation, i.e. the array observes one field at a time, with no mutual coupling or contamination between neighbouring pointings. Contributions from secondary lobes of the primary beam are neglected, which would otherwise introduce weak signals from bright sources outside the main field of view.

We do not model primary beam attenuation in our visibility simulations. In real observations, sources near the beam edge appear dimmer due to the primary beam response, but this mainly affects source counts rather than shape measurements. Since primary beam effects can be accurately corrected during calibration or imaging (see e.g. Ch. 19 of Taylor et al. 1999), we neglect them here for simplicity.

3.4 Galaxy stamps

In addition to the full-sky patches, we generate individual galaxy cutouts, or stamps, which are used for training and testing our deblending pipeline. These stamps are classified as either isolated or blended.

Isolated stamps are straightforward to produce: GALSIM generates individual images for each galaxy when simulating the full patch. We then crop and resample these stamps to a uniform size of 128 × 128 pixels to ensure consistency for training and testing. GALSIM blended stamps are constructed from the full patch images. For each central galaxy of interest, we extract a 128 × 128 pixel cutout centered on that galaxy. These cutouts can include multiple neighbouring galaxies, which may be partially or fully present within the stamp. The fluxes of neighbouring sources can be higher or lower than that of the central galaxy. Importantly, even if some nearby galaxies are not fully visible within the extracted stamp, their flux can still contribute to the central galaxy’s signal due to the extended light profiles. This approach ensures that the blended stamps capture realistic source configurations, including overlapping flux contributions and a variety of relative flux ratios, providing a robust training set for the deblender. We do not include instrumental noise or PSF effects in the blended stamps, as these are expected to be effectively corrected by our reconstruction framework.

We quantify the degree of source overlap for each target galaxy using the blendedness parameter, B, following Bosch et al. (2018), $$\mathcal{B}=1-\frac{{\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\cdot {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}}{{\mathbf{i}}_{\mathrm{b}\mathrm{l}}\cdot {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}},$$(5)

where i_iso is the isolated stamp of the target galaxy, i_bl is the corresponding blended stamp containing neighbouring sources, and the dot, ‘•’, denotes the dot product. The parameter B therefore measures the fractional contamination of the target galaxy’s flux by overlapping sources, with B ~ 0 indicating an almost isolated galaxy and B ~ 1 corresponding to strong contamination from neighbouring sources.

4 Methodology

Our methodology consists of three main modules. The first module performs faceting, where the measurement set is divided into smaller patches centered on individual galaxies. Each faceted dirty image, along with its corresponding PSF, is then reconstructed using the HQS-PnP algorithm introduced in Paper I. HQS-PnP is a PnP framework based on the half-quadratic splitting (HQS) method that reconstructs high-fidelity galaxy images from dirty image-PSF pairs.

The second module is a deblending network inspired by a VAE, designed to remove contaminating sources within each reconstructed facet. The final module performs shape estimation using an E(2)-equivariant CNN for feature extraction, combined with a PSF encoder that mitigates any residual PSF artefacts.

4.1 Deblender network

Deep learning-based deblending has been extensively studied in the optical domain, where crowded scenes pose significant challenges for various astrophysical studies. A variety of architectures have been proposed, including generative adversarial networks (Reiman & Göhre 2019; Hemmati et al. 2022), residual dense neural networks (Wang et al. 2022), and instance segmentation approaches such as Mask R-CNN (Burke et al. 2019). Among probabilistic methods, VAE-inspired deblenders have demonstrated strong performance, achieving fast and accurate deblending and, in some cases, outperforming classical approaches such as SCARLET (Melchior et al. 2018) in multiband optical settings (Arcelin et al. 2020; Biswas et al. 2025).

Motivated by these successes, we adopt a VAE-inspired deblender for radio galaxy stamps. Existing optical networks cannot be applied directly to our data, as they were trained on multiband stamps of smaller size and included PSF and noise corruptions. In our framework, a separate deconvolution module first addresses these effects, allowing the deblender to focus solely on removing overlapping emissions. We further enhance the network with multi-head attention layers to capture spatial correlations more effectively (see Sect. 4.1.2) and employ a simplified, single-stage training strategy instead of the two-step training procedure used in prior works (see Sect. 4.1.5). These modifications enable an efficient adaptation of the VAE framework while meeting the specific requirements of our pipeline.

4.1.1 Theory

Our deblender network is inspired by the standard VAE framework and is trained to map blended galaxy stamps to their corresponding isolated images. Unlike a conventional VAE (Kingma & Welling 2019), in which both the encoder and decoder are conditioned on the same observation, our network operates on paired inputs and targets: the encoder receives a blended stamp i_bl, while the decoder reconstructs the corresponding isolated galaxy i_iso. As a result, this network is not a VAE in the strict sense, but its training objective retains the functional form of an evidence lower bound (ELBO).

We let z denote a latent variable drawn from a standard multivariate Gaussian prior, z ~ N(0,I). The generative process is defined by the conditional likelihood p_θ(i_iso|z), parameterized by a decoder network with weights θ that maps latent variables z to reconstructed isolated images. Because the true posterior p_θ(z|i_bl) is intractable, we introduce an encoder q_φ(z|i_bl) with parameters φ to approximate it using a tractable Gaussian distribution, learning both the mean and variance.

The network is trained to find parameters θ̂ and φ̂ that maximise the marginal likelihood of the isolated stamp, p_θ(i_iso). Direct evaluation of this likelihood is intractable, so we instead optimise an ELBO-like lower bound: $$\begin{array}{rl}\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})& \ge {\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})}[\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\mid \mathbf{z})]\\ & -{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\parallel p(\mathbf{z})].\end{array}$$(6)

The first term in Eq. (6) encourages accurate reconstruction of the isolated galaxy, while the second term, the Kullback-Leibler (KL) divergence, regularizes the latent space. Following the β-VAE framework (Higgins et al. 2016), we introduce a weighting factor β on the KL term to balance reconstruction fidelity and latent regularisation. In practice, we find β = 10⁻⁴ provides an effective trade-off, preserving fine morphological details in blended galaxies. A smaller β is preferred in our setting because the loss is defined as a pixelwise sum rather than a mean (see Eq. (14)), which empirically improves optimisation stability.

For a full derivation of this generalised ELBO-like lower bound and a comparison with a standard VAE, we refer to Appendix A. In particular, because the encoder is conditioned on the blended input i_bl while the decoder is conditioned on the isolated target i_iso, the resulting bound is formally looser than a standard ELBO. Nevertheless, the reconstruction and KL terms provide sufficient guidance for the network to learn robust latent representations of the central galaxy, effectively deblending the image for downstream shape measurement tasks.

4.1.2 Network architecture

The encoder and decoder networks use a standard CNN architecture augmented with attention layers. Attention mechanisms dynamically weigh spatial features, allowing the network to focus on the most relevant regions of the image.

Formally, we let the input feature map be X ∈ R^C×H×W. We first flatten the spatial dimensions to obtain $${\mathbf{X}}_{\mathrm{f}\mathrm{l}\mathrm{a}\mathrm{t}}\in {\mathbb{R}}^{C\times N},N=HW.$$(7)

We then compute queries, keys, and values via learnable linear projections at each spatial location, implemented as 1 × 1 convolutions, $$\mathbf{Q}=W_Q{\mathbf{X}}_{\mathrm{f}\mathrm{l}\mathrm{a}\mathrm{t}},\mathbf{K}=W_K{\mathbf{X}}_{\mathrm{f}\mathrm{l}\mathrm{a}\mathrm{t}},\mathbf{V}=W_V{\mathbf{X}}_{\mathrm{f}\mathrm{l}\mathrm{a}\mathrm{t}},$$(8)

where W_Q, W_K, W_V ∈ R^C×C are the learned linear transformations, and Q, K, V ∈ R^C×N.

The attention weights are then obtained through $$\mathbf{A}=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(\frac{{\mathbf{Q}}^{\mathrm{\top }}\mathbf{K}}{\sqrt{C}}\right)\in {\mathbb{R}}^{N\times N},$$(9)

which encodes the attention between each pair of spatial locations. Finally, the output feature map is computed as $$\mathbf{Y}=\mathbf{V}\mathbf{A}\in {\mathbb{R}}^{C\times N}.$$(10)

MHA extends this by splitting the feature channels into h subspaces (heads), each learning a different attention pattern. The outputs from all heads are concatenated along the channel dimension and projected back to the original channel dimension, $$\mathrm{M}\mathrm{H}\mathrm{A}(\mathbf{X})=W_O\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}[{\mathbf{Y}}_1,{\mathbf{Y}}_2,\dots ,{\mathbf{Y}}_h],$$(11)

where W_O ∈ R^C×C is the output projection layer. The resulting tensor is reshaped to R^C×H×W to recover the original feature map dimensions. A residual connection is then applied after scaling by a learnable residual factor γ_r : $$\mathbf{Z}=\mathbf{X}+{\gamma }_r\mathrm{M}\mathrm{H}\mathrm{A}(\mathbf{X}).$$(12)

Combining CNNs with attention allows the model to capture both local and global dependencies efficiently, and has shown improved performance in various image processing tasks such as deblurring and segmentation (Bello et al. 2019; Guo et al. 2022; Kuang et al. 2025). In this work, we adopt a hybrid architecture that integrates convolutional feature extraction with localised self-attention near the latent space, enabling the model to retain the inductive biases of CNNs while benefiting from the increased representational capacity of attention mechanisms. This design provides a balance between performance and computational efficiency. The detailed network configuration, including the placement of MHA layers and the complete model diagram, is provided in Appendix B.

Transformers, which rely entirely on attention mechanisms, have also seen rapid development in recent years. In particular, Vision Transformer (ViT) architectures and their variants, such as DeiT and Swin Transformer (Dosovitskiy et al. 2021 ; Touvron et al. 2021; Liu et al. 2021), have demonstrated state-of-the-art performance across a wide range of computer vision tasks. Several of these approaches, along with related attention-based models, have also been successfully applied to astronomical imaging problems, including source classification, morphology estimation, and cosmological parameter inference (Bhavanam et al. 2024; Kakadia et al. 2025; Parlange et al. 2026).

In the context of deblending, Zhang et al. (2024) adapted the Cross-Aggregation Transformer (CAT) module originally proposed by Chen et al. (2022), introducing the CAT-deblender for optical galaxy images. To provide a direct comparison, we retrain this architecture on our radio galaxy dataset. As shown in Sect. 5.1, this fully transformer-based model achieves lower reconstruction fidelity than our proposed network, while also exhibiting slower convergence. This may reflect both the data-hungry nature of transformer models and the fact that CAT represents an earlier or more specialized transformer design rather than the latest generation of architectures. More recent transformer-based models have not yet been extensively explored for galaxy deblending, and evaluating them in this context remains an important direction for future work.

4.1.3 Image normalisation

Image normalisation ensures stable and efficient training by keeping pixel values within a consistent range. Standard approaches include min-max normalisation (rescaling to [−1,1]) and standardisation (zero mean, unit variance). However, for astronomical deblending, these conventional methods are ill-suited: they overemphasize bright galaxies while suppressing faint ones, making deblending harder, and using raw intensities leads to unstable training. A good scaling operator should therefore map the flux values into a standard range while preserving the relative contrast between sources of varying brightness.

Previous works address this using nonlinear transformations such as the inverse hyperbolic sine (Hemmati et al. 2022) or tanh(sinh⁻¹(∙)) (Zhang et al. 2024), as well as min-max scaling (Wang et al. 2022). We benchmarked several of these approaches and found the following scaling operator to yield the best training stability and deblending performance: $$\mathcal{S}(\mathbf{i})\equiv {\sinh }^{-1}\left(\frac{\mathbf{i}}{\mathrm{f}\mathrm{a}\mathrm{c}}\right).$$(13)

Here, i denotes either the input image (the blended stamp) or the target image (the isolated stamp). The parameter fac is a scaling factor selected to map typical flux values at the μJy level approximately into the [0,1] range; empirically, we found that fac = 5 × 10⁻⁷ provides the best performance.

4.1.4 Training

For training, we use 40 patches from WideSet and reserve 5 separate patches for validation. Only galaxies with flux densities above 10 μJy are included in the training/validation set, since fainter sources are typically masked by nearby bright objects and dominated by thermal noise in the dirty image, making their reconstructions unreliable. Since these galaxies would not likely be used for shape measurement, training the deblender on them is unnecessary. Nevertheless, faint sources are retained in the simulations, as their presence can influence the quality of both reconstruction and deblending from neighbouring, brighter galaxies.

After flux cuts, the dataset contains ~600 000 training and ~80000 validation stamps. The large sample size is crucial for ensuring stable convergence and effectively training attentionbased architectures, which generally require substantial data to learn robust feature dependencies (Vaswani et al. 2017).

Since the model operates in the scaled domain, the reconstruction loss compares network outputs to scaled isolated images. Assuming a 16-dimensional multivariate Gaussian latent space with mean vector μ and variance vector σ, the KL term simplifies to a standard analytical form. The final training objective is: $$\begin{array}{r}\hat{\theta },\hat{\varphi }=\underset{\theta ,\varphi }{\arg min}{\mathbb{E}}_{{\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}},{\mathbf{i}}_{\mathrm{b}\mathrm{l}}}[\sum _{\mathrm{p}\mathrm{i}\mathrm{x}}(\mathcal{S}({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}r)-\mathcal{S}({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}){)}^2\\ -\frac{\beta }{2}\sum _{j=1}^{16}(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2)].\end{array}$$(14)

Here, i_iso and i_bl denote the isolated and blended stamps, respectively. S(i_iso) represents the decoder output (see the architecture diagram in Figure B.1), corresponding to the scaled deblended stamp. To recover the deblended stamp in its original flux scale, we apply the inverse scaling transformation to this decoder output.

4.1.5 Alternative VAE-inspired deblending strategies

Previous works have proposed two-step training strategies for VAE-based deblenders. In Arcelin et al. (2020), the network was first trained to encode and decode isolated stamps. After convergence, the decoder was frozen, and the encoder was retrained so that blended stamps mapped to the same latent representations as their isolated counterparts. Similarly, Biswas et al. (2025) pretrained the network on isolated galaxies and then retrained a second encoder to align the latent vectors of blended inputs with those of the isolated stamps by minimising a Euclidean distance in latent space.

We explored both of these two-stage training strategies for our deblender. In practice, they substantially increased training time without improving reconstruction fidelity compared to a simpler, single-stage approach. We therefore adopt single-stage training for our deblender network. The lack of improvement may reflect differences in the data used by previous studies, which relied on multi-channel images which included explicit PSF and noise corruptions. In those settings, separating representation learning from deblending can help the encoder disentangle instrumental effects before aligning latent representations.

In the full MADNESS pipeline (Biswas et al. 2025), the VAE-deblended stamps serve only as an initial step for a maximum a posteriori (MAP) reconstruction. The MAP procedure then combines all deblended stamps into a wide field image and evaluates a global likelihood against the full observation, effectively enforcing consistency across overlapping regions. Implementing a similar strategy for radio interferometric data would be considerably more complex. Each facet would first need to be deconvolved and then deblended to obtain a reconstructed image, and overlapping facets would have to be merged carefully to form a wide field image. This wide field reconstruction would then be compared with the observed visibilities, which involves degridding operations. To maximise the posterior, these steps would need to be repeated multiple times as the algorithm iteratively moves towards the maximum a posteriori solution, requiring repeated gridding and degridding operations to transition between image and visibility domains. While such an approach would likely improve the overall fidelity of the reconstructed galaxies, it introduces substantial computational and implementation challenges and is beyond the scope of the present work.

4.2 Image deconvolution

In Paper I, we introduced HQS-PnP, a deep learning-based PnP algorithm for radio interferometric image reconstruction. We demonstrated that it can produce high-fidelity reconstructions in a comparable time to traditional deconvolution algorithms such as MS-CLEAN (Cornwell 2008), while achieving substantially better image quality. In the present work, we retain the same reconstruction framework but extend it to the wide field regime.

Applying HQS-PnP directly to wide field observations poses significant computational challenges. In this work, we adopt a conservative oversampling factor of 2, resulting in an image of size 25 200 × 25 200 pixels. Performing reconstruction on such large images is inefficient due to the cost of FFT-based grid-ding operations as well as the repeated application of the learned denoiser over a very large image. While in principle this could be addressed using strategies such as distributed GPU computation, this is not necessary in our setting since we are not interested in reconstructing the full field, but only a small number of target galaxies that are sparsely distributed within it. This makes full-field processing both unnecessary and inefficient.

A common strategy in radio imaging is therefore to divide the problem into smaller regions. CPU-based imagers such as WSCLEAN⁴ (Offringa et al. 2014) follow a similar idea, accelerating wide field deconvolution by splitting large images into smaller sub-images that are processed independently and later recombined, a strategy often referred to as parallel or faceted deconvolution. These approaches typically first form a wide field dirty image and then partition it into multiple sub-images while carefully handling boundaries to avoid flux leakage across regions. Although such strategies are useful for full-field reconstruction, they are less suitable in our case, where target galaxies are sparse. Instead, we adopt a simpler and more uniform strategy in which fixed-size image patches (facets) are constructed around each target source.

4.2.1 Faceting

Rather than reconstructing the entire FoV at once, we form localised dirty images centered on individual galaxies, referred to as facets. Faceting approximates a wide field as a collection of small, narrow-field patches, allowing efficient inference while maintaining high-fidelity recovery. This strategy has previously been used to address direction-dependent effects and wide field imaging (Tasse et al. 2018; Thouvenin et al. 2023; Wilber et al. 2023).

For each galaxy, the visibilities are phase-shifted from the original phase-tracking center to the galaxy position by multiplication with a complex phase-term, $$V^{\prime }(u,v,w)=\exp [2\pi i(u\mathrm{\Delta }{l}_g+v\mathrm{\Delta }{m}_g+w\mathrm{\Delta }{n}_g)]\cdot V(u,v,w),$$(15)

where ∆l_g, ∆m_g are the galaxy’s direction cosines relative to the phase center, and $\mathrm{\Delta }{n}_g=\sqrt{1-\mathrm{\Delta }{l}_g^2-\mathrm{\Delta }{m}_g^2}$. This procedure assumes that galaxy positions within the measurement set are known, typically obtained using source-finding tools such as PYBDSF (Mohan & Rafferty 2015).

Phase shifting alone does not fully correct for the noncoplanar baseline (w-term) effect: the sky remains projected onto facet planes tangent to the original phase-tracking center rather than each galaxy center, limiting the usable field of view (Taylor et al. 1999, Ch. 19). To mitigate this, Kogan & Greisen (2009) introduced a first-order approximation that incorporates the w-term into shifted (u, v) coordinates: $$u^{\prime }=u-\frac{\mathrm{\Delta }{l}_g}{\mathrm{\Delta }{n}_g},v^{\prime }=v-\frac{\mathrm{\Delta }{m}_g}{\mathrm{\Delta }{n}_g}.$$(16)

With this approximation, each facet’s visibilities can be treated using a simple two-dimensional Fourier transform over (u′, v'). This is valid for sufficiently small facets, where higher-order residuals of the w-term are negligible. Tasse et al. (2018) extend this with a second-order polynomial approximation, but such corrections are only needed for fields larger than 1 deg² and are therefore omitted in this work.

After re-centering, the visibilities are inverted on a coarser uv grid to form faceted dirty images. Increasing the uv cell size ∆u reduces the field of view in the image domain, since FoV ~ 1∕∆u. This allows us to construct compact 128 × 128 dirty images, which are significantly cheaper to process during deconvolution and subsequent downstream steps due to smaller FFT sizes. The trade-off is an increase in aliasing, loss of small-scale structure from smoothing and interpolation, and a broader effective PSF with more pronounced sidelobes. In particular, emission outside the field of view can fold back into the image, contaminating the facet and making reconstruction more challenging.

While a finer uv grid (and hence larger facets) could in principle mitigate these effects, we do not observe improved reconstruction quality in practice. This is likely because the denoiser is trained on 128 × 128 inputs, and its performance degrades when applied to larger images, leading to suboptimal denoising. As a result, the 128 × 128 facet size provides the best overall performance, both in terms of reconstruction quality and computational efficiency. We further investigate the impact of facet size in Appendix C.

A related faceting strategy was proposed by Rivi & Miller (2022), where the facet size scales with source flux and size, assigning larger facets to brighter and more extended galaxies. However, such an approach requires prior estimates of these properties to determine the appropriate facet size for each source. In practice, this introduces an additional preprocessing step at the measurement-set level, where fluxes and sizes must be estimated from partially calibrated or dirty data, potentially introducing uncertainties and increasing pipeline complexity.

In contrast, we adopt a fixed uv resolution that produces uniform 128 × 128 pixel facets for all sources. This choice ensures consistency across the deblending pipeline and enables efficient batching for GPU-based processing. For the range of galaxy sizes considered here, this facet size is sufficient to capture the dominant morphological features relevant for shape estimation. While very faint outer isophotes of particularly extended sources may lie beyond the facet boundary, previous results (Paper I) show that the core shape information remains well preserved. The resulting faceted dirty images and their corresponding PSFs are then passed to the HQS-PnP algorithm to remove the effects of the dirty beam.

4.2.2 Parameter setting

The HQS-PnP algorithm uses DRUNet (Zhang et al. 2020) as its learned image denoiser. We adopt the implementation provided by the DEEPINVERSE⁵ library (Tachella et al. 2023) and initialise the network with the provided pretrained weights. These weights encode a strong general prior but are trained on natural images, and thus do not capture the morphology or noise characteristics of radio galaxies. To obtain a strong prior adapted to radio data, we fine-tune the denoiser using blended galaxy stamps drawn from 40 patches of WideSet.

At each iteration, 32 stamps are sampled, rescaled to the unit range, and corrupted with Gaussian noise whose amplitude is randomly drawn between 0 and 0.6. This produces substantially noisier inputs than those used in the original DRUNet training, where the maximum noise level is 50/255. The denoiser is trained to reconstruct the clean stamps from their noisy counterparts using an ℓ₁ loss. Fine-tuning runs for 20000 iterations with a cosine learning rate schedule that decays from 10⁻⁴ to 10⁻⁹. This yields a much stronger, domain-specific prior that is essential for high-fidelity deconvolution of radio images. The fine-tuned denoiser is kept fixed for all subsequent steps.

With the denoiser parameters fixed, the remaining components of the HQS-PnP algorithm require setting four hyperparameters: the number of iterations K, the regularisation weight λ, and the denoising strengths at the first and last iterations, σ₁ and σ_K. Following Paper I, we set K = 30 and tune the remaining hyperparameters using the OPTUNA library (Akiba et al. 2019). The search is performed over 250 trials and is configured to maximise the peak signal-to-noise ratio (PSNR) of the reconstructed central galaxy. Because our goal is accurate recovery of the target galaxy rather than the full blended scene, the tuning metric is evaluated after both deconvolution and deblending. Specifically, the HQS-PnP output is passed through the trained deblender, and the resulting isolated stamp is compared with the ground truth. Optimising directly against the blended stamp would implicitly encourage the reconstruction to model contaminating neighbours. As many of these sources are faint, this would bias the solution towards weaker denoising and reduce fidelity for the target galaxy. Evaluating the metric after deblending therefore provides a more targeted criterion for hyperparameter selection.

Evaluation is performed on a single validation patch from WideSet that is unseen by both the denoiser and the deblender, ensuring a fully independent validation set. Using this procedure, we obtain the optimal hyperparameters σ₁ = 3.70σ, σ_K = 0.38σ, and λ = 3.11. An analysis using the fANOVA importance calculator (Hutter et al. 2014) implemented in OPTUNA shows that σ_K is by far the most influential hyperparameter, with an importance of 95%. This is expected in a PnP framework, as the final denoising strength directly controls the balance between data fidelity and prior enforcement at convergence, and therefore has the strongest impact on the quality of the final reconstructed image.

4.3 Shape measurement network

We adopt the same shape measurement architecture introduced in Paper I. The PSF autoencoder retains its original weights, since the PSFs in our setting are deterministic and consistent with those used previously, making their latent representation directly transferable. In contrast, the feature extraction block is retrained from scratch. This is necessary because the earlier feature extractor was trained on images that did not reflect the biases introduced by blending or by the updated reconstruction pipeline. Retraining allows it to adapt to the new imaging conditions and accounts for both sources of contamination. For this stage we use 40 patches from WideSet. These patches are entirely separate from those used to train the denoiser and the deblender, ensuring that the shape model is not exposed to any data processed by the upstream components. This separation prevents information leakage and avoids biases that could arise if the shape estimator learned to compensate for artefacts introduced earlier in the pipeline.

We follow the same two-stage training strategy as in the previous work. The model is first pretrained on 300 000 noiseless galaxy images and subsequently refined using 110000 reconstructed image and PSF pairs. During the refinement stage, it is crucial to restrict the training set to stamps with reliable reconstructions, as poorly reconstructed examples can slow convergence and bias the learned shape estimates.

To enforce this, we apply a quality cut requiring that the peak of the dirty image corresponding to the isolated source be at least three times the image-plane noise level σ, $$\mathrm{m}\mathrm{a}\mathrm{x}({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\star \mathbf{h})\ge 3\sigma ,$$(17)

where h is the corresponding PSF. Additionally, we exclude galaxies with flux below 50 μJy. This filtering removes sources likely to produce unreliable reconstructions, as including all stamps would otherwise degrade training stability and reduce the robustness of the resulting shape estimates. We adopt the same selection criterion as in Paper I, thus allowing us to compare results with the previous work and evaluate how the shape measurement bias changes when moving from isolated to wide field measurements.

5 Results

5.1 Deblending

We evaluate our deblender network on the true galaxy images from DeepSet, ensuring that reconstruction errors arise solely from the deblending process and not from PSF effects. In practice, PSF convolution would further degrade the reconstruction quality, and we examine this in Sect. 5.2.

For benchmarking, we compare our network against two alternatives. The first is a standard VAE-based network composed solely of convolutional layers, without multi-head attention (dubbed as VAE-CNN). Such purely convolutional architectures are commonly used for optical galaxy deblending (Arcelin et al. 2020; Biswas et al. 2025); the only difference from our deblender (dubbed as VAE-MHA) is the inclusion of the MHA layer. The second comparison model is the fully transformerbased CAT-deblender proposed by Zhang et al. (2024) (CAT). While powerful, fully attention-based networks are considerably slower to train than CNN-based models. To ensure a fair comparison, all three models are trained for an equal duration of 300 hours on a single V100 GPU, with the model achieving the lowest validation loss selected for inference. The same training and validation datasets are used for all three networks.

Since weak-lensing studies require high-fidelity galaxy shapes, we quantify reconstruction quality using metrics sensitive to structural differences rather than global flux errors:

1. PSNR: PSNR measures the peak error in reconstruction quality and is defined as the logarithm of the ratio between the highest pixel squared value of the ground-truth image (the isolated stamp, i_iso) and the mean squared error between the ground-truth and the reconstructed image (the deblended stamp, î_iso), expressed in decibels (dB): $$\mathrm{P}\mathrm{S}\mathrm{N}\mathrm{R}=10{\log }_{10}\frac{max({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}{)}^2}{⟨({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}r-{\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}{)}^2⟩}\mathrm{d}\mathrm{B}.$$(18)

2. Ellipticity error: the ellipticity error is the Euclidean distance between the ellipticity measured from the reconstructed stamp ε̂ and the true ellipticity ε^true: $$\mathrm{\Delta }ϵ=\Vert \widehat{ϵ}-{ϵ}^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}{\Vert }_2.$$(19)

   For the analysis presented in this section, galaxy shapes are measured using the HSM module of GALS IM, which estimates shapes using adaptive image moments. Using this standard approach allows us to assess shape degradation with a method independent of our dedicated shape measurement network, thereby avoiding any bias in the evaluation of the reconstruction performance.

All metrics are computed on non-scaled stamps, even though the networks are trained on and predict scaled images. This ensures that the metrics reflect the true reconstruction fidelity relevant for weak-lensing.

Figure 2 shows example reconstructions for four representative galaxies from DeepSet, arranged in order of increasing blendedness. Each stamp includes the corresponding PSNR and ∆ε values. In Figure 3, we present the distributions of these metrics across the test sample, including the median and inter-quartile ranges (IQRs) for each network.

We also investigate how the reconstruction quality varies with several intrinsic galaxy properties, including blendedness, flux, shape, and size. Among these, the reconstruction fidelity shows the strongest dependence on blendedness, as illustrated in Figure 4.

Our results indicate that the VAE-MHA model achieves the best overall performance across all evaluated metrics, followed closely by VAE-CNN, while CAT performs significantly worse. Although VAE-MHA and VAE-CNN yield broadly comparable results, the MHA variant consistently outperforms the CNN version across all blendedness bins. This demonstrates the benefit of incorporating attention layers, which allow the network to capture long-range spatial correlations and contextual dependencies - capabilities that are especially valuable in crowded fields.

The poorer performance of CAT can be attributed to its deterministic, non-generative architecture. As a UNet trained solely with a reconstruction loss, CAT lacks a latent representation that can model the underlying data distribution or probabilistically disentangle overlapping sources. Consequently, while it performs adequately in low-blendedness regimes where the target and companion are well separated, it struggles as emissions overlap. In such cases, CAT often produces non-physical morphologies that combine both galaxies rather than isolating the intended target. We additionally experimented with longer training schedules for CAT; however, this did not yield any meaningful improvement in performance, indicating that the limitations are intrinsic to the architecture rather than a consequence of insufficient optimisation. This behaviour highlights the importance of generative modelling for deblending, where capturing uncertainty and disentangling sources are crucial for reliable reconstruction.

In terms of computational efficiency, the two VAE-based architectures also compare favourably, with inference times of approximately 1.6 ms per galaxy for VAE-CNN and 1.7 ms for VAE-MHA, whereas CAT is considerably slower at around 35 ms per galaxy.

Even with these differences, all three networks face challenges in the most extreme scenarios - such as when a faint target lies near a bright companion or is heavily blended. In these cases, reconstructions are often poor, with the networks recovering only the dominant, brighter galaxy and missing the fainter target entirely. While, in principle, such cases could be identified if the intrinsic flux and blendedness of all galaxies were known, neither quantity is reliably measurable from real observations. Flux estimates are contaminated by nearby emission, and true blendedness requires access to the unblended scene. These difficulties are further compounded by biases in source detection and centroiding; when a bright galaxy lies nearby, source-finding algorithms tend to assign centroids closer to the brighter object, biasing the deblender towards reconstructing it instead of the intended target (Dawson et al. 2016). Although such extreme configurations are relatively rare, they represent a potential source of systematic bias and warrant further study.

Fig. 2 Examples of four galaxies from the DeepSet, arranged by increasing blendedness B. Top row : blended inputs (titles show B). Second row : isolated stamps (titles show flux of the central galaxy). Bottom rows : reconstructions from the three networks, with titles indicating PSNR and ellipticity error ∆ε for each stamp.

Fig. 3 Comparison of reconstruction metrics for the three models. Top : PSNR. Bottom : ellipticity error ∆ε (scaled by 100 for clarity). Each curve shows the KDE, the dashed line indicates the median, and the legend reports the median with the IQR. The shapes are measured using GALSIM.

Fig. 4 Deblending performance as a function of blendedness for the three models. The top panels show binned box plots for the three models in terms of PSNR and ellipticity error, while the bottom panel shows the one-dimensional distribution of blendedness. The shapes are measured using GALSIM.

Fig. 5 Examples of four galaxies from the DeepSet. Top row : isolated stamps (titles indicate flux of the central galaxy). Second row : blended inputs (titles show B). Third row : dirty facet images. Fourth row : deconvolved stamp output by HQS-PnP. Fifth row : Residual facet images. Bottom row : final reconstructed stamps, with titles reporting PSNR and ∆ε for each stamp. The shapes are measured using GALSIM.

5.2 Image reconstruction

In this section, we present the results of the full reconstruction pipeline, which recovers isolated galaxy stamps from dirty facets and their associated PSFs. The procedure begins by dividing the wide field visibility data into smaller facets, as described in Sect. 4.2. Each dirty facet is reconstructed using the HQS-PnP algorithm, yielding an intermediate deconvolved image, which is then processed by the trained VAE-MHA model to perform source deblending and produce the final reconstructed images î.

Figure 5 shows four representative examples from DeepSet, selected at random, illustrating the different stages of the reconstruction. For each galaxy, we display the isolated stamp, the blended stamp, the corresponding faceted dirty image, the deconvolved facet image i̇, the residual dirty image, and the final reconstructed stamp î. The residual image i_R is defined as $${\mathbf{i}}_R=F^{†}G(\mathbf{v}-G^{†}F\tilde{ı}),$$(20)

where F denotes the Fourier transform operator, υ the measured visibilities, and G the gridding operator for 128 × 128 facets.

As shown in Figure 5, HQS-PnP reliably reconstructs the bright central regions of galaxies, while the reconstruction quality degrades towards the outer, low surface brightness regions where the signal is dominated by instrumental noise. As a consequence, the deconvolved images tend to appear rounder for fainter sources, whereas bright galaxies are recovered more accurately. The deconvolution also suppresses most contaminating sources with lower brightness, indicating that once a sufficiently clean deconvolution is obtained, even a relatively weak deblender is adequate to isolate the target galaxy.

To isolate the impact of source contamination and wide field effects on the reconstruction, we performed an additional analysis using isolated stamps. In this test, galaxies from DeepSet were simulated as truly isolated sources at their original locations, with visibilities generated under the assumption of no contamination from neighbouring galaxies or wide field effects. These isolated galaxies were reconstructed using the same HQS-PnP procedure, without applying deblending, since the sources are isolated by construction. We find that the resulting reconstructions are consistently of higher fidelity than their wide field counterparts, both in terms of PSNR and ellipticity error.

Figure 6 shows the corresponding dirty images and reconstructions for the same galaxies presented in Figure 5, while Table 2 provides a quantitative summary of the reconstruction fidelity for both wide field and isolated cases. The reported statistics correspond to the median and IQR of the PSNR and ellipticity error ∆ε across the sample. Just as in the previous section, we use the HSM module from GALSIM for shape measurement. As expected, the isolated reconstructions achieve systematically higher PSNR and lower ellipticity errors, reflecting the absence of contamination and the improved recovery of galaxy structure.

We further examine how the reconstruction fidelity depends on intrinsic galaxy properties, including peak flux, size, Sér-sic index, distance from the phase centre, and blendedness. The strongest dependencies are observed with peak flux and galaxy size. Galaxies with higher peak flux are reconstructed with higher PSNR and lower ellipticity error, indicating more faithful recovery of both intensity and shape information.

In contrast, larger galaxies tend to exhibit lower PSNR, indicating that HQS-PnP reliably recovers the bright central regions while the reconstruction quality degrades in the outer, low surface brightness regions dominated by instrumental noise. This dependence of reconstruction fidelity on galaxy size is consistent with that reported in Paper I. Despite this trend in PSNR, we do not observe a strong correlation between galaxy size and ellipticity error. This indicates that, at least for moment-based shape measurements, the central isophotes contain sufficient information to yield robust ellipticity estimates even when the outer structure is only partially recovered.

These trends are summarised in Figure 7. We observe similar behaviour for both wide field and isolated reconstructions, with the isolated case achieving systematically better performance across all tested bins.

It is also worth noting that the PSNR and ellipticity errors of the final reconstructed stamps remain significantly worse than those obtained for the purely deblended images shown in Figure 2, even in the isolated case. This indicates that the dominant limitation of the full reconstruction pipeline arises from the deconvolution step, rather than from the subsequent deblending. Motivated by this observation, we explored whether the deconvolution stage itself could be further improved by recovering additional structure from the residual visibilities.

One potential strategy to improve the deconvolution is to further process the residual images in order to recover additional structure, following an approach analogous to a major-minor cycle framework. In practice, however, the residuals produced by HQS-PnP are largely noise dominated, with the remaining signal primarily associated with contaminating sources. Additional deconvolution of these residuals does not lead to improved reconstructions and can even degrade the results by reintroducing flux from neighbouring galaxies.

An alternative approach is to subtract the model visibilities of the reconstructed facets from the observed visibilities, $${\mathbf{v}}_R=\mathbf{v}-G^{†}F\sum _{f=1}^{N_{\mathrm{f}\mathrm{a}\mathrm{c}\mathrm{e}\mathrm{t}}}{\tilde{ı}}_f,$$(21)

where i̇_f denotes the deconvolved image for facet f, and N_facet are the total number of facets. The residual visibilities υ_R can then be re-faceted and deconvolved in an attempt to recover additional flux. While this procedure can yield modest improvements, it approximately doubles the computational cost and requires careful tuning of the denoising strength to avoid reconstructing contaminating sources. Therefore, we adopt a single deconvolution of each dirty facet followed by deblending, which we use to produce the final reconstructed stamps.

While the metrics discussed above, in particular the ellipticity error ∆ε, provide a useful and unbiased indication of how reconstruction errors propagate into galaxy shape estimates, they are derived from a standard moment-based shape measurement that is not optimised for the reconstructed stamps. To perform a more precise and application-specific assessment, the final reconstructed images are therefore passed to our dedicated shape measurement network, which we use in the following section to evaluate shape recovery in detail.

Fig. 6 Reconstruction of isolated galaxy stamps. Top row : dirty images generated from the same galaxies as in Figure 5. Bottom row : corresponding reconstructions using HQS-PnP.

5.3 Shape measurement

In this section, we evaluate galaxy shape recovery by the shape measurement network using the final reconstructed stamps produced by the reconstruction module. We apply the DeepShape shape measurement network to both the wide field reconstructions and the isolated reconstructions described in Sect. 5.2, allowing us to quantify the impact of wide field effects and residual source contamination on shape estimation.

For consistency, we use the same trained network weights for both cases. We verified that using the network weights exclusively on isolated reconstructions does not lead to a significant change in performance, and therefore adopt a single model for all experiments. We apply the same S/N selection as used during training, retaining only galaxies with total flux above 50 μJy and a dirty-image peak exceeding three times the noise level, as defined in Eq. (17).

Figure 8 shows the shape measurement residuals for the isolated and wide field cases as a function of the true input ellipticity. The contours, estimated using kernel density estimation, enclose 50 and 90% of the data respectively. For the isolated case, both contours are compact and well-centered around zero residual, indicating consistent shape recovery across the ellipticity range. For the wide field case, while the inner contour remains compact, the outer contour is notably more elongated, reflecting a broader scatter in the residuals. This is consistent with the overall worse reconstruction quality in the wide field configuration, as evidenced by the lower PSNR and higher shape measurement errors reported in Table 2.

To quantify shape recovery, we compute a set of shape recovery metrics that capture both bias and dispersion in the measurements. in particular, we estimate the multiplicative bias m̂_α and additive bias ĉ_α for each ellipticity component by fitting a linear relation between the true ellipticity ${ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}$ and the measurement residual ${\widehat{ϵ}}_{\alpha }-{ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}$, with α ∈ (1,2) corresponding to the two ellipticity components.

In addition to bias, we compute the ellipticity error Ae, defined in Eq. (19), where the predicted ellipticities are obtained from the shape measurement network rather than from GALSIM. This metric provides a direct measure of the typical deviation between predicted and true shapes. To characterise the overall dispersion of the measurements, we also compute the Pearson correlation coefficient ρ_α for each ellipticity component, $${\rho }_{\alpha }=\frac{\mathrm{c}\mathrm{o}\mathrm{v}({\widehat{ϵ}}_{\alpha },{ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}})}{\sigma ({\widehat{ϵ}}_{\alpha })\sigma ({ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}})},$$(22)

where $\mathrm{c}\mathrm{o}\mathrm{v}({\widehat{ϵ}}_{\alpha },{ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}})$ denotes the covariance between the predicted and true ellipticities, and $\sigma ({\widehat{ϵ}}_{\alpha })$ and $\sigma ({ϵ}_{\alpha }^{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}})$ are their respective standard deviations. Values of ρ_α close to unity indicate a tight correlation and low scatter in the recovered shapes.

The resulting shape recovery metrics are summarised in Table 3. For brevity, we report the Pearson correlation coefficient and bias parameters for the first ellipticity component only, noting that the second component exhibits similar behaviour.

As expected, the isolated reconstructions yield the most accurate shape estimates, with low ellipticity error, negligible bias, and stronger correlation with the ground truth. In contrast, the wide field case exhibits increased scatter and a larger multiplicative bias, reflecting the residual impact of source blending and reconstruction errors. Overall, these results demonstrate that while the shape measurement network remains robust under wide field conditions, effective source separation and high-fidelity reconstruction are critical to controlling bias and dispersion in the final shape estimates.

We note that the measured linear bias exceeds the SKA-Mid weak-lensing requirement (Brown et al. 2015). Nevertheless, this level of bias is comparable to that reported in other wide field studies; for example, RadioLensfit (Rivi & Miller 2018) found a bias of roughly 0.07 for similar galaxy populations and observed a marked increase when transitioning from isolated to wide field measurements.

To explore this further, we attempted to apply RadioLensfit directly to our simulated datasets, but adapting it proved challenging, likely due to differences in the simulation setup. In contrast to previous RadioLensfit studies, where the simulated fields contained only sources within the target flux range, our simulations include a substantial population of galaxies below 10 μJy. We also include cases in which galaxy emission overlaps in the image plane, which may lead to errors in model-based fitting. We tested smaller fields and lower-density simulations, but were unable to run RadioLensfit for our wide field measurement sets. Moreover, our simulations generate visibilities from wide field images using wide field degridders, which can introduce artefacts into the simulated visibility measurement sets. Fully accounting for these effects, or adapting RadioLensfit to handle them, would be required for a consistent comparison.

The increased bias in wide field measurements is likely driven by flux mixing in the visibility domain. While faceting partially isolates individual sources, contamination from other objects within the primary beam can leak into a given facet through PSF sidelobes. To quantify this effect, we define a diagnostic factor f as $$f=\frac{\mathrm{m}\mathrm{a}\mathrm{x}({\mathbf{i}}_{D,f})-\sigma }{\mathrm{m}\mathrm{a}\mathrm{x}({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\star \mathbf{h})},$$(23)

where i_D,f is the dirty facet image. This factor provides a heuristic measure of whether the brightest peak in the facet is dominated by the target galaxy (f ≤ 1) or contaminated by other sources (f > 1). Figure 9 shows the multiplicative bias slope m̂ as a function of f.

As illustrated, bias increases rapidly with increasing f, confirming that strong contamination within a facet significantly degrades shape measurements. When f ≤ 1, the bias approaches the level observed for isolated galaxies, consistent with effective source isolation. In this sense, f can be interpreted as a visibilitydomain analogue of the image-plane blendedness parameter B. Unlike B, which captures overlap within a local stamp, f is sensitive to contamination arising from distant sources through extended PSF sidelobes, an effect that is intrinsic to wide field radio interferometric imaging.

We emphasise that this factor is only a coarse diagnostic. It captures whether a source brighter than the target dominates the facet, but does not account for the cumulative effect of multiple fainter contaminants that can still perturb the reconstruction. Defining a statistic that fully characterises such visibility-domain contamination remains non-trivial. Moreover, the deblending network is trained on image-plane blends and relies on the faceting stage to suppress wide field artefacts. When residual contamination persists within a facet, it propagates through the reconstruction and cannot be fully corrected by the deblender, ultimately degrading the accuracy of the inferred shape measurements.

We also examine the dependence of shape measurement accuracy on intrinsic galaxy properties. Figure 10 shows a strong correlation between total flux and ellipticity error ∆ε: fainter galaxies systematically exhibit larger shape errors in both isolated and wide field cases, while isolated reconstructions consistently achieve lower errors across all flux bins. We observe a weaker dependence on peak flux, which is largely correlated with total flux, and no significant dependence on galaxy size or Sérsic index.

Figures 9 and 10 together indicate that shape measurement accuracy improves substantially when applying quality cuts to the galaxy sample. In a realistic weak-lensing pipeline, preprocessing would identify candidate source positions, excluding galaxies with very low flux or strong contamination from neighbours. Within our framework, removing such poorly reconstructed objects does not influence the accuracy of the remaining high-quality objects, since shape estimation is performed independently for each galaxy. This independence may not hold for global model-fitting approaches, where faint sources can potentially bias the inferred parameters of even bright galaxies.

At the same time, if source-finding algorithms can reliably detect low-flux galaxies, including them in shear analyses may still be beneficial. Even though individual measurements are noisy, the increased sample size can improve the overall statistical precision of shear estimates (Rivi & Miller 2018). A quantitative assessment of this trade-off is beyond the scope of this work and is left for future investigation.

Fig. 7 Dependence of reconstruction fidelity on galaxy properties. The top row shows PSNR and the bottom row shows ∆ε. The first column shows the median metric in bins of peak flux. The second column shows the median metric in bins of galaxy size, with shaded regions indicating the IQR. The third column shows the KDE of each metric; a dashed line marks the median. The shapes are measured using GALSIM.

Fig. 8 Shape measurement residuals for the isolated and wide field cases as a function of the input ellipticity for the first component ε1. The two-dimensional contours are estimated using kernel density estimation. The coloured dashed lines indicate the best-fit linear relations for each case. The panels on the right show the corresponding 1D marginal distributions of the residuals.

Fig. 9 Estimated bias slope as a function of factor f.

Fig. 10 Dependence of the ellipticity error as a function of total flux. The full DeepSet sample is used, without imposing flux or peak cuts, to assess performance across the full dynamic range.

5.4 Time consumption

The dominant computational cost of our framework arises during the training phase. Three neural networks are trained: the VAE-MHA deblender, the DRUNet denoiser used within the HQS-PnP reconstruction, and the shape measurement network. The VAE-MHA deblender is trained for approximately 300 hours on a V100 GPU, although the validation loss typically converges after ~65 hours. The DRUNet is initialised from pre-trained weights and requires only fine-tuning on radio galaxy images, which takes approximately 10 hours on an A100 GPU. The two-stage training of the shape measurement network requires a further ~80 hours on an A100 GPU. As is typical for supervised learning approaches, this training cost is incurred only once and can be amortised over large survey volumes.

In contrast, the inference stage of the framework is computationally efficient. We benchmark inference on a node equipped with an AMD EPYC 7262 CPU (252 GB RAM) and a Quadro RTX 6000 GPU. Faceting and gridding of the visibilities are performed on the CPU and constitute the primary bottleneck, taking approximately 0.68 s per galaxy. The reconstruction stage, which includes the HQS-PnP deconvolution and the subsequent image plane deblending, requires approximately 0.1 s per galaxy for a given dirty facet and PSF. The shape measurement network then requires an additional 0.003 s per galaxy for each reconstructed image and PSF pair.

Overall, the complete DeepShape pipeline, from visibility measurement sets and source positions to final shape estimates, takes roughly 0.85 s per galaxy, including data and model loading as well as transfers between stages. This represents a substantial speed-up compared to existing visibility plane approaches. For example, RadioLensfit reports an inference time of ~10 s per galaxy in its serial implementation (Rivi & Miller 2018) (on an Intel Xeon E5-2650 CPU). Overall, this demonstrates that our framework provides a fast, fully end-to-end solution for shape measurement in realistic wide field radio observations, making it well-suited for the data volumes and source densities expected from next-generation surveys.

6 Conclusions and future work

6.1 Conclusions

In this work, we extend the DeepShape framework introduced in Paper I to wide field radio interferometric observations, with a particular focus on the challenge of source separation in crowded fields. Our primary aim is to demonstrate that accurate radio weak-lensing shape measurement can be achieved within a fully deep learning based, end-to-end pipeline that operates efficiently on wide field visibility data.

The proposed framework consists of three key modules. First, a faceting-based strategy partitions the wide field visibility data into localised sky regions, enabling partial isolation of individual sources, followed by deconvolution using the original HQS-PnP algorithm. Second, residual contamination from neighbouring sources in the image plane is addressed using a learned deblender. We investigate three deblending architectures: a standard convolutional variational autoencoder based network (VAE-CNN), an enhanced version incorporating multi-head attention (VAE-MHA), and a fully transformer-based deblender (CAT). Our results demonstrate that the VAE-based models consistently outperform CAT in both reconstruction fidelity and computational efficiency. Among them, VAE-MHA achieves the best overall performance across all tested metrics, highlighting the importance of attention mechanisms for capturing long-range spatial correlations and contextual information in crowded radio fields.

We then evaluate the full reconstruction pipeline, combining faceted deconvolution and learned deblending, and compare it against the original DeepShape framework applied to isolated galaxies. While the isolated case consistently yields the best performance, the wide field pipeline exhibits similar qualitative trends across reconstruction metrics, indicating that the proposed source separation strategy effectively mitigates, though does not fully eliminate, the impact of source entanglement.

Finally, we retrain the original DeepShape shape measurement network on the reconstructed galaxy stamps and assess the resulting shape estimates. As expected, shape recovery in wide field conditions shows increased scatter and bias compared to the isolated case. In particular, the multiplicative bias degrades from the per cent level for isolated galaxies to approximately ten per cent in wide field measurements. Despite this degradation, the achieved bias is comparable to that reported by state-of-the-art visibility plane approaches, demonstrating that competitive shape measurement accuracy can be retained within an image plane, deep learning driven framework.

Crucially, our method achieves this performance with a per-galaxy inference time of approximately 0.85 seconds, representing a speed-up of more than an order of magnitude relative to traditional visibility domain techniques. This combination of competitive accuracy, explicit source separation, and substantial computational efficiency establishes our framework as a viable and scalable solution for radio weak-lensing analyses in the era of next-generation surveys such as SKA-Mid.

6.2 Future work

Several avenues exist for extending this work. A key next step is to integrate the framework with an automated source-finding stage, such as PYBDSF or a deep learning based alternative. In this study, source positions are assumed to be known, but in realistic observations, they must be inferred from noisy visibilities. Errors in source localisation can propagate through the faceting and reconstruction stages, degrading deblending performance and potentially biasing the recovered shapes. Incorporating source finding will therefore allow a more realistic evaluation of the end-to-end pipeline.

A second important extension is to propagate shape measurement uncertainties into shear estimation. While low-flux galaxies exhibit noisier shape estimates, their inclusion increases the effective source density and may improve statistical precision. Exploring the impact of different quality cuts and weighting schemes, as well as the ability of source-finding algorithms to identify poorly reconstructed galaxies, will be essential.

Ultimately, we aim to combine source finding, reconstruction, deblending, shape measurement, and shear estimation into a single end-to-end pipeline operating directly on radio interferometric visibility data. Such an approach would enable direct shear estimation from visibility measurement sets and provide a scalable framework for next-generation radio weak-lensing surveys.

Data availability

Our implementation, results, and usage instructions are available in the repository https://github.com/priyamvad-tripathi/DeepShape2/. The version corresponding to this paper is archived as a tagged release https://github.com/priyamvad-tripathi/DeepShape2/releases/tag/deepshape2-aap-release

Acknowledgements

This research was funded by the Agence Nationale de la Recherche (ANR) under the project ANR-22-CE31-0014-01 TOSCA. This work was granted access to the HPC resources of IDRIS under the allocation AD010415840 and AD010415840R1 by GENCI.

References

Appendix A ELBO-like formulation for the deblender network

The objective of our VAE-inspired deblender network is to maximise the log-likelihood of the isolated galaxy stamp log p_θ(i_iso). We start from the KL divergence between the encoder posterior q_φ(z|i_bl) and the true posterior of the latent given the corresponding isolated stamp p_θ(z|i_iso): $$\begin{array}{rl}{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\Vert p_{\theta }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})]& ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})}[\log q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\\ & -\log p_{\theta }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})].\end{array}$$(A.1)

Using Bayes’ theorem for the true posterior, $$p_{\theta }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})=\frac{p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\mid \mathbf{z})p(\mathbf{z})}{p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})},$$(A.2)

we can rewrite the KL divergence as $$\begin{array}{rl}{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\Vert p_{\theta }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})]& ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})}[\log q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\\ & -\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\mid \mathbf{z})\\ & -\log p(\mathbf{z})+\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})].\end{array}$$(A.3)

Rearranging gives the log-likelihood of the isolated stamp: $$\begin{array}{rl}\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})& ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})}[\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\mid \mathbf{z})]\\ & -{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\Vert p(\mathbf{z})]\\ & +{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\Vert p_{\theta }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})].\end{array}$$(A.4)

Since the KL divergence is always non-negative, dropping the final term yields a valid lower bound on the log-likelihood: $$\begin{array}{rl}\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}})& \ge {\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})}[\log p_{\theta }({\mathbf{i}}_{\mathrm{i}\mathrm{s}\mathrm{o}}\mid \mathbf{z})]\\ & -{\mathbb{D}}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(\mathbf{z}\mid {\mathbf{i}}_{\mathrm{b}\mathrm{l}})\parallel p(\mathbf{z})].\end{array}$$(A.5)

This objective has the same functional form as the standard VAE ELBO, but the conditioning differs in an important way. In a conventional VAE, both the encoder posterior and the true posterior are conditioned on the same observation i_iso, and the final KL term in Eq. A.4 can vanish when the variational family is sufficiently expressive. In our formulation, however, the encoder is conditioned on the blended input i_bl while the generative model is conditioned on the corresponding isolated target i_iso. As a result, the dropped KL term generally does not approach zero, and the bound is formally looser than the standard ELBO.

Despite this mismatch, the objective remains effective for training the deblender. The reconstruction term directly encourages the latent code inferred from the blended image to generate the correct isolated galaxy, while the KL regularisation maintains a well-behaved latent space. Empirically, this provides sufficient pressure for the encoder to learn a representation of the central source that is robust to contamination. In practice, the network learns to treat neighbouring flux as a nuisance factor and to encode the dominant morphology of the target galaxy, which is the quantity of interest for downstream shape measurement.

Appendix B Deblender architecture

The deblender employs a simple VAE-like architecture that combines convolutional feature extraction with lightweight selfattention near the latent space. The encoder consists of six convolutional blocks with progressively increasing channel depth (from 16 to512) and downsampling from 128 × 128 to 4 × 4. Each block uses PReLU activations, dropout for regularisation, and batch normalisation at the dense stage. Two Multi-Head SelfAttention (MHA) layers are inserted after the 8 × 8 and 4 × 4 feature maps to capture global dependencies. The flattened features are linearly projected to a 512-dimensional representation, which is subsequently mapped to the latent mean μ and variance σ, each of dimension 16.

The decoder mirrors this process with linear up-projection followed by six transposed convolution layers that gradually reconstruct the 128 × 128 image. A third MHA layer is placed after the second transposed convolution to enhance global coherence during reconstruction. All layers use PReLU activations except the final transpose convolution, which employs ReLU. The full architecture is illustrated in Figure B.1. For the VAE-CNN model, we simply remove the three MHA layers. For the CAT deblender, we use the same UNet structure proposed in (Zhang et al. 2024) but set the number of channels to 1.

Appendix C Effect of facet size

We evaluate the impact of facet size on both reconstruction quality and computational cost. To this end, we randomly select 128 galaxies from DeepSet and generate faceted dirty images of sizes 64 × 64, 128 × 128, 256 × 256, and 512 × 512 directly from the simulated visibility measurement set. For each case, reconstruction is then performed using HQS-PnP.

To assess reconstruction quality, we compare the reconstructed images with the corresponding ground truth images (blended stamps) using PSNR. For a fair comparison across different facet sizes, both the reconstructed and ground truth images are centrally cropped to 64 × 64, ensuring that the evaluation focuses on the target source. We do not perform deblending or downstream shape measurement in this experiment, as the goal here is to isolate the effect of facet size on the deconvolution step alone.

In addition, we measure the total runtime per galaxy, including the time required for gridding the visibilities at different resolutions, forming the dirty image and PSF, and performing reconstruction using HQS-PnP. The results are summarised in Fig. C.1.

As shown in the figure, increasing the facet size leads to a significant rise in computational cost. This is primarily driven by the increased cost of gridding on finer uv grids and the repeated application of the denoiser on larger images.

In terms of reconstruction quality, we find that 128 × 128 facets achieve the best performance. This is expected, as the denoiser is trained on images of this size and is therefore optimally tuned for this input scale. Other facet sizes yield only marginally lower PSNR, indicating that the method is reasonably robust to moderate changes in input size. However, for larger facet sizes, we observe a slight degradation in performance. This can be attributed to the multi-scale structure of the UNet denoiser, where the effective receptive field and learned feature hierarchies are adapted to a specific input scale. When applied to significantly larger images, the denoiser may not capture finescale structures with the same fidelity, leading to suboptimal reconstruction.

Fig. B.1 Architecture of the deblender network. The width of each block represents the number of feature channels, while the horizontal length corresponds to the spatial resolution.

Fig. C.1 Total runtime per galaxy (left axis) and median PSNR of reconstructed images (right axis) as a function of facet size.

Overall, these results support the choice of 128 × 128 facets, as they achieve the best reconstruction quality while also maintaining significantly lower computational cost.

All Tables

All Figures

	Fig. 1 Example sky patch from WideSet. For visualisation, only the central 0.1 deg2 region is shown. The colourmap is scaled as log10(S + 10−9), where S denotes the flux density of each pixel in units of Jy per pixel, to enhance contrast and improve visual clarity.
In the text

	Fig. 2 Examples of four galaxies from the DeepSet, arranged by increasing blendedness B. Top row : blended inputs (titles show B). Second row : isolated stamps (titles show flux of the central galaxy). Bottom rows : reconstructions from the three networks, with titles indicating PSNR and ellipticity error ∆ε for each stamp.
In the text

	Fig. 3 Comparison of reconstruction metrics for the three models. Top : PSNR. Bottom : ellipticity error ∆ε (scaled by 100 for clarity). Each curve shows the KDE, the dashed line indicates the median, and the legend reports the median with the IQR. The shapes are measured using GALSIM.
In the text

	Fig. 4 Deblending performance as a function of blendedness for the three models. The top panels show binned box plots for the three models in terms of PSNR and ellipticity error, while the bottom panel shows the one-dimensional distribution of blendedness. The shapes are measured using GALSIM.
In the text

Fig. 5 Examples of four galaxies from the DeepSet. Top row : isolated stamps (titles indicate flux of the central galaxy). Second row : blended inputs (titles show B). Third row : dirty facet images. Fourth row : deconvolved stamp output by HQS-PnP. Fifth row : Residual facet images. Bottom row : final reconstructed stamps, with titles reporting PSNR and ∆ε for each stamp. The shapes are measured using GALSIM.

In the text

	Fig. 6 Reconstruction of isolated galaxy stamps. Top row : dirty images generated from the same galaxies as in Figure 5. Bottom row : corresponding reconstructions using HQS-PnP.
In the text

Fig. 7 Dependence of reconstruction fidelity on galaxy properties. The top row shows PSNR and the bottom row shows ∆ε. The first column shows the median metric in bins of peak flux. The second column shows the median metric in bins of galaxy size, with shaded regions indicating the IQR. The third column shows the KDE of each metric; a dashed line marks the median. The shapes are measured using GALSIM.

In the text

	Fig. 8 Shape measurement residuals for the isolated and wide field cases as a function of the input ellipticity for the first component ε1. The two-dimensional contours are estimated using kernel density estimation. The coloured dashed lines indicate the best-fit linear relations for each case. The panels on the right show the corresponding 1D marginal distributions of the residuals.
In the text

	Fig. 9 Estimated bias slope as a function of factor f.
In the text

	Fig. 10 Dependence of the ellipticity error as a function of total flux. The full DeepSet sample is used, without imposing flux or peak cuts, to assess performance across the full dynamic range.
In the text

	Fig. B.1 Architecture of the deblender network. The width of each block represents the number of feature channels, while the horizontal length corresponds to the spatial resolution.
In the text

	Fig. C.1 Total runtime per galaxy (left axis) and median PSNR of reconstructed images (right axis) as a function of facet size.
In the text