ptychopy: GPU framework for ptychographic data analysis

The architecture of the software is based on the model-view-controller (MVC) model, which includes three components, a user view interface for presenting and accepting information, a backend module for processing the information, and a controller component to manage the data, rules and logic to connect the interface to the backend module, as shown in Fig. 1. The user view interface and controller component are implemented using python and the backend modules are implemented with CPython, C++, and CUDA. It provides python APIs for doing ptychographic reconstruction using three algorithms ePIE, DM, LSQML on GPU.

The frontend component is implemented with the PyQt5 graphics toolkit for the GUI and visualization functionality. The major functionalities include providing input for the reconstruction parameters and visualizing the reconstruction results (see Fig. 1). For the reconstruction parameters, there are four categories, which are:

• •

  Provide the interface for the experimental parameters such as beam energy, the guess of probe size, the number of probe modes, detector distance, scan dimensions and step size, and axis for flipping the scan direction.

• •

  Define the parameters for preprocessing diffraction data, such as the center of the diffraction patterns, the cropping array size, the pixel bit-depth for detecting saturated pixels, the threshold value to remove the noise, the rotation angles of the diffraction patterns, and the number of diffraction patterns per data file.

• •

  Specify the parameters for the multiple GPU configuration, such as the GPU node index, the overlapping size where the data is shared between multiple GPUs, the share-frequency of data shared among multi-GPUs.

• •

  Set up the parameters for phase retrieval, such as the reconstruction algorithm (currently ePIE, DM and LSQML are available), and the number of iterations for phase retrieval, after which iteration to update the probe. The probe and/or object guess can also be specified in the GUI if they are available.

Most of the parameters for experiment scans and diffraction patterns are saved in MDA files by EPICS, which then can be directly imported into our GUI interface so that the users do not have to enter the parameters one by one. The editable text boxes on the GUI also allow users to tune the parameters to do customized reconstructions.

The ePIE, DM and LSQML algorithms use a set of far-field diffraction pattern data to reconstruct both the probe function and the object function simultaneously. The iterative procedure involves the computation of multiple Fast Fourier Transforms (FFTs) over thousands of diffraction patterns, making it suitable for parallel analysis on GPU. Compared to DM, ePIE needs a smaller memory footprint than DM, since DM has to keep a copy of the previous calculation in memory and calculate the exit wavefronts of all the scan points at once. But DM would run faster than ePIE on the reconstruction since DM can parallelize the calculation of the scan points at the same time on GPUs. LSQML’s memory footprint requirement is between ePIE and DM while taking most time among three algorithms to produce the best quality of the images. Depending on the requirement on the reconstruction speed and the quality of the resultant image, the user could choose ePIE, DM or LSQML algorithm based on their own need.

The size of the ptychography datasets is increasing rapidly due to the developments of highly efficient scanning and fast detectors, which brings up the need of massive computation power. Homogeneous many-core processors with massively parallel structure, such as GPUs, are emerging in broad application areas and has been providing massive computation power for processing large blocks of data in parallel. CUDA is a parallel computing platform and programming model created by NVIDIA, which gives the developers access to the virtual instruction set and memory of the parallel computational elements in GPUs via a C programming environment. GPU programs using CUDA have been widely implemented on many scientific problems and have been improving the program performance dramatically. Therefore, in the ptychopy reconstruction module, we implemented the algorithms with C and CUDA as the backend for the performance enhancement purposes.

The backend module is a CPython module which provides ePIE, DM, and LSQML reconstruction algorithm options for phase retrieval. All the reconstruction algorithms are implemented using CUDA and C++ for performance consideration. CPython is a python interpreter implemented with C, providing the python interface for the CUDA code and also serves as the glue for the data reconstruction backend and other modules that are implemented in a python environment.

The controller component uses threads for pipelining the reconstruction procedure and visualization procedure. The main thread is responsible for accepting the parameters from the GUI and reading the data from the HDF5 file. The worker thread is calling the backend reconstruction module and saving the results as CSV files. In addition, the controller component handles all the data saving, loading, and pipelining between the GUI and the reconstruction.

The data collected at APS synchrotron beamlines is saved in HDF5 format and the parameters for the ptychography scan (such as energy, scan dimensions, step size, and rotation angle) and detector parameters are saved in the scan MDA file. ptychopy will read these parameters as default from the MDA file via its controller module. These default parameters can also be changed with customized ones if they are specified in the GUI. For batch reconstruction, the MDA files of the selected scans can be loaded as a list and the controller will read the MDA file list and corresponding parameters in a queue. In addition, to reduce the collected data size, the HDF5 datasets are compressed (by LZ4 or zlib) during data collection, which are automatically uncompressed when the data is imported for reconstructions. The collected data are usually a set of diffraction patterns and each HDF5 saves certain number of diffraction patterns. At APS beamlines, we save each row of the scan diffraction patterns into a separate file and use the row number as the HDF5 filename suffix. Our data analysis module of ptychopy will read the diffraction patterns from the HDF5 with the corresponding row and column numbers as the scan positions into GPU memory. The number of rows and columns of the scan positions can be defined using scanDims parameters showed in the following script example.

Ptychopy provides flexibility in that it can be used in either GUI or script mode. The interface of the GUI mode shown in Fig. 1 conveniently allows users who have little command line interface experience to specify the scan and reconstruction parameters by choosing the scan MDA files and/or editing those parameters in the GUI. If the parameters are not specified, the default value saved in the MDA file will be used for the reconstruction. After the reconstruction, the final result will be saved in CSV files, and meanwhile, the reconstructed results will be shown on the visualization window. The reconstruction results can also be viewed anytime using the image browser window on the GUI.

In the script mode, the module may be integrated with beamline controls (such as PyEpics) and data management packages to achieve streamlined data collection and processing. The APIs for ePIE, DM and LSQML are illustrated in the Fig. 3 for running as the whole mode or step mode. For whole mode, the reconstruction will run as a whole without interruption after passing parameters as para to the function. For the step mode, the result during the reconstruction could be evaluated for a specified number of iterations. Taking ePIE as the example, para in the function epie(para) indicates the reconstruction parameters which are passed to the function and epie(para) will handle all the reconstruction on the GPU. If the value needs to be checked during the reconstruction, the step APIs in the Fig. 3 could be used. Since the reconstruction algorithm is running on the GPU, retrieving the result from GPU needs to transfer the value back to CPU side and will take more time than running ePIE as a whole. If performance matters most, the whole mode should be used instead of the step mode for running the algorithm.

Listing 1 shows a script example using the ePIE algorithm with a simulated image about how to use ptychopy module in the python script. The simulated images and test scripts are found in the source packages. As shown in this example, the module accepts a string with the reconstruction parameters and returns the reconstruction results. The parameters can be found correspondingly in the GUI as a hint text box when one moves the cursor over a GUI component.

The first step is to initialize the reconstruction process and GPU device status with the epieinit() function. This includes initializing reconstruction parameters, IO processing, data loading and copying, initializing sample and probe array, copying the data from CPU side to GPU side, and preprocessing the data, such as diffraction pattern normalization. Then the number of iterations is defined and the ePIE reconstruction is launched with the epiestep() function. The phase retrieval procedure is running on the GPU side for a specified number of iterations. Before the reconstruction is done, the updated object and probe functions still reside on the GPU. To get the object and probe results, epieresobj() and epieresprobe() can be called anytime to return the data as a numpy array. Therefore, for the live view of the reconstruction process the updated probe and object results can be displayed at a defined iteration frequency from the GPU side. After the reconstruction, epiepost() function should be called to clean up the whole process and GPU device.

For machines with multiple GPUs, ptychopy provides the feature to run reconstruction jobs on multiple GPUs in two ways. As shown in Fig. 1, the list box provides a list of GPUs that the destination machine has, and the reconstruction job will be put into a queue which is associated with each chosen GPU in the GUI. And then the GPU will use a first-in first-out order to process the assigned job as shown in Fig. 2. This one-to-one GPU allocation is very useful for a large number of jobs with small data size (one GPU can handle one scan). In addition, the job could also run on multiple GPUs and scans with data size larger than one GPU memory would benefit from this. Either way, the controller uses both a worker thread and main thread for controlling the job allocation, running, and visualization. For the first way, the number of worker threads depends on the number of GPUs chosen in the GUI for reconstruction. For each GPU, the controller will create a worker thread for the specified GPU and a separate job queue. For the second way, the reconstruction job will be all put into one thread and use all the GPUs available. The worker threads will record all the job running information and are responsible for triggering the backend CPython module for the reconstruction whenever there are still jobs left in the queue for reconstruction.

Ptychopy has a logging feature to save the reconstruction procedure and the reconstruction parameter information in real-time or after analysis. The information will be redirected to STDOUT and saved under the log folder as ptychopy.log. After the reconstruction, the parameter information and the reconstructed result images could be sent to the email account that is specified on the GUI email address text box. The result image is converted to 8-bit TIFF format in advance, since the resulted reconstruction image could be more than the limited size for the email attachment.