This algorithm package consists of four algorithms for phase retrieval:
Phase retrieval is a well-known and important inverse problem in optics that seeks to recover the phase of light according to the intensity profile of the light, i.e., image measured by camera, after the light propagates through an optical system. Phase retrieval is a notoriously ill-posed problem that does not have a generic analytical solution. Therefore, we developed two physics and data driven hybrid methods for high-accuracy phase retrieval. GS direct unrolling is a state-of-the-art method with high accuracy but very slow inference speed when used for the optical system with a complex forward path. Therefore, FourierGSNet was developed from GS direct unrolling to enhance the inference speed on complex optical systems, e.g., the laser beam shaping system of M290 PBF-LB/M machine, while maintaining its high accuracy. We also developed two physics-driven fitting methods for the beam shaping system, which are even slower than GS direct unrolling but do not need training to achieve good accuracy.
The algorithms can be used by cloning the GitHub repository and installing the required dependencies.
FourierGSNet: FourierGSNet is an efficient Gerchberg-Saxton (GS) algorithm deep unrolling method. Unlike GS direct unrolling, instead of unrolling the original optical system, we unroll a simplified system based on Fourier transform. We extract physical knowledge from the simplified system and inject it into a neural network that recovers phase for the real system. This maintains the accuracy brought by the physical knowledge while significantly reducing computation. FourierGSNet addresses the inference speed limitation on complex optical systems and the lack of a system-agnostic solution of GS algorithm unrolling technique.
GS direct unrolling: GS direct unrolling is a state-of-the-art high-accuracy phase retrieval technique. It combines the physics-driven GS algorithm with data-driven deep learning method by unrolling a few GS algorithm iterations into cascaded layers and inserting deep learning layers after each unrolled iteration. The physics knowledge embedded in the unrolled GS iterations brings high phase recovery accuracy while the entire structure is still a neural network that is much faster than an iterative method which typically needs hundreds of iterations for predicting one phase map. When used for different optical systems, the forward and the inverse simulations in each unrolled GS iteration need to be implemented specifically for the system. When the simulation for the system is computationally intensive, inference speed of the entire method becomes slow. Since the simulation in the unrolled GS iterations must be both invertible and differentiable to support both the GS algorithm and deep learning, we directly employed the tool "Invertible differentiable batched GPU-accelerated simulation for the laser beam shaping system of EOS M290 PBF-LB-M system" from the AIMS 5.0 AI Toolbox for the beam shaping system. For the X-ray imaging system and the CDI system, we developed invertible differentiable simulations for them specifically.
Gradient descent method: This is a pure physics-driven fitting method of phase retrieval for the beam shaping system. The forward simulation involved in fitting loops is the same simulation in GS direct unrolling. The method parameterizes the phase aberration in the Spatial Light Modulator (SLM) with 12 Zernike coefficients. Mode decomposition with 12 Zernike polynomials as basis is used to decompose any aberration in SLM as the linear combination of the basis. The magnitude coefficient scaling a Zernike polynomial in the linear combination is the corresponding Zernike coefficient. The Zernike polynomials and the linear combination are implemented with Pytorch to be an extra neural network layer inside the simulation to keep the differentiability. The Zernike coefficients are set as learnable parameters and an MSE (Mean Square Error) loss function is constructed between the simulated beam shape and the target distorted beam shape. Finally, the Adam optimizer of Pytorch is used to perform gradient descent on Zernike coefficients for the minimization of the loss function to update the Zernike coefficients in order for the consistency between the simulated beam shape and the target. When the simulated beam shape is close enough to the target, the linear combination of the Zernike polynomials scaled by fitted Zernike coefficients is the recovered SLM phase aberration.
Newton-Raphson method: This method is the same as gradient descent method except the loss function and the optimizer. The loss function employed for this method is the negative log likelihood function of a Poisson distribution, as in this method we employed an image formation model that models the arrival of photons on each pixel as a queueing problem which typically follows a Poisson distribution. Minimizing the loss function can maximize the conditional probability of observing the target beam shape given the simulated beam shape. Therefore, this method is significantly different from gradient descent method because this method is not a simple parameter fitting but a Maximum Likelihood Estimation (MLE). To handle this more complex physics-informed loss function, gradient descent optimizers are not enough to make it converge. Therefore, Newton-Raphson optimization is used to minimize the loss function, which not only utilize the first order derivatives but also the second-order derivatives. All other parts are the same as the gradient descent method.
The algorithms operate in the following stages:
These algorithms are intended for:
Within the AIMS5.0 context, these algorithms are especially relevant for scenarios where fast and precise beam shaping, aberration detection, and correction are needed, and where fast and accurate X-ray imaging and coherent diffractive microscopy imaging are needed. FourierGSNet can provide rapid imaging and precise light control for industries such as biomedical industry, microscopy for laboratories and semiconductor industry, laser industry, and additive manufacturing industry.
If you use these four algorithms in research or development, please cite the corresponding paper.
@article{yan2026efficient,
title={Efficient Gerchberg--Saxton algorithm deep unrolling for phase retrieval with a complex forward path},
author={Yan, Shengyuan and Holenderski, Mike and Meratnia, Nirvana},
journal={Advanced Photonics Nexus},
volume={5},
number={2},
pages={026005--026005},
year={2026},
publisher={Society of Photo-Optical Instrumentation Engineers (SPIE)}
}