2 Wavefront sensing

The wavefront sensor that was simulated in this example was a quadriwave lateral shearing interferometer (QWLSI), which is known for its high resolution and reconstruction fidelity. Nonetheless, our method can in general be applied to any kind of wavefront retrieval technique, including the popular Shack–Hartmann sensor or multi-plane techniques, such as Gerchberg–Saxton phase retrieval.

A lateral shearing interferometer (LSI) measures the spatially varying phase of a light beam, and was first applied to the measurement of ultrashort laser pulses in the late 1990s^[ Reference Chanteloup, Druon, Nantel, Maksimchuk and Mourou ^{14 ]}. The LSI works by creating multiple copies of the laser pulse and shearing them laterally relative to each other before their interference pattern is captured on a sensor. Due to the shear, information about the spatial gradient of the wavefront is encoded in the interferogram. This can then be extracted using Fourier filtering^[ Reference Takeda, Ina and Kobayashi ^{15 ]} and stitched together to form the wavefront via methods such as modal reconstruction^[ Reference Dai, Tang, Wang, Sasaki and Feng ^{16 ]} or Fourier integration^[ Reference Velghe, Primot, Guérineau, Cohen and Wattellier ^{17 ]}. The most popular implementation is the aforementioned QWLSI. By generating and shearing four ( $2\times 2$ ) copies of the pulse under investigation, this setup also enables the extraction of two separate pairs of orthogonal gradients, meaning two distinct estimates for the wavefront can be found, providing error estimation. This property is highly desirable for a sensor based on CS, because inevitable noise in the measurement can corrupt the wavefront with reconstruction artifacts. This would, for instance, be the case in a two-plane Gerchberg–Saxton algorithm. In contrast, redundancy of phase information in the QWLSI provides direct validation and, thus, makes the wavefront retrieval much more resilient to noise. A sketch illustrating the concept of the QWLSI is shown in the red box of Figure 1.

Figure 1 Schematic of the experimental setup that was simulated. The pulse first travels through a quadriwave lateral shearing interferometer, yielding a hypercube of interferograms, a slice of which is shown in the green box. The hypercube is then passed through a CASSI setup. This consists of a random mask and a relay system encompassing a prism, before the coded shot is captured with the camera. This diagram is not to scale.

2.1 QWLSI simulation

A physical implementation of the QWLSI usually consists of a phase grating with ‘pixels’ of alternating phase arranged in a checkerboard pattern^[ Reference Velghe, Primot, Guérineau, Haïdar, Demoustier, Cohen and Wattellier ^{18 ]}, which leads to dominant diffraction in $2\times 2$ copies of the beam. Instead of simulating this process, we consider an idealized setup where we analytically generate the four copies. We begin with creating the pulse of interest by defining a spatial-spectral intensity, ${I}_0$ , and phase, ${\phi}_0$ :

(1) $$\begin{align}{E}_0\left(x,y,z=0,\omega \right)=\sqrt{I_0\left(x,y,z=0,\omega \right)}{e}^{i{\phi}_0\left(x,y,z=0,\omega \right)}.\end{align}$$

The pulse is copied four times, and each copy’s field is propagated to the detector plane according to the following rules. Note that for brevity, when the $z$ index is not stated, $z=0$ and the $\omega$ index will be suppressed from the electric field.

2.1.1 Propagation

Considering the jth copy has travelled a distance $\Delta z$ , its poloidal angle is ${\theta}_{{j}}$ and its azimuthal angle is $\zeta$ , then one finds its displacement in the $x$ and $y$ directions to be as follows:

$$\begin{align*}\Delta {x}_{{j}}\left(\omega \right)=\Delta z\sin \left({\theta}_{{j}}\right)\sin \left(\zeta \left(\omega \right)\right),\\\Delta {y}_{{j}}\left(\omega \right)=\Delta z\cos \left({\theta}_{{j}}\right)\sin \left(\zeta \left(\omega \right)\right).\end{align*}$$

In a QWLSI, the poloidal angles are ${\theta}_{{j}}\in \left\{0,\frac{\pi }{2},\pi, \frac{3\pi }{2}\right\}$ . In Figure 1, one identifies the azimuthal angle, $\zeta$ , as that between the white copy lines and the central yellow line. This is related to both the pitch of the grating $\Lambda$ and the wavelength $\lambda$ by the following:

$$\begin{align*}\zeta \left(\omega \right)=\arcsin \left(2\pi \frac{\lambda }{\Lambda}\right).\end{align*}$$

The resulting electric field of the copy is as follows:

(2) $$\begin{align}{E}_{{j}}\left(x,y,\Delta z\right)=\frac{1}{4}\sqrt{I_0\left(x-\Delta {x}_{{j}},y-\Delta {y}_{{j}}\right)}{e}^{i{\phi}_0\left(x-\Delta {x}_{{j}},y-\Delta {y}_{{j}}\right)}.\end{align}$$

