Towards generalizable deep ptychography neural networks

TL;DR

This work tackles the robustness gap in deep-learning-based ptychography by introducing a probe-centric, unsupervised training workflow that pairs experimentally measured probes with synthetic objects. The proposed PtychoPINN-torch integrates physics-informed inductive biases, overlap constraints, and a multi-probe learning regime to achieve generalization across multiple beamlines and instruments, with a synthetic data regime as small as ~28k examples achieving competitive performance. The study demonstrates single-experiment transfer, multi-probe representations, and frequency-sensitive effects of synthetic objects, showing meaningful reconstruction fidelity improvements and substantial real-time inference speedups compared to conventional iterative methods. The approach promises practical deployment of experiment-steering models capable of real-time feedback under dynamic conditions, while highlighting probe learning as a dominant factor in performance and paving the way for larger probe libraries and more expressive architectures.

Abstract

X-ray ptychography is a data-intensive imaging technique expected to become ubiquitous at next-generation light sources delivering many-fold increases in coherent flux. The need for real-time feedback under accelerated acquisition rates motivates surrogate reconstruction models like deep neural networks, which offer orders-of-magnitude speedup over conventional methods. However, existing deep learning approaches lack robustness across diverse experimental conditions. We propose an unsupervised training workflow emphasizing probe learning by combining experimentally-measured probes with synthetic, procedurally generated objects. This probe-centric approach enables a single physics-informed neural network to reconstruct unseen experiments across multiple beamlines; among the first demonstrations of multi-probe generalization. We find probe learning is equally important as in-distribution learning; models trained using this synthetic workflow achieve reconstruction fidelity comparable to those trained exclusively on experimental data, even when changing the type of synthetic training object. The proposed approach enables training of experiment-steering models that provide real-time feedback under dynamic experimental conditions.

Figures (19)

• Figure 1: An overview of the data and PtychoPINN-torch training process. (a) Example experiment datasets, shown by their reconstructed objects (gray) and probes (pink). Note that only the phase structure in the center of the phase image associated with non-zero amplitude values (upper probe image) are valid.(b) Examples of synthetic object classes used to create synthetic training data. (c) The training process for PtychoPINN-torch. Training data is either sourced from experiment measurements or from synthetic data created from combining synthetic objects and reconstructed experiment probes. The data format containing diffraction images and probe positions is used directly for both training and inference. Training uses the full model -- the neural network combined with physics-based constraints -- while inference uses only the neural network mapping component.
• Figure 2: Comparison of models trained on synthetic (PS and SS) and experimental data (PE). a) From left to right: Reconstructions of datasets TP2, IC2, NCM and W from models trained using only TP2 raw data. PS_TP2 exceeds or equals PE_TP2 on test datasets from the same instrument (NCM, IC2), while a supervised model, SS_TP2, transfers poorly across the board. All models generalize poorly on the W dataset, whose probe function differs greatly from the training probe function (due to being from a different instrument) and cannot be generalized to. b) Enlarged regions of the TP2 (top) and IC2 (bottom) datasets, with ground truth, PE and PS reconstructions. PE misses high frequency features in both training set (TP2) and test set (IC2). c) FRC comparison of 3 models on the IC2 dataset, with baseline PE_IC2 (model trained exclusively on IC2 experimental data). PS_TP2, despite being trained on a probe from a different experiment, performs nearly identically. PE_TP2 performs the worst of the three, due to limited object variety in the TP2 dataset which does not generalize to the IC2 dataset.
• Figure 3: a) Schematic representation of probe-dependent conditional mappings for our training scenarios, where learned mappings for individual probes are largely independent of each other. As additional probes are added to the training dataset, the model is forced to learn a joint mapping that generalizes across all training probes. b) Reconstructions for datasets (instruments): W (HXN), FLY1 (Velociprobe), IC2 (Velociprobe) and LFP (Cosmic) under 3 training schemes with an increasing number of training probes from left to right. Under probe-excluded single training, models perform poorly when training and testing probe differ. When the testing probe is added to training in probe-included dual, the reconstruction quality remains high, up to 4 distinct training probes (unified multi-probe). c) FRC-AUC scores organized by experimental dataset per row (see b). Bars represent probe-excluded single training (gray), dual-probe training (orange), multi-probe training (green), and test probe-only training (blue). X-axis labels describe probes used in probe-excluded single training.
• Figure 4: a) Images of synthetic objects dead leaves (DL), blurred white noise (BWN), procedural (PR) and simplex noise (SN). Each object image is accompanied by a 2-dimensional power spectral density (PSD) plot with an inset of the 1-dimensional integrated (PSD), showing differences in frequency statistics amongst synthetic images. b) Reconstruction results on datasets FLY1, W, IC2 and LFP for models trained exclusively on one synthetic object class. Each row corresponds to the synthetic class shown in a. c) Example reconstruction and PSD for simplex noise, showing high-frequency encoding via the probe rather than object. d) FRC curves of reconstructions from b showing differences in reconstruction quality reflecting the frequency statistics of the training datasets.
• Figure S1: Ablation study results of PtychoPINN-torch's inductive biases. Each row represents a training or testing dataset trained with the following models from left to right: supervised, PINN, vanilla PtychoPINN and PtychPINNv2.