A semi-implicit relaxed Douglas-Rachford algorithm (sir-DR) for Ptychograhpy

TL;DR

The paper introduces a semi-implicit relaxed Douglas-Rachford algorithm (sir-DR) that accelerates convergence and improves reconstruction quality in ptychography, especially with sparse data, outperforming existing methods like ePIE and rPIE.

Contribution

The paper presents a novel semi-implicit relaxed Douglas-Rachford algorithm for ptychography, demonstrating improved convergence and reconstruction with sparse data, and provides open-source implementation.

Findings

sir-DR converges faster than ePIE and rPIE.

sir-DR achieves higher reconstruction quality.

sir-DR successfully reconstructs images with sparse measurements.

Abstract

Alternating projection based methods, such as ePIE and rPIE, have been used widely in ptychography. However, they only work well if there are adequate measurements (diffraction patterns); in the case of sparse data (i.e. fewer measurements) alternating projection underperforms and might not even converge. In this paper, we propose semi-implicit relaxed Douglas Rachford (sir-DR), an accelerated iterative method, to solve the classical ptychography problem. Using both simulated and experimental data, we show that sir-DR improves the convergence speed and the reconstruction quality relative to ePIE and rPIE. Furthermore, in certain cases when sparsity is high, sir-DR converges while ePIE and rPIE fail. To facilitate others to use the algorithm, we post the Matlab source code of sir-DR on a public website (www.physics.ucla.edu/research/imaging/sir-DR). We anticipate that this algorithm can be generally applied to the ptychographic reconstruction of a wide range of samples in the physical and biological sciences.

Equations

$$|\mathcal{F}(P{O}_n)|=\sqrt{I_n}\text{mbox}forn=1,..,N.$$

$$O(x+r_n)=O_n(x)\text{mbox}if{r}_n+x\in {\Omega }_n\text{mbox}forn=1,...,N$$

$$\mathcal{T}$$

$$\mathcal{S}$$

$$\underset{\Psi }{\min }{i}_{\mathcal{S}}(\Psi )+i_{\mathcal{T}}(\Psi )$$