2.1.2 Tilt

As diffraction occurs at an angle $\zeta$ , the grating imparts a tilt onto the copy. This translates to an additional phase shift dependent on both the spectral and spatial domains:

$$\begin{align*}\Delta {\phi}_{{j}}\left(x,y,\omega \right)=k\left(\omega \right)\left(x\cos \left({\theta}_{{j}}\right)+y\sin \left({\theta}_{{j}}\right)\right)\sin \left(\zeta \right),\end{align*}$$

where $k=\frac{2\pi }{\lambda }$ is the wavevector of the pulse. This tilt is crucial in reconstruction as it provides a high-frequency modulation that separates the gradients in Fourier space.

Combining these two effects and summing over copies, we obtain the final changes to the field:

(3) $$\begin{align}\begin{array}{l} E\left(x,y,\Delta z\right)=\frac{1}{4}\sum \limits_{{{j}}=1}^4\;\sqrt{I_0\left(x-\Delta {x}_{{j}},y-\Delta {y}_{{j}}\right)}\\ {}\kern7em \cdot {e}^{i(\phi (x-\Delta {x}_{{j}},y-\Delta {y}_{{j}})+\Delta {\phi}_{{j}}(x,y))}. \end{array}\end{align}$$

At the Talbot self-imaging plane, $\Delta z=2{\Lambda}^2/\lambda$ , one has a hypercube of interferograms. An example of a one frequency channel slice is shown in the green box of Figure 1.

In other applications one would collapse the cube onto a sensor at this point; however, this would eliminate the chance of retrieving the spectrally resolved phase. Instead, as discussed in Section 3, we use SCI to aid in the capturing of the cube.

3 Snapshot compressive imaging

CS describes the highly efficient acquisition of a sparse signal from fewer samples than would classically be required according to the Nyquist theorem by utilizing optimization methods to solve underdetermined equations. SCI is an example of CS, capturing 3D data on a 2D sensor in a single shot.

Fundamentally, there are two requirements to be fulfilled for CS to work. Firstly, the signal must be sparse on some basis and, secondly, the signal must be sampled on the basis that it is incoherent with respect to the sparse basis^[ Reference Candes ^{21 ]}. The first condition was hypothesized to be satisfied given the fact that laser wavefronts are known to be well-expressed with a few coefficients of the Zernike basis. When one does not have prior knowledge about which basis the signal is sparse on, the second condition is often solved by performing random sampling. Whilst being trivial for 2D data, in the context of SCI it is challenging, as the 3D hypercube must be randomly sampled onto a 2D sensor. To do so, nearly all research in this area uses hardware based on the coded aperture snapshot compressive imaging (CASSI) system^[ Reference Gehm, John, Brady, Willett and Schulz ^{22 ,} Reference Wagadarikar, John, Willett and Brady ^{23 ]}.

3.1 CASSI

The hypercube is first imaged onto a coded aperture. This is a binary random mask with each pixel transmitting either 100% or 0% of the light. The cube is then also passed through some dispersive media, for example, a prism or grating, before being captured by a sensor, resulting in what is known as the coded shot. The effect of this optical system is that when the hypercube reaches the detector plane, each spectral channel is encoded with a different coded aperture, thereby approximating random sampling across the whole cube. It is then possible for a reconstruction algorithm to retrieve the cube. A diagram of a CASSI system is shown in the yellow box in Figure 1, with an example of a coded shot for an interferogram hypercube shown on the far left of Figure 3. The setup can easily be simulated by multiplying the cube by the mask, then shifting the channels of the cube according to the amount of (angular) dispersion imparted onto them, and finally summing over the spectral axis.

Mathematically, the CASSI system discussed above is summarized into a matrix $\boldsymbol{\varPhi}$ , which operates on $\boldsymbol{m}$ , a vectorized representation of the hypercube, to give $\boldsymbol{n}$ , a vectorized coded shot:

(4) $$\begin{align}\boldsymbol{n}=\boldsymbol{\varPhi} \boldsymbol{m}.\end{align}$$

In order to reconstruct $\boldsymbol{m}$ , one can solve the following:

(5) $$\begin{align}\tilde{\boldsymbol{m}}={\mathrm{argmin}}_{\boldsymbol{m}}\left[\underset{\kern0.1em \mathrm{data}\kern0.17em \mathrm{term}\kern0.1em }{\underbrace{{\left\Vert \boldsymbol{n}-\boldsymbol{\varPhi} \boldsymbol{m}\right\Vert}^2}}+\underset{\kern0.1em \mathrm{regularizer}\kern0.1em }{\underbrace{\eta \mathrm{\mathcal{R}}\left(\boldsymbol{m},\psi \right)}}\right].\end{align}$$

