PtyRANNOSAUR: Ptychography with Robust Artificial Neural Networks Optimized for Sub-Angstrom Accuracy and Ultrafast Reconstruction

Technical Summary of PtyRANNOSAUR

Problem Statement
Electron ptychography is a powerful technique for recovering phase images of 2D or 3D objects from 4D scanning transmission electron microscopy (STEM) datasets, enabling deep sub-angstrom resolution and high dose efficiency. However, its widespread adoption is currently hindered by the immense computational cost of data processing. State-of-the-art reconstruction relies on iterative physics-based solvers (e.g., PtyRAD) that utilize automatic differentiation to optimize forward models of electron scattering. While accurate, these methods require hundreds to thousands of iterations, often taking hours on modern GPUs (e.g., NVIDIA A100). This latency precludes real-time sample feedback during data acquisition and limits experimental throughput. Conversely, existing single-shot machine learning (ML) methods offer speed but fail to account for complex electron scattering physics, microscope optical variations, or scan position errors, resulting in inferior resolution and quality compared to iterative methods. There remains an unmet need for a reconstruction method that operates at near-real-time speeds while maintaining the accuracy, generalizability, and resolution of iterative solvers across diverse material systems and experimental conditions.

Methodology
The authors introduce PtyRANNOSAUR, a data-driven neural network framework designed to reconstruct atomic-resolution electron ptychography data in seconds. The core methodology involves training specialized convolutional autoencoders to map 4D-STEM data (specifically, a 5×5 grid of convergent beam electron diffraction patterns) directly to 2D projected phase images.

• Training Data Generation: Instead of using stock images or simple simulations, the authors generated a massive database of 6.4 million input-target pairs. The targets are 2D projected phase images (3 Å × 3 Å) derived from 16,000 stable crystal structures sourced from the Materials Project. The inputs are corresponding sets of 25 simulated CBED patterns generated using multislice electron scattering simulations in abTEM.
• Physics-Informed Simulation: The training data explicitly incorporates complex physical phenomena, including multiple scattering (via the multislice method), spatial partial coherence (modeled via Gaussian convolution to simulate source size), and variations in experimental parameters such as accelerating voltage, convergence angle, defocus, and sample thickness. This ensures the network learns the inverse mapping for realistic, imperfect experimental conditions rather than idealized scenarios.
• Model Specialization: Three distinct models were trained to cover different experimental regimes:
  • Model 2: Optimized for ultra-thin samples (e.g., 2D materials) at 80 kV with a thickness of 2 ± 0.6 nm.
  • Model 10 & Model 20: Optimized for standard TEM lamella at 300 kV with nominal thicknesses of 10 ± 2 nm and 20 ± 2 nm, respectively.
• Loss Function: The networks are trained using the Structural Similarity Index Measure (SSIM) as a loss function, which quantifies luminance, contrast, and structural similarities between the network output and the ground truth phase image.
• Inference and Stitching: During inference, the network processes overlapping local patches of the 4D-STEM data. These patches are stitched together to form the full image. To address stage drift and probe position errors common in experimental data, the authors implemented an optional post-processing position correction module using phase cross-correlation, which runs in approximately 20 seconds on a CPU.

Key Contributions

1. Ultrafast Reconstruction: PtyRANNOSAUR achieves reconstruction speeds of approximately 5 seconds per dataset on a single NVIDIA A100 GPU, representing a 10–100x speedup over standard iterative methods (which take ~30 minutes).
2. Robustness to Experimental Variability: The models demonstrate significant tolerance to variations in defocus (±8 nm), scan step size (±0.12 Å), sample thickness (±10 nm), and convergence angle (±2 mrad), eliminating the need for the manual hyperparameter tuning typically required by iterative solvers.
3. Handling of Complex Physics: Unlike previous single-shot ML approaches, PtyRANNOSAUR successfully accounts for multiple scattering and spatial partial coherence, enabling accurate reconstructions of thicker samples (up to ~25 nm) and complex geometries.
4. Accessibility: The authors provide trained models via a user-friendly Jupyter notebook interface, allowing users to perform reconstructions without simulating data or training new models, while also offering the capability to train custom models for specific microscope configurations.

Results
The authors validated PtyRANNOSAUR on both experimental and simulated datasets:

• Experimental Performance: The method was tested on diverse specimens including bulk sapphire, Al/AlOx/Al thin film multilayers, gold nanoparticles on amorphous carbon, and twisted bilayer MoS2. In all cases, PtyRANNOSAUR produced atomic resolution images comparable to, and in some cases exceeding, the quality of iterative PtyRAD reconstructions. Notably, for twisted bilayer MoS2, the network resolved atomic columns spaced 0.4 Å apart, achieving a resolution of < 0.5 Å, which was barely resolved in the iterative reconstruction.
• Defect Detection: The network successfully identified single-atom defects, such as sulfur vacancies in MoS2, without requiring material-specific training.
• Simulated Validation: On a validation set of 2,950 unseen materials, the median material-level SSIM loss dropped to 0.020 after stitching, with the network faithfully reproducing periodic lattice structures and relative intensities.
• Robustness: The models maintained high accuracy across the tested parameter ranges. While performance degraded outside these ranges (e.g., extreme defocus or thickness), the network still captured the coarse structure, acting as a rapid diagnostic for data quality issues.

Significance and Claims
The paper claims that PtyRANNOSAUR enables "near-live, state-of-the-art electron ptychography reconstructions." By reducing the computational bottleneck from hours to seconds, the framework allows for real-time feedback during data acquisition, enabling users to locate features of interest and optimize acquisition parameters on the fly. The authors position the tool not only as a standalone reconstruction method but also as a generator of priors to accelerate iterative solvers or as a rapid screening tool for large datasets (e.g., in-situ experiments or tomography). The work demonstrates that deep learning can effectively learn the complex, non-linear inverse problem of electron ptychography, achieving deep sub-angstrom resolution (< 0.5 Å) with robustness to experimental imperfections, thereby democratizing access to high-resolution microscopy.