Feature-domain phase retrieval: ascension from image domains to feature domains

This is the MATLAB code for the implementation of feature-domain phase retrieval (FD-PR), an optimization framework for intensity-based wavefront recovery.

Feature-domain phase retrieval is a wavefront retrieval engine that recovers for a broad class of wavefront through unique non-convex, high-dimensional, feature-domain optimization with arbitrary constrains.

Optimizations of the non-convex loss function are regarded as supervised learning, solved by complex backpropagation.

News

• 2025/03/31: ✨ Our paper has been selected to be featured on the Inside Front Cover
  
• 2025/01/30: ✨ Our paper has been accepted by Advanced Science!
  
• 2024/05/05: 🔥 We released our MATLAB codes!
  

Contents

This repository contains the implementation of FD-PR for two wavefront tasks which are
(1) Feature-domain Fourier Ptychography
(2) Coded Ptychography
(3) Computational Holography

• Feature-domain phase retrieval: ascension from image domains to feature domains
  • Contents
  • How does it work?
    • Feature-domain likelihood
    • Extended Hybrid input-output (eHIO) modulus for Plug-and-Play constraints
    • Learning the wavefronts using Optimizers
  • Results
    • 1. Phase retrieval under unknown aberrations
    • 2. Coded ptychography
    • 3. Recovery with arbitrary constraints

How does it work?

The FD-PR begins with a general task for wavefront recovery, where one or a series of intensity measurements (observation, ob), $\mathbf{I}_1^{obs}, \mathbf{I}_2^{obs}, \dots \mathbf{I}_n^{obs}, \dots, \mathbf{I}_N^{obs}, (n = 1, 2, 3, \dots)$ were collected, with the corresponding image formation model (forward model)

$$\mathbf{I}_n = \text{Degrading}\left( \sum_{m=1}^{M} \left| \mathbf{A}_{n,m}\mathbf{x} \right|^2 \right)$$

describing the image formation progress of the optical system. $\text{Degrading}( \cdot)$ denotes arbitrary corrupt processes. $\mathbf{x}$ is the target wavefront to be recovered. The matrix $\mathbf{A}_{n,m}$ denotes a known linear process that converts the complex amplitude $\mathbf{x}$ to the $n$-th intensity measurement.

It is assumed that the final measured intensity is the summation of several intensity (a total of $M$) of incoherent waves, where $M$ can be a function of $n$.

The flowchart of FD-PR is depicted in the title figure, where the loss function for wavefront recovery comprises two blocks:
(1) The first block is the feature-domain augmented likelihood block that uniquely maximizes the data likelihood in image's feature-domain.
(2) The second block is the constraint block which implements extended-HIO (eHIO), providing plug-and-play interfaces for arbitrary customized constraints.

Feature-domain likelihood

The feature-domain likelihood is the core of FD-PR, which is established on image's feature extracted by invertible feature-extracting operators. The idea is that the image's feature is the inherent properties of image which is more robust to image degrading than image itself.With the feature-domain information, the likelihood function can better utilize the data, improving the robustness of recovery algorithm.

$${\Large \mathcal{L}_{Likelihood} = \mathcal{D} \left [ \mathbf{\Theta} \mathcal{S} \left(\mathbf{I}_n^{obs}\right),\ \mathbf{\Theta}\mathcal{S}\left( \mathbf{I}_n^{pre} \right) \right ]}$$

where $\mathbf{I}_{n}$ is the model predicted intensity. $\mathcal{D}(\mathbf{x},\mathbf{y}), \mathcal{D} \ge 0$ denotes an arbitrary differentiable likelihood or fidelity function measuring the \textit{distance} between the model prediction $\mathbf{x}$ and observation $\mathbf{y}$. $\mathcal{S}(\mathbf{x})$ is a scaling operation that adjusts the dynamics range of prediction and observations. $\mathbf{\Theta}$ is an manually-selected invertible feature extraction operator.

Extended Hybrid input-output (eHIO) modulus for Plug-and-Play constraints

The FD-PR framework incorporates the extended HIO, serving as a generalized Gerchberg-Saxton (GS) algorithm. The GS alternates between object and Fourier domain constraints to minimize the error between prediction and observation. If the object-domain constraints was treated as the likelihood-optimization, it would be natural to treat the Fourier-domain constraints as the prior/penalty-optimization.

In FD-PR, we optimize the likelihood using complex back-propagation, treating $\mathbf{x}^{t}$ as the input and $\mathbf{x}^{t+1}$ as the output after each gradient step. By reintroducing HIO, we insert custom constraints on $\mathbf{x}$ during gradient descent, leading to a refined penalty function that enhances the reconstruction quality.

Ptychography reconstruction.
Top left: FD-PR + TV; Top right: FD-PR + Second-order TV; Bottom left: FD-PR + Median filter; Bottom right: WASP.

The animation shows FD-PR on Ptychography with different denoisers including TV-denoiser, Second-order TV denoiser, and Median filter. The results are compared with WASP.

Learning the wavefronts using Optimizers

The FD-PR bears resemblance to training a deep neural network in a supervised manner, in which the target wavefront is learned from a series of intensity observations by minimizing the loss function through complex back-propagation.

• For information of optimizers please refer Optimizing gradient descent.
  
• For Python implementation of optimizers please refer Optimizers.
  
• Usually, the optimizers are designed for real-valued variables and cannot be directly applied to complex-variable in our case, a little modifications to the codes of the optimizers are needed, please refer this discussion.

Results

1. Phase retrieval under unknown aberrations

The following GIF shows how FD-PR implementation works for a Fourier ptychography experiment, which retrieves the phase pattern of a quantitative phase target together with the aberration of the pupil function.
Sample codes and data are available in FD-PR-FPM

2. Coded ptychography

The FD-PR can be applied to Coded ptychography as well. Data can be found in DATA. For FD-PR, the user needs first save all image file to .tif format in a foler named 'raw_data', then you can select a specific area for reconstruction.

Please refer this PAPER for the implementation of codes and data.

3. Recovery with arbitrary constraints

The FD-PR can achieve twin-image free inline holography by using different denoising constraints.

License and Citation

This framework is licensed under the MIT License. Please see LICENSE for details.

If you use it in your research, we would appreciate a citation via

@article{zhang2025high,
  author = {Zhang, S. and Pan, A. and Sun, H. and Tan, Y. and Cao, L.},
  title = {High-Fidelity Computational Microscopy via Feature-Domain Phase Retrieval},
  journal = {Advanced Science},
  year = {2025},
  volume = {12},
  pages = {2413975},
  doi = {10.1002/advs.202413975},
  url = {https://doi.org/10.1002/advs.202413975}
}