The first term on the right-hand side is labelled the data term, and enforces that the hypercube must match the coded shot when captured. This alone would be an underdetermined system, so a regularization term, parameterized by $\psi$ , is added that restricts the solution space and selects the correct hypercube.

Most methods that have been developed to solve this non-convex equation can be sorted into two classes: iterative algorithms or end-to-end neural networks. The former offers good generalization but lacks abstraction capability and is slow, whilst deep nets are fast and have been shown to learn almost any function, but can be prone to overfitting^[ Reference Yuan, Brady and Katsaggelos ^{24 ]}. A middle ground that offers state-of-the-art performance is deep unrolling.

3.2 Deep unrolling

While an end-to-end neural net would attempt to solve Equation (5) directly, if it were possible to split the equation, the data term can actually be solved analytically. This is desirable as it alleviates the abstraction needed to be done by the network, resulting in greater generalization and better parameter efficiency^[ Reference Monga, Li and Eldar ^{25 ]}. To perform such a separation, half quadratic splitting is employed. Firstly, an auxiliary variable $\boldsymbol{p}$ is substituted into the regularization term, with Equation (6) being equivalent to Equation (5). Then, the constraint is relaxed and replaced by a quadratic loss term:

(6) $$\begin{align}\hspace{-5pt}\widehat{\boldsymbol{m}},\kern1pt\widehat{\boldsymbol{p}}={\mathrm{argmin}}_{{\boldsymbol{m}},{\boldsymbol{p}}}\left[{\left|\boldsymbol{n}-\boldsymbol{\varPhi} \boldsymbol{m}\right|}^2+\eta R\left(\boldsymbol{p}\right)\right]\kern0.24em \mathrm{s}.\mathrm{t}.\ \boldsymbol{m}=\boldsymbol{p},\end{align}$$

(7) $$\begin{align}\hspace{20pt}\approx{\mathrm{argmin}}_{{\boldsymbol{m}},{\boldsymbol{p}}}\left[|\boldsymbol{n} - \boldsymbol{\varPhi}\boldsymbol{m}|^{2} + \eta R(\boldsymbol{p}) + \beta |\boldsymbol{m} - \boldsymbol{p}|^{2}\right].\end{align}$$

Here, $\beta$ is a variable that controls the strength of the constraint. High values of $\beta$ will strongly enforce $\boldsymbol{m}=\boldsymbol{p}$ and approximate the subject-to statement.

The benefit of this problem formulation is that it is then possible to split Equation (7) into two minimization sub-problems in $\boldsymbol{m}$ and $\boldsymbol{p}$ , and effectively separate the data term from the regularization term. When minimized iteratively, the following sub-problems can approximate Equation (7):

(8) $$\begin{align}{\widehat{\boldsymbol{p}}}^{k+1}={\mathrm{argmin}}_{\boldsymbol{p}}\left[\beta {\left|\boldsymbol{p}-{\boldsymbol{m}}^k\right|}^2+\eta R\left(\boldsymbol{p}\right)\right]\sim \mathcal{S}\left({\boldsymbol{m}}^k\right),\end{align}$$

(9) $$\begin{align}{\widehat{\boldsymbol{m}}}^{k+1}={\mathrm{argmin}}_{\boldsymbol{m}}\left[{\left|\boldsymbol{n}-\boldsymbol{\varPhi} \boldsymbol{m}\right|}^2+\beta {\left|{\boldsymbol{p}}^{k+1}-\boldsymbol{m}\right|}^2\right].\end{align}$$

Equation (9) is a convex equation and can be solved via a conjugate gradient algorithm, which provides better stability than solving analytically. On the right-hand side of Equation (8), $\mathcal{S}$ represents that a neural network will be used to solve the equation.

The deep unrolling process is shown in Figure 2(b) panel (i). Firstly ${\boldsymbol{m}}^{(0)}$ is initialized: ${\boldsymbol{m}}^{(0)}={\boldsymbol{\varPhi}}^\mathrm{T}\boldsymbol{n}$ . Then the two equations are solved for a fixed number of iterations, with the same architecture neural net being used to represent Equation (8) in each iteration. However, the network has its own set of weights for each iteration, hence the unrolling of the algorithm. The architecture of the network will be discussed in Section 4.

Figure 2 A diagram showing the full reconstruction process of the wavefront from the coded shot. (a) A flow chart of the reconstruction process. (b) (i) The deep unrolling process, where sub-problems ① and ② are solved recursively for 10 iterations. Also shown is the neural network structure used to represent $\mathcal{S}\left({\boldsymbol{m}}^k\right)$ . (ii) The training curve for the deep unrolling algorithm. Plotted is the training and validation PSNR for the 3D ResUNet prior that was used, as well as the validation score for a local–nonlocal prior. Here is demonstrated the superior power of 3D convolutions in this setting. (i) The network design for the Xception-LSI network. The Xception* block represents that the last two layers were stripped from the conventional Xception network. (c) (ii) The training curve for Xception-LSI for training and validation sets, with the loss shown in log mean squared error. Also plotted is the validation loss when further training the model on the deep unrolling reconstruction of the data (transfer).

4 Network architecture

This section contains the architectures of the neural networks that were used. They will be discussed in the order they are used in the reconstruction process, which is displayed in the flow chart of Figure 2(a). Firstly, the deep unrolling algorithm performs reconstruction of the interferogram hypercube from the coded shot, and secondly another network, Xception-LSI, reconstructs the spatial-spectral wavefront from the hypercube.

5 Training data generation

To represent the initial pulse, a total of 300 cubes were generated with dimensions $\left({n}_{{x}}\times {n}_{{y}}\times {n}_{\omega}\right)=\left(512\times 512\times 31\right)$ . The data was randomly split at a ratio of $4:1:1$ into training, validation and test sets, respectively. The wavelength range considered was 750–850 nm, representing a broadband Ti:sapphire laser, giving $\Delta \lambda \approx 3.23$ nm. For each cube, the wavefront for each channel was first initialized to a randomly weighted sum of 15 Zernike basis functions. Then, to simulate an STC, one Zernike function was chosen and was made to vary either linearly or quadratically with frequency. Indeed, common STCs, such as pulse front tilt and pulse front curvature, can be represented in this way^[ Reference Jolly, Gobert and Quere ^{1 ]}. The mean amplitude of this coefficient was also made to be higher. This choice of Zernike coefficients is arbitrary, but allows for a demonstration that the method can identify all Zernike basis functions. The intensity of slices of the cube was set to an image taken of a real laser intensity.

Each cube was then processed according to Figure 1. Firstly, it was passed through the QWLSI simulation (see Section 2.1), yielding a hypercube of interferograms – these are the training labels for the deep unrolling algorithm. This hypercube was then passed through the SCI simulation, yielding a coded shot – the training data. The wavefront was reconstructed via the process in Figure 2(a). The interferogram hypercube was reconstructed via deep unrolling, before being passed into the Xception-LSI network to predict the spectral Zernike coefficients.

The pitch of the LSI was set to $\Lambda =80$ μm, and the dispersion of the prism, measured at the camera plane, was set to 1 pixel per channel (each channel having a width of 3.23 nm).

Before being passed through the deep unrolling network, the cubes and coded shots were split spatially into 64 $\times$ 64 oblique parallelepiped patches, allowing for a one-to-one reconstruction between the input and output^[ Reference Wang, Sun, Fu, Kim and Huang ^{28 ]}. The initial learning rate was set to 0.01 and decayed by 10% every five epochs. The total number of epochs was 70, and the batch size was 8.

The Xception-LSI network was fed individual channels of the ground truth interferogram hypercubes and predicted Zernike coefficients. Normal random noise ( $\mathcal{N}(\mu =0,\sigma =0.1$ )) was applied to the input, to make the model robust to noise produced by the SCI reconstruction. The initial learning rate was set to ${10}^{-5}$ and decayed by 10% every five epochs. The total number of epochs was 40, and the batch size was 16. Once trained on the ground truth hypercubes, the model was trained on interferogram hypercubes that had been reconstructed by deep unrolling, for the further eight epochs. The aim of this transfer learning was to allow the network to account for any systematic noise in the SCI reconstruction, resulting in a more accurate wavefront reconstruction.

7 Summary and outlook

In this report we have demonstrated the possibility of combining a wavefront sensor with SCI in order to achieve a single-shot measurement of the spatial-spectral phase. Crucially, it has been shown that SCI has the ability to reconstruct modulated signals, such as those produced by a QWLSI.

A natural progression to this study is to realize the results in an experimental setting, where challenges arise from the more complicated dispersion, transfer functions and noise. Other further work could include extending the deep learning LSI analysis to the hyperspectral setting. By passing the network a hypercube of interferograms, rather than individual slices, it may be possible to exploit spectral correlations in order to improve accuracy and detect STCs more easily. Also, work can be done on testing the model with a more varied set of Zernike polynomials. Finally, there has been recent interest in the possibility of spreading phase contrast imaging to a hyperspectral setting. However, current methods take many seconds to capture a hypercube of phase^[ Reference Ba, Tsang and Mertz ^{32 ]}. The proposed method would be able to collect information with higher spectral resolution in a single shot, allowing for dynamic events to be recorded hyperspectrally.