PtyRAD: A High-performance and Flexible Ptychographic Reconstruction Framework with Automatic Differentiation

TL;DR

PtyRAD is an open-source, high-performance framework for electron ptychography that leverages automatic differentiation and optimization techniques to improve reconstruction speed, accuracy, and ease of use, facilitating broader adoption and research.

Contribution

The paper introduces PtyRAD, a flexible, efficient, and open-source software framework that integrates automatic differentiation and advanced optimization for electron ptychography reconstruction.

Findings

Achieves up to 17x faster reconstruction times compared to existing tools.

Provides accurate reconstructions without manual parameter tuning.

Includes novel real-space depth regularization for complex datasets.

Abstract

Electron ptychography has recently achieved unprecedented resolution, offering valuable insights across diverse material systems, including in three dimensions. However, high-quality ptychographic reconstruction is computationally expensive and time consuming, requiring a significant amount of manually tuning even for experts. Additionally, essential tools for ptychographic analysis are often scattered across multiple software packages, with some advanced features available only in costly commercial software like MATLAB. To address these challenges, we introduce PtyRAD, an open-source software framework offers a comprehensive, flexible, and computationally efficient solution for electron ptychography. PtyRAD provides seamless optimization of multiple parameters--such as sample thickness, local tilts, probe positions, and mixed probe and object modes--using gradient-based methods with automatic differentiation (AD). By utilizing PyTorch's highly optimized tensor operations, PtyRAD achieves up to a 17x speedup in reconstruction time compared to existing packages without compromising image quality. In addition, we propose a real-space depth regularization, which avoids wrap-around artifacts and can be useful for twisted two-dimensional (2D) material datasets and vertical heterostructures. Moreover, PtyRAD integrates a Bayesian optimization workflow that streamlines hyperparameter selection. We hope the open-source nature of PtyRAD will foster reproducibility and community-driven development for future advances in ptychographic imaging.