Parallel ptychographic reconstruction

Abstract

Ptychography is an imaging method whereby a coherent beam is scanned across an object, and an image is obtained by iterative phasing of the set of diffraction patterns. It is able to be used to image extended objects at a resolution limited by scattering strength of the object and detector geometry, rather than at an optics-imposed limit. As technical advances allow larger fields to be imaged, computational challenges arise for reconstructing the correspondingly larger data volumes, yet at the same time there is also a need to deliver reconstructed images immediately so that one can evaluate the next steps to take in an experiment. Here we present a parallel method for real-time ptychographic phase retrieval. It uses a hybrid parallel strategy to divide the computation between multiple graphics processing units (GPUs) and then employs novel techniques to merge sub-datasets into a single complex phase and amplitude image. Results are shown on a simulated specimen and a real dataset from an X-ray experiment conducted at a synchrotron light source.

Fig. 1

Simplified ptychography experiment setup showing a Cartesian grid used for the overlapping raster scan positions.

Fig. 2

GPU memory layout of a 2D M×N matrix. The gray area represents the aligned matrix zero padding, where W = warp size (32 threads on most devices). The padded area is always ignored in matrix operations. CUDA blocks operate on matrix slices of size s such as the region outlined by the dashed red line. Memory requests for a row of the padded matrix are served by a single cache line.

Fig. 3

The diffraction patterns are subdivided and distributed among available GPUs. Pairwise stitching is performed on the separate reconstructions attained by phase retrieval.

Fig. 4

Object array sharing through neighborhood exchange between 16 GPUs. The overlap (halo), highlighted in blue, is defined in terms of additional scan points assigned to each GPU sub-dataset.

Fig. 5

(a) Normalized RMS error of final reconstructions achieved by different GPU configurations using the asynchronous, synchronous, and synchronous with halo=2 implementations. (b) Magnitude and phase of the object wavefront retrieved from simulated data using the asynchronous version and 128 GPUs. (c) Magnitude and phase of the object wavefront retrieved from simulated data using the synchronous version and 32 GPUs. (d) Magnitude and phase of the object wavefront retrieved from simulated data using the synchronous version, 32 GPUs, and a halo region of 2 additional scan point rows and columns.

Fig. 6

Performance plots on synthetic data. Left: Total running time (in seconds) of different GPU configurations. Right: The scaling efficiency plotted as a percentage of linear scaling.

Fig. 7

(a) Phase image of the reconstructed object transmission function, with the 1 μm×1 μm scan region highlighted in red. (b) Fourier ring correlation (FRC) plot showing a spatial resolution of 16 nm in the phase of the exit surface wave. (c,d) The recovered illumination function of two probe modes.

Fig. 8

Performance plots on real data. Left: Total running time (in seconds) of different GPU configurations. Right: The scaling efficiency plotted as a percentage of linear scaling.