Pty-Chi: A PyTorch-based modern ptychographic data analysis package

TL;DR

Ptychography reconstruction benefits from both analytical and gradient-based approaches; Pty-Chi merges these within a unified PyTorch framework, enabling differentiable forward models, multi-GPU parallelism, and modular components for advanced corrections like orthogonal probe relaxation and multislice modeling. The package provides a broad set of reconstructors (PIE, LSQML, DM, BH, and AD-based) and supports extended workflows such as deep image priors and plug-and-play ADMM, demonstrated through case studies under limited coherence, OPR, and data decimation. Benchmarking on a multi-GPU server confirms scalable parallelization and varied performance profiles across reconstructors, highlighting the trade-offs between speed, memory, and accuracy. The work emphasizes extendability, vendor-agnostic acceleration, and community-driven development to advance computational ptychography at synchrotron facilities and beyond.

Abstract

Ptychography has become an indispensable tool for high-resolution, non-destructive imaging using coherent light sources. The processing of ptychographic data critically depends on robust, efficient, and flexible computational reconstruction software. We introduce Pty-Chi, an open-source ptychographic reconstruction package built on PyTorch that unifies state-of-the-art analytical algorithms with automatic differentiation methods. Pty-Chi provides a comprehensive suite of reconstruction algorithms while supporting advanced experimental parameter corrections such as orthogonal probe relaxation and multislice modeling. Leveraging PyTorch as the computational backend ensures vendor-agnostic GPU acceleration, multi-device parallelization, and seamless access to modern optimizers. An object-oriented, modular design makes Pty-Chi highly extendable, enabling researchers to prototype new imaging models, integrate machine learning approaches, or build entirely new workflows on top of its core components. We demonstrate Pty-Chi's capabilities through challenging case studies that involve limited coherence, low overlap, and unstable illumination during scanning, which highlight its accuracy, versatility, and extensibility. With community-driven development and open contribution, Pty-Chi offers a modern, maintainable platform for advancing computational ptychography and for enabling innovative imaging algorithms at synchrotron facilities and beyond.

Figures (7)

• Figure 1: Main building blocks and layers of Pty-Chi.
• Figure 2: The UML diagram of Pty-Chi.
• Figure 3: Reconstructions delivered by the LSQML reconstructor of Pty-Chi with various settings. (a) Random batch sampler. (b) Compact batch sampler. (c) Compact batch sampler with momentum acceleration. The former two cases exhibit halo around the object and horizontal stripe artifacts around the wall of the capillary tube. A clean and artifact-free reconstruction was only obtainable with the configurations in (c). Due to the large dynamic range of the reconstructed phase in (c), features inside the pipette are presented with seemingly reduced contrast. We show a selected portion of the features with smaller dynamic range in the inset to show that the real contrast is preserved.
• Figure 4: Reconstructions delivered by LSQML with (a) and without (b) OPR. The reconstruction without OPR exhibit ripple-like artifacts [highlighted by dashed boxes in (a)] due to the variation of illumination across scan lines. By modeling this variation with OPR, the artifacts are effectively suppressed.
• Figure 5: An illustration of reconstruction based on deep image prior. Instead of directly solving the object function, Pty-Chi's AD reconstructor solves for the parameters of a neural network that maps a fixed noise image into the supposed object. The rest of the forward model stays the same as traditional model-based reconstruction.