Vision Transformer for Multi-Domain Phase Retrieval in Coherent Diffraction Imaging

TL;DR

This work tackles the challenging strong-phase, multi-domain phase retrieval problem in Bragg coherent diffraction imaging by introducing an unsupervised Fourier ViT that performs global spectral token mixing directly on diffraction magnitudes. The model combines a shallow CNN front-end, multiscale Fourier attention, and a CNN decoder to reconstruct complex real-space density (amplitude and phase) within a fixed support, trained with a hybrid Fourier-space loss. On synthetic Voronoi-domain data, Fourier ViT achieves low or near-zero reciprocity-space error and can resolve up to 19 domains, while maintaining robustness under realistic noise; experiments on STO and La$_{2-x}$Ca$_x$MnO$_4$ demonstrate competitive to superior performance relative to iterative and CNN baselines, with comparable or better chi-squared and high fidelity to domain structures. Overall, Fourier ViT provides a fast, unsupervised pathway to reliable multi-domain BCDI reconstructions, offering robustness to initialization and practical utility for in situ analyses, with future directions including explicit partial-coherence modeling and uncertainty quantification.

Abstract

Bragg coherent diffraction imaging (BCDI) phase retrieval becomes rapidly difficult in the strong-phase regime, where a crystal contains distortions beyond half a lattice spacing. An important special case is the phase domain problem, where blocks of a crystal are displaced with sharp jumps at domain walls. The strong-phase, here defined as beyond $\pm π/2$, generates split Bragg peaks and dense fringe structure for which classical iterative solvers often stagnate or return different solutions from different initialisations. Here, we introduce an unsupervised Fourier Vision Transformer (Fourier ViT) to solve this block-phase, multi-domain phase-retrieval problem directly from measured 2D Bragg diffraction intensities. Fourier ViT couples reciprocal-space information globally through multiscale Fourier token mixing, while shallow convolutional front and back-ends provide local filtering and reconstruction. We validate the approach on large-scale synthetic datasets of Voronoi multi-domain crystals with strong-phase contrast under realistic noise corruptions, and on experimental diffraction from a $\mathrm{La}_{2-x}\mathrm{Ca}_x\mathrm{MnO}_4$ nanocrystal. Across the regimes considered, Fourier ViT achieves the lowest reciprocal-space mismatch ($χ^2$) among the compared methods and preserves domain-resolved phase reconstructions for increasing numbers of domains. On experimental data, with the same real-space support, Fourier ViT matches the iterative benchmark $χ^2$ while improving robustness to random initialisations, yielding a higher success rate of low-$χ^2$ reconstructions than the complex convolutional neural network baseline.

Figures (7)

• Figure 1: Architecture of the proposed Fourier ViT model for BCDI phase retrieval. (a) The input diffraction magnitude ($64\times64$ pixels) is passed through (b) a shallow convolutional feature extractor to produce 128 channels ($128\times64\times64$). (c) Each image is partitioned into patches and embedded as tokens ($8\times8$ patches shown for clarity; all experiments use $16\times16$ patches), giving a token sequence of shape $(64\times128)$. This sequence is processed by (d) a multi-layer Vision Transformer with multi-scale Fourier attention (frequency mixing at 1:4, 1:2 and 1:1 scales) and (e) an FFT-based global convolution block. (f) The transformer output is reshaped into a latent feature map ($128\times8\times8$) and decoded by (g) a convolutional upsampling path with a skip connection to reconstruct the complex crystal field, yielding (h) $64\times64$ real-space amplitude and phase maps consistent with the measured diffraction pattern.
• Figure 2: Fourier ViT reconstruction of a synthetic 10-domain crystal from diffraction only. (a,d,g) Ground truth diffraction magnitude (64×64 pixels), crystal amplitude and phase (40×40 pixels); (b,e,h) corresponding Fourier ViT reconstructions. (c) Radial log intensity against $q$ profiles of ground truth and reconstructed diffraction, showing agreement across the full $q$-range. (f,i) Histograms of diffraction $\chi^{2}$ over 100 random initialisations for known crystal amplitude (f) and unknown amplitude (i) phase retrieval.
• Figure 3: Noise robustness of the Fourier ViT. Columns show three types of noise applied to the same diffraction pattern: Gaussian noise, Poisson noise, and partial coherence. For each model, L1, L3 and L5 denote increasing noise strengths (levels 1, 3 and 5); green percentages give $\chi^2_{\mathrm{n}}$ (in %), the input diffraction error relative to the clean ground truth. Top row shows the noisy input diffraction amplitudes (64×64 pixels). Middle and bottom rows show the reconstructed crystal amplitude and phase zoomed in 40×40 pixels, respectively; both are masked to the support and phase is in radians.
• Figure 4: Comparison study of phase retrieval on experimental data from a multi-domain LCMO-500 nanocrystal. (a) Measured BCDI diffraction magnitude (64×64 pixels). (b--d) Reconstructed crystal's amplitude and (e--g) phase from the iterative method, Fourier ViT, and C-CNN. Scale bar is 500 nm. (h--j) Corresponding $\chi^{2}$ histograms over 200 runs per method, each run with a random initial phase. $\chi^{2}$ range is between $0.0\%$ and $2.0\%$.
• Figure 5: Simulation pipeline for synthetic crystals and diffraction patterns. (a) Random seed placement within the support. (b) Voronoi phase domain generation from seeds, with phases wrapped to $[-\pi,\pi]$. (c) Construction of the crystal amplitude mask, max-normalised to $[0,1]$. (d) Corresponding diffraction magnitude computed by fast Fourier transform (FFT), max-normalised to $[0,1]$. All images are $64\times64$ pixels.