$$i_{\mathcal{S}}(\Psi )=\{\begin{array}{ll}0& \Psi \in \mathcal{S}\\ \infty & \text{otherwise}\end{array}$$

$${\Psi }_n^{\prime }$$

$$\{P^{k+1},O_n^{k+1}\}$$

$${\Psi }_n^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$\{P^{k+1},O_n^{k+1}\}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$\text{mbox}Step1:O_n^{k+1}$$

$$\text{mbox}Step2:P^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$O_n^{k+1}$$

$$P^{k+1}$$

$$\displaystyle O_{n}^{k+1}=O^{k}+\beta_{O}\frac{\overline{P^{k}}\big{(}\Psi_{n}^{\prime}-\Psi_{n}^{k}\big{)}}{(1-\beta_{O})\|P^{k}\|_{\max}^{2}+\beta_{O}|P^{k}|^{2}}$$

$$\underset{\Psi }{\min }f(\Psi )+g(\Psi )$$

$$\displaystyle\Psi^{k+1}=\Psi^{k}+\mathrm{prox}_{tf}\big{(}2\;\mathrm{prox}_{tg}(\Psi^{k})-\Psi^{k}\big{)}-\mathrm{prox}_{tg}(\Psi^{k})$$

$$\displaystyle\Psi^{k+1}=\Psi^{k}+\Pi_{\mathcal{T}}\big{(}2\;\Pi_{\mathcal{S}}(\Psi^{k})-\Psi^{k}\big{)}-\Pi_{\mathcal{S}}(\Psi^{k})$$

$$\displaystyle\Psi^{k+1}=\mathrm{prox}_{tf}\Big{(}(1+\sigma)\;\mathrm{prox}_{tg}(\Psi^{k})-\sigma\Psi^{k}\Big{)}+\sigma\Big{(}\Psi^{k}-\mathrm{prox}_{tg}(\Psi^{k})\Big{)}$$

$$\underset{{\Psi }_n}{\min }\sum _{n=1}^N\Vert |\mathcal{F}{\Psi }_n|-\sqrt{I_n}{\Vert }^2+i_{\mathcal{S}}({\Psi }_n)$$

$${\mathrm{prox}}_{tf}({\Psi }^k)$$

$$\displaystyle=\frac{\Psi^{k}+t\mathcal{F}^{-1}\Big{[}\sqrt{I_{n}}\arg(\mathcal{F}\Psi)\Big{]}}{1+t}$$

$$\displaystyle=(1-\tau)\Psi^{k}+\tau\mathcal{F}^{-1}\Big{[}\sqrt{I_{n}}\arg(\mathcal{F}\Psi)\Big{]}$$

$${\Psi }^{k+1}$$

$$\displaystyle=\tau\Big{(}\sigma\Psi^{k}+\Pi_{\mathcal{T}}\big{(}(1+\sigma)\pi_{\mathcal{S}}(\Psi^{k})-\sigma\Psi^{k}\big{)}\Big{)}+\Big{(}1-\tau(1+\sigma)\Big{)}\Pi_{\mathcal{S}}(\Psi^{k})$$

Full text

††journal: osac††articletype: Research Article

1 Introduction

Since the first experimental demonstration in 1999 [1], coherent diffraction imaging (CDI) through directly inverting far-field diffraction patterns to high-resolution images has been a rapidly growing field due to its broad potential applications in the physical and biological sciences [2, 3, 4, 5]. A fundamental problem of CDI is the phase problem, that is, the diffraction pattern measured only contains the magnitude, but the phase information is lost. In the original demonstration of CDI, phase retrieval was performed by measuring the diffraction pattern from a finite object. If the diffraction intensity is sufficiently oversampled [6], the phase information can be directly retrieved by using iterative algorithms [7, 8, 9, 10, 11]. Ptychography, a powerful scanning CDI method, relieves the finite object requirement by performing 2D scanning of an extended relative to an illumination probe and measuring multiple diffraction patterns with each illumination probe overlapped with its neighboring ones [12, 13]. The overlap of the illumination probes not only reduces the oversampling requirement, but also improves the convergence speed of the iterative process. By taking advantage of ever-improving computational power and advanced detectors, ptychography has been applied to study a wide range of samples using both coherent x-rays and electrons [2, 5, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]. More recently, a time-domain ptychography method was developed by introducing a time-invariant overlapping region as a constraint, allowing the reconstruction of a time series of complex exit wave of dynamic processes with robust and fast convergence [24].

Algorithms for ptychography have been studied exhaustively in theory and practice. The majority following non-convex optimization approaches [25, 26, 27], while a few follow convex relaxation [28]. In recent years, powerful ptychographic algorithms have been developed to handle partial coherence [29], solve for the probe [30, 13, 31], correct positioning errors [32, 33, 34], reduce noise [35, 36], and deal with multiple scattering [37, 38].

These iterative algorithm can be generally divided into three classes: i) the conjugate gradient (CG) [30, 34], ii) extended ptychography iterative engine (ePIE) [31], and iii) difference map (DM) [13], whereas the last two have a close relationship. ePIE is an alternating projection algorithm, while DM is built on both projection and reflection which is believed to provide a momentum to speed up the convergence. The relaxed average alternating reflection (RAAR) [8] is a relaxation of DM and has been shown to be effective in phase retrieval [39]. All algorithms except ePIE take a global approach, i.e. using the entire collection of diffraction patterns to perform an update of the probe and object in each iteration. In contrast, ePIE goes through the measured data sequentially to refine the probe and object. However, ePIE has a slow convergence due to the step size restriction which may cause divergence if violated. To fix this issue, rPIE was proposed, in which regularization is used for stability [40]. The significant results also show that rPIE obtains a larger field of view (FOV) than ePIE.

In this paper, we show that DM and RAAR can also be implemented locally, similarly to ePIE. We then apply non-convex optimization tools to improve the robustness, convergence and FOV of the ptychography reconstruction. The proposed algorithm incorporates two techniques. The first modifies the update of the probe and object as the algorithm iterates via semi-implicit method or Proximal Mapping. The second technique is the implementation of relaxed Douglas Rachford, a generalized version of DM and RAAR, on the local scale.

2 The proposed algorithm

Given N measured diffraction patterns at N positions, the ptychographic algorithm aims to find a 2D object O and probe P that satisfy the overlap constraint and the Fourier magnitude constraint

