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.