As the spatial resolution (r) of an imaging system improves, the depth of focus (DOF) decreases as r²/ λ.  As Argonne’s APS Upgrade enables higher-resolution imaging, researchers must account for diffractive blurring in the image of the thick samples that X-rays will be able to image. In this project, we use Fresnel multislice or finite-difference methods to account for wave propagation using an initial guess of an object under study, and then solve for the object using a model-based optimization approach. We compute the gradient values, which are crucial in numerical optimization methods, using automated and flexible methods represented by automatic differentiation (AD).  Because AD is used for network training, it is part of many AI toolkits that are already built for large-scale data handling on supercomputers.  We have completed a detailed comparison of the forward model approaches (Fresnel multislice versus finite-difference methods), showing that our AD-based approach can be used to recover beyond-DOF objects.

The team used the Argonne Leadership Computing Facility to develop parallelized code to address beyond-DOF imaging. We have shown that the Fresnel multislice method and the finite-difference approach implemented using PETSc both show high accuracy and excellent scaling for large problems in the forward model.  We also used TensorFlow, AutoGrad, and PyTorch to implement the AD-based solution of the object (the image reconstruction problem) for ptychography, as well as for other imaging methods including point-projection near-field holography.  We have developed a software framework Adorym with the flexibility to use our approach for a variety of imaging problems (the framework is described in this paper).  Finally, we are now exploring the use of the alternating direction of multipliers method (ADMM) to improve our ability to tackle larger problems and include neural network priors for improved performance.