[TABLE]

Where $O_{n}$ is the object at the $n^{th}$ scan position. Here, we omit the spatial variables for a simple notation and use the notations $P,\,O_{n}$ and $I_{n}$ for both continuous and discrete cases. The absolute value, multiplication, division, conjugate, and square root operators are applied element-wise on $P,\,O_{n}$ and $I_{n}$ which represent 2D complex matrices of the same size in the discrete case. We can argue that $O_{n}$ is the object restricted to a sub-domain $\Omega_{n}$. The overlap constraint can be mathematically interpreted as

[TABLE]

where $\{r_{n}\}_{n=1}^{N}$ are displacement vectors. In short notation, we write $O_{n}=O|_{\Omega_{n}}$ to imply the object is restricted to sub domain $\Omega_{n}$. The equivalent constraint in the discrete case is the agreement between the sub-matrix of $O$ and $O_{n}$. We find a better representation of the problem by introducing the exit wave variable $\Psi=PO$. By denoting the Fourier measurement constraint set $T$ and the overlap object constraint set $S$, we have

[TABLE]

Then we write the ptychography problem in a minimization fashion

[TABLE]

where $i_{\mathcal{S}}(\Psi)$ and $i_{\mathcal{T}}(\Psi)$ are the indicator functions of sets $\mathcal{S}$ and $\mathcal{T}$ respectively, defined as

[TABLE]

To solve Eq. (3), an alternating projection method is proposed. At each iteration, we select a random position $n$ and update $\Psi_{n}$

[TABLE]

The Frobenius norm is used in this minimization problem and entire paper unless a different norm is specified. The minimization of Eq. (6) is difficult due to instability. One way to solve this non-convex problem is to minimize each variable while fixing the other ones.

[TABLE]

This approach is unstable because of the division. A cut-off method is used to avoid the divergence and zero-division. A modification is recommended by adding a penalizing least square error term (i.e. regularizer)

[TABLE]

The idea of regularization appears throughout the literature such as proximal algorithms [41, 42]. Eq. (9) is more reliable to solve than Eq. (6) but is still very expensive since the variables are coupled. $P^{k+1}$ and $O_{n}^{k+1}$ can be solved via a Backward-Euler system derived from Eq. (9).

[TABLE]

ePIE proposes a simple approximation by linearizing the system so that it can be solved sequentially.

[TABLE]

The system is solved by alternating direction methods (ADM) [43]. The remaining part is to choose appropriate step sizes $t$ and $s$ to ensure stability. ePIE suggests $t=\beta_{O}/\|P^{k}\|_{\max}^{2}$ and $s=\beta_{P}/\|O^{k+1}\|_{\max}^{2}$ where $\beta_{O},\;\beta_{P}\in(0,1]$ are normalized step sizes. The max matrix norm is the element-wise norm, taking the maximum in absolute values of all elements in the matrix. The final version of ePIE is given by

[TABLE]

We will exploit the structure of Eq. (10) to give a better approximation.

2.1 A semi-implicit algorithm

We replace the minimization of Eq. (9) by two steps

[TABLE]

This results in a better approximation to the linearized system of Eq. (11) and simpler than the Backward-Euler Eq. (10)

[TABLE]

In this uncoupled system, we can derive a closed form solution for each sub-problem.

[TABLE]

By choosing the step sizes $s$ and $t$ as in the ePIE algorithm and normalizing the parameters $\beta_{O}$ and $\beta_{P}$, we obtain

[TABLE]

This formula can be explained as a weighted average between the previous update $O_{n}^{k}$ and $\displaystyle\frac{\Psi_{n}^{k}}{P^{k}}$. The object update is similar to the rPIE algorithm when rewriting it as

[TABLE]

The difference is rPIE does not have the parameter $\beta_{O}$ in front of the fraction. i.e. rPIE has a larger step size than sir-DR. This helps converge faster but might also get trapped in local minima. The regularization (weighted average) in sir-DR is more mathematically correct and enhances the algorithm’s stability.

In the next section, we apply the Douglas-Rachford algorithm to solve for the exit wave $\Psi$.

2.2 The relaxed Douglas-Rachford algorithm

The Douglas-Rachford algorithm was originally proposed to solve the heat conduction problem [44], which represents a composite minimization problem

[TABLE]

The iteration consists of

[TABLE]

Over the past decades, this accelerated convex optimization algorithm has been exhaustively studied in both theory and practice with many applications [45, 46, 47, 48, 49, 50].Here we apply the algorithm to the ptychographic phase retrieval. Note that the Douglas-Rachford algorithm reduces to Difference Map (DM) when $f=i_{\mathcal{T}}$ and $g=i_{\mathcal{S}}$ are characteristic functions of constraint sets $\mathcal{T}$ and $\mathcal{S}$, respectively

[TABLE]

We realize that the reflection operator $2\Pi_{\mathcal{S}}-I$ helps to accelerate the convergence in the convex case and escape local minima in the non-convex case. However this momentum, caused by reflection, might be too large and can lead to over-fitting. Therefore, we relax the reflection by introducing the relaxation parameter $\sigma\in[0,1]$

[TABLE]

Since the experimental measurements are contaminated by noise, a direct projection of measurement constraint is not an appropriate approach. We thus relax the Fourier magnitude constraint by a least square penalty

[TABLE]

Recall that $\mathrm{prox}_{tf}(\Psi)$ has a closed form solution

[TABLE]

where $\tau=t/(1+t)\in(0,1)$ exclusively is the normalized step size. Combining this result with DM, we obtain

[TABLE]

When we let $\beta=1-\tau$ and $\sigma=1$, the update reduces to RAAR. Therefore, we show that relaxed Douglas-Rachford is a generalized version of RAAR. We now move to our main algorithm.

2.3 The sir-DR algorithm

In combination of the semi-implicit algorithm and relaxed Douglas Rachford algorithm, we propose the sir-DR algorithm, shown in Fig. 1.

In this algorithm, we only apply the semi-implicit method on $O_{n}^{k}$ while $P^{k}$ can be integrated with the Forward Euler (gradient descent) method. $\tau\in[0,1]$ is chosen to be small. In most of our experiments, we select $\tau\approx 0.1$, while the choice of $\sigma$ depends on the specific problem. In many cases, $\sigma=1$ works very well (full reflection). But in some specific cases, large $\sigma$ might cause divergence or small recovered FOV. We decrease $\sigma$ in these cases, for example $\sigma=0.5$. We choose $\beta_{O}=0.9$ in most cases. $\beta_{P}$ is chosen to be large at the beginning ($\beta_{P}=1$) and decreases as a function of iteration. This adaptive step-size has been introduced as a strategy for noise-robust Fourier ptychography [51].

3 Experimental results

3.1 Reconstruction from simulated data

To examine the sir-DR algorithm, we simulate a complex object of $128\times 128$ pixels with a cameraman and a pepper images representing the amplitude and the phase, respectively (Fig. 2).The circular aperture is chosen as probe with a radius of 50 pixels. We raster scan the aperture over the object with a step size of 35 pixels, resulting in 4x4 scan positions. The overlap is therefore 56.4%, the approximate lower limit for ePIE to work. Poisson noise is added to the diffraction patterns with a flux of $10^{8}$ photons per scan position. We use $R_{noise}$ to quantify the relative error with respect to the noise-free diffraction patterns

[TABLE]

where $P^{0}$ and $O^{0}$ is the noise-free model and the $L_{1,1}$ matrix norm represents the sum of all elements in absolute value of the matrix. The above flux results in $R_{noise}=3.73\%$. Fig. 3 shows that three algorithms (ePIE, rPIE and sir-DR) all successfully reconstruct the object in the case where the overlap between adjacent positions is high and the noise level is low.

As a baseline comparison, Fig. 3 shows that all three algorithms correctly reconstruct the object in the ideal case when the overlap between adjacent positions is high and the noise level is low.

Next, we apply the three algorithms to the reconstruction of sparse data, which is centrally important to reducing computation time, data storage requirements and incident dose to the sample. We increase the scan step size to 50 pixels while keeping the same field of view, which reduces the number of diffraction patterns to $3\times 3$. Consequently, the overlap is reduced to 39.1%. Not only is the overlap between adjacent positions low, but the total number of measurements is also small, creating a challenging data set for conventional ptychographic algorithms. Fig. 4 show that sir-DR can work well with sparse data, while ePIE and rPIE fail to reconstruct the object faithfully.

3.2 Reconstruction from experimental data

3.2.1 Optical laser data

As an initial test of sir-DR with experimental data, we collect diffraction patterns from an USAF resolution pattern using a green laser with a wavelength of 543 nm. The incident illumination is created by a $15\,\mu m$ diameter pinhole. The pinhole is placed approximately 6 mm in front of the sample, creating a illumination wavefront on the sample plane that can be approximated by Fresnel propagation. The detector is positioned 26 cm downstream of the sample to collect far-field diffraction patterns. We raster scan across the sample with a step size of $50\,\mu m$ and acquire 169 diffraction patterns. We perform a sparsity test by randomly choosing 85 diffraction patterns (50% density) and run ePIE, rPIE, and sir-DR on this subset with 300 iterations. If we assume the probe diameter is to where the intensity falls to 10% of the maximum, then the the overlaps are 73% and 46.4% for the full and sparsity sets respectively. Fig. 5 shows that rPIE and sir-DR obtain a larger FOV than ePIE as both use regularization. Furthermore, sir-DR removes noise more effectively and obtains a flatter background than ePIE and rPIE. We monitor the R-factor (relative error) to quantify the reconstruction, defined as

[TABLE]

$R_{F}$ is 16.94%, 13.95% and 13.28% for the ePIE, rPIE and sir-DR reconstructions, respectively.

3.2.2 Synchrotron radiation data

To demonstrate the applicability of sir-DR to synchrotron radiation data, we reconstruct a ptychographic data set collected from the Advanced Light Sources [16]. In this experiment, 710 eV soft x-rays are focused onto a sample using a zone plate and the far-field diffraction patterns are collected by a detector. A 2D scan consists of 7,500 positions, which span approximately $10\times 4\,\mu m$. The sample is a portion of a HeLa cell labeled with nanoparticles, which is supported on a graphene-oxide layer. Fig. 6 shows the ePIE reconstruction of the whole FOV of the sample. To compare the three algorithms, we choose a subdomain of a $4\times 4\,\mu m$ region, consisting of 2,450 diffraction patterns. With the same assumption, the overlap is computed to be 79.5%.

Fig. 7 show the ePIE, rPIE, and sir-DR reconstructions, respectively. With 300 iterations, all three algorithms converge to images with good quality. When reducing the number of iterations to 100, we observe that sir-DR converges faster and reconstruct a larger FOV than ePIE. The individual nanoparticles, which serve as a resolution benchmark, are better resolved in the sir-DR reconstruction than the ePIE and rPIE ones. Furthermore, the reconstruction by ePIE contains artifacts as a faint square grid, which is removed by rPIE and sir-DR.

We next perform a sparsity test by randomly picking 980 out of 2,450 diffraction patterns, i.e. a reduction of data by 60%. The corresponding overlap of the sparsity set is 50.8%. Fig. 8 shows the reconstructions by ePIE, rPIE, and sir-DR with 300 iterations. Both the ePIE and rPIE reconstructions exhibit noticeable degradation. In particular, the nanoparticles are not well resolved. But the sir-DR reconstruction has no noticeable artifacts noise and the individual nanoparticles are clearly visible. The quality of sir-DR reconstruction with 60% data reduction is still comparable to that of ePIE using all the diffraction patterns

4 Conclusion

In this work, we have developed a fast and robust ptychographic algorithm, termed sir-DR. The algorithm relaxes Douglas-Rachford to improve robustness and applies a semi-implicit scheme (semi-Backward Euler) to solve for the object and to expand the reconstructed FOV. Using both simulated and experimental data, we have demonstrated that sir-DR outperforms ePIE and rPIE with sparse data. Being able to obtain good ptychographic reconstructions from sparse measurements, sir-DR can reduce the computation time, data storage requirement and radiation dose to the sample.

Acknowledgments

This work was supported by STROBE: A National Science Foundation Science & Technology Center (DMR-1548924). J.M. also acknowledges the support by the Office of Basic Energy Sciences of the US DOE (DE-SC0010378).