Reconstructing the Scattering Matrix from Scanning Electron Diffraction Measurements Alone

Three-dimensional phase contrast imaging of multiply-scattering samples in X-ray and electron microscopy is extremely challenging, due to small numerical apertures, the unavailability of wavefront shaping optics, and the highly nonlinear inversion required from intensity-only measurements. In this work, we present a new algorithm using the scattering matrix formalism to solve the scattering from a non-crystalline medium from scanning diffraction measurements, and recover the illumination aberrations. Our method will enable 3D imaging and materials characterization at high resolution for a wide range of materials.


INTRODUCTION
Phase contrast imaging is widely used in light [1,2], xray [3,4], and electron microscopy [5,6], due to its high efficiency and resolution. By using coherent radiation to illuminate a sample, we can resolve very small changes in a sample's local index of refraction through the interference of the illumination wave fronts that the accumulated phase shifts produce [7]. However, because we can only directly measure the probability density of a illumination wave function (given by the wave intensity, or amplitude squared), phase contrast imaging is a fundamentally nonlinear measurement process: we must indirectly infer the underlying relative phase shifts induced by the sample [8].
Various approximations can make phase contrast microscopy data easier to interpret. The first is by assuming that the sample is a pure phase object, i.e. it does not modulate the illumination wave function amplitude directly, and so any variations in the measured intensity can be directly ascribed to changes in the sample's index of refraction [9]. However this assumption does not guarantee uniqueness in all cases, due the possibly of phase wrapping [10]. An even stronger assumption is the weak phase object approximation (WPOA), where the sample's transmission function is assumed to be a small imaginary perturbation a known carrier wave [11]. When the WPOA holds, the linear relation implied between specimen potential and measured intensity allow constructive and unambiguous solution. Another commonly used simplification in phase contrast microscopy is the projection approximation * philipp.pelz@berkeley.edu † cophus@gmail.com (PA), where all scattering is assumed to originate from an infinitesimally thin 2D plane [12,13]. The various different approximations above hold for a wide range of samples of interest and are therefore very useful in practice [14].
However, phase contrast imaging of many samples cannot be approximated by any of the above assumptions. Transmission electron microscopy (TEM) in particular often violates these assumptions, due to high scattering cross section of electrons with matter [15]. Instead, these scattering processes can typically only be modeled by a framework that includes multiple scattering [16]. The equations describing multiple scattering for a paraxial wave function can be approximately solved with the multislice algorithm [17], which has also been used as a model for inverse scattering in many experimental configurations in light, X-ray-and electron microscopy.
While the inverse multislice model has been successfully applied to image thick, multiply scattering specimens in light microscopy [18][19][20][21], its use in X-ray [22][23][24][25][26] and electron microscopy [27][28][29][30] has been limited to proofof-principle demonstrations with less than 10 slices or weakly scattering samples. This is mainly due to the fact that the optical systems in X-ray and electron microscopy have relatively small numerical apertures, such that the information recorded from a single view covers only a small fraction of reciprocal space [31][32][33]. This problem can be overcome either by enforcing strong prior knowledge about the underlying scattering potential in the form of sparsity constraints or the proper choice of slice separation [30], or by performing tomographic experiments [34][35][36].
Another framework that incorporates multiple scattering is the scattering matrix (S-matrix) formalism [37,38]. In electron microscopy, the S-matrix formalism has been used to efficiently calculate diffraction results with single crystals [38] and for scanning TEM (STEM) experiments [39,40], and to retrieve projected potentials of strongly scattering samples in a two-step approach. First, the S-matrix is retrieved from a series of intensity measurements. Second, the projected structure is retrieved. The proposed experimental methods for retrieval of the S-matrix from intensity measurements range from measurements with different crystal thicknesses and sample tilts [41], different sample tilts alone [42][43][44], wavelength variation [45], large-angle rocking beam diffraction [46], and scanning diffraction with a convergent beam [47]. Only the last two of these approaches have been experimentally demonstrated [46,48], and only on single-crystal structures.
In the visible light wavelengths, S-matrix retrieval and subsequent singular value decomposition allows the identification of transmission eigenchannels [49] in strongly scattering materials and maximization of energy transport [50] through the system. Phase retrieval of the S-matrix is performed by real-space phase- [51] or amplitude-modulation [52,53], 4-phase interferometry [49], or full-field Mach-Zehnder interferometry [54] with input-and output channels in the plane-wave basis. The input and output channels of the S-matrix are often represented in real-space, achieved by imaging the output plane with a CCD camera.
Our contribution in this work is three-fold: first, we develop the measurement operator to calculate scanning diffraction intensities of arbitrary samples from a given S-matrix and derive its adjoint operator. Second, we formulate a phase retrieval algorithm that retrieves the S-matrix of arbitrary samples from a series of scanning diffraction measurements with different modulations of the illumination aperture (e.g. a defocus series). Third, we formulate a relaxation of the phase retrieval algorithm for samples that do not require the full S-matrix to be reconstructed.

II. RECONSTRUCTING THE S-MATRIX
A. Theory of phase contrast imaging Phase contrast microscopy with coherent light or matter waves defined by the wavefunction |ψ r typically uses a series of interferometric measurements to invert a partial differential equation of the form where i is the imaginary constant, ∇ 2 ⊥ is the twodimensional Laplace operator, V (r ) is the threedimensional potential over the real space coordinates r = (r, z), and a and b are real-valued constant prefactors. The formal operator solution to this equation for a wave function that has propagated a distance ∆z through the potential is given by [55], In the scattering matrix formalism, the entire process of multiple scattering is modeled by multiplication with the complex-valued linear operator S, The S-matrix formalism has a wide range of applications in describing the interaction of coherent waves with multiply scattering objects [56].
All the previously discussed methods for S-matrixretrieval at high resolution have in common that they require a crystalline sample to solve for either the scattering matrix or the structure factors. The interferometric methods developed for light optics rely on the ability to precisely manipulate phases and/or amplitudes of the S-matrix input channels and such precise control of the electron and X-ray optics is not yet feasible. In the following section, we describe our iterative reconstruction scheme from scanning diffraction measurements for S-matrix-retrieval.

B. A real-space S-matrix measurement model
Previous work for retrieving the S-matrix from scanning diffraction measurements modeled the formation of the diffraction pattern intensity in the far-field of the sample, given a coherent probe |ψ at position ρ, with an intensity measurement given by [47] In this work, we use the approximation that the wave function has a finite support after propagating through the specimen potential. To use this approximation as a constraint in an inversion algorithm, we need to represent the S-matrix in real space: Here S r,h is the S-matrix that maps Fourier-space input coefficents at wave-vectors h (we refer to these as the "beams" of the S-matrix) to real-space output coefficients at positions r. A previous experiment [48] used Eq. 5 and a series of defocus modulations to retrieve the phases of S q,h for a set of h vectors separately, and then used symmetry relations of S q,h to find the relative phases between the different S-matrix columns. Whereas that approach is only valid for crystalline samples, we use only self-consistency in the measured data and retrieve all amplitudes and phases of S r,h simultaneously. We also introduce a real-space compactness constraint on the scattered probes produced by the scattering matrix, equivalent to the method of Fourier-interpolating the Smatrix [39]. We introduce the cropping operator, a two dimensional rectangular function of width ∆ centered about each probe scan position ρ, which transforms Eq. 5 into (8) The fact that the cropping operator acts on all Smatrix beams equally leads to a self-consistent solution when measurements are taken with overlapping probe positions.

C. Phase retrieval of the S-matrix
We now describe an algorithm to retrieve all amplitudes and phases of S r,h simultaneously, given a set of phase modulations {χ d (h)} d=1,...,D of the probe-forming aperture, using only self-consistency in the measured data. Let the detector be sampled with M 1 × M 2 pixels. We perform a scan with K positions and D different probes and label a single position with k and a single defocus with d. Then the measured intensities have the dimension I ∈ R K·D·M1·M2 . For ease of notation, we enumerate all B samples in |h| < h max with indices b = 1, ..., B. The S-matrix measurement operator maps the B beams of the Smatrix of sampled on a discrete grid of N 1 × N 2 pixels and the D probes to K · D diffraction patterns of size we first define the measurement operator for position k and probe d: where [·] V is a vectorization from 2D to 1D. We have also introduced the linear cropping operator C k,d := C ρ k,d : C B×N1×N2 → C B×M1×M2 , which extracts a real-space patch of size M 1 × M 2 from each beam of a given S-matrix at the position with index k for the phase modulation d. The measurement operator for the full experiment is just the operators for each probe and position stacked on top of each other: T We can then write the forward model for the measured intensities of a series of D scanning diffraction experiments taken with different probes as Given this forward model and a set of intensity measurements I we can formulate the phase retrieval problem for blind S-matrix inversion as Find S ∈ C B×N1×N2 and Ψ ∈ C D×M1×M2 Subject to |A(S, Ψ)| 2 = I.
If the wave functions Ψ are known, the problem of finding S from a set of measurements I is a classical phase retrieval problem. There is a rich history of a algorithmic developments to solve the phase retrieval problem. Historically the first were algorithms based on alternating projections onto non-convex constraint sets [57][58][59]. Since these algorithms lack theoretical convergence guarantees, more recently convex relaxations were developed [60,61] which provide a convergence guarantee, but use a prohibitive amount of memory. More recently, Bayesian accelerated gradient methods [62] and methods based on the alternating direction method of multipliers (ADMM) [63] have become popular. Since the wave functions Ψ d are usually not known precisely in advance, the problem turns into multi-objective optimization. Additionally, in the presence of noise, it is beneficial to the reconstruction quality to include the noise model of the detector in the optimization. Since most advanced detectors in X-ray and electron microscopy are counting detectors, the noise statistics follow a Poisson distribution: I ∼ Poisson(y).
Here we choose an amplitude-based cost function as an approximation to the Poisson likelihood, due to its better convergence behaviour and divergence-free derivative [64,65]: where · 2 is the l 2 norm and y are far-field amplitudes of the current model. We use the ADMM algorithm [66] to solve the joint optimization problem of S and Ψ. The augmented Langrangian of the S-matrix retrieval problem is where we have introduced the auxiliary variables z ∈ C K·D·M1·M2 and Λ ∈ C K·D·M1·M2 , which link the dataloss term with the model-loss term. We seek to solve for S and Ψ such that L(S, Ψ, z, Λ) is minimized: ADMM decouples the joint problem into subproblems and solves them step by step: 1. Ψ l+1 = arg min Ψ L Ψ β := arg min Ψ L β (S l , Ψ, z l , Λ l ) 2. S l+1 = arg min S L S β := arg min S L β (S, Ψ l+1 , z l , Λ l ) 3. z l+1 = arg min z L β (S l+1 , Ψ l+1 , z, Λ l ) The subproblems with respect to Ψ and S both involve the adjoint of the measurement operator A, which for a single measurement is given by for a fixed Ψ, and A for a fixed S b . We solve the subproblems with respect to Ψ and S with gradient descent.
where γ 1 , γ 2 ∈ R are gradient descent step sizes. We found that one gradient step per iteration is usually enough for fast convergence. The gradient is given by See the Appendix B for a detailed derivation.
The subproblem w.r.t. z was solved elsewhere [67]. The solution is The full ADMM algorithm is then given as: Algorithm 1 Joint S-matrix and probe retrieval via ADMM Input: measured intensities I ∈ R K×D×M 1 ×M 2 scan positions ρ ∈ R K×D×2 initial Fourier space probe phases χ 0 ∈ C D×B step sizes γ1, γ2, β ∈ R Initialize: set (N1, N2) = max(rs)+M M · M such that the plane waves e ih·r have periodic boundary conditions calculate I mean = 1

III. SIMULATED S-MATRIX PHASE RETRIEVAL
In this section, we use forward simulations to validate our S-matrix phase retrieval algorithm. We also examine the algorithm dependence on the sampling density and calibration.

A. Sampling and calibration dependence
To demonstrate that our algorithm can reconstruct Smatrices of realistic samples, we simulate a 4D-STEM focal series of the sample shown in Fig. 1 a), as it may appear in a tomography experiment. The sample contains two decahedral Ag nanoparticles of 3.   defocus at the top of the sample. The detector size was set to 128 × 128 pixels, yielding an angular resolution of 0.31 µrad and S-matrix dimensions of S ∈ C 5973×256×256 .
We ran Algorithm 1 for 500 iterations, utilizing 48 NVIDIA V-100 GPUs. After 200 minutes, the reconstruction converged to an normalized root mean square error (NRMSE) of 4 % and an R-factor of 0.1 %. Nine selected S-matrix beams from the reconstruction are shown in Fig. 2 a, and the ground-truth S-matrix is shown in 2c. To investigate the convergence properties under varying number of measurements and calibration errors, we used a smaller test sample, consisting of 16 randomly distributed Germanium atoms in a volume of 5Å × 5Å × 100Å, shown in Fig. 2d. The convergence angle for the following tests was chosen as 30 mrad, with a detector spanning 60 mrad, and the diffraction patterns were sampled on a 20 × 20 pixel detector, yielding Smatrix dimensions of S ∈ C 177×60×60 , and the defocus step was chosen as 2 nm.
For the following investigations we fix the scan step to Nyquist sampling. First we investigate the converge behaviour with respect to the number of measured defoci. Fig. 2e and f show the R-factor, and the NRMSE as a function of iterations and number of defoci measured. We define the oversampling factor as and the bright-field oversampling factor as One can see that for 2 defocus measurements, the NRMSE diverges slowly, and for 3 measurements the NRMSE does not converge monotonously with the Rfactor. While the oversampling factor O lies above the number 4 typically needed for successful phase retrieval, the number of phase modulations that each beam receives, O BF , is below the threshold. For this case, a more heterogeneous sample than the crystalline objects considered in previous work, the reconstruction does not stably converge in these cases. This could be due to the small defocus steps used and will be investigated in the future.
We also investigate the dependence of the probe refinement on the level of defocus miscalibration and residual uncorrected probe aberrations. Fig. 2 h) shows the mean errors of 30 reconstructions performed with defocus errors ∆C 1 drawn form a normal distribution with a standard deviation of 10 %, 20 % and 30 % of the defocus step, axial coma with a standard deviation of 100 nm, three-fold astigmatism with a standard deviation of 20 nm, spherical aberration with a standard deviation of 4 µm, and star aberration with a standard deviation of 4 µm. Although convergence takes roughly twice as many iterations S-matrix-reconstruction with miscalibrated aberrations, for all miscalibration values a probe reconstruction error of less than 10 % was achieved.

B. Reconstructing the projected S-matrix
Consider the scattering matrix for a phase object, which is a valid approximation for a thin and weakly scattering sample [14], with specimen potential V (r). The analytic expression for each component will be, So every S-matrix component will be the same except for the multiplicative phase ramp of e −2πih·r . As we consider thicker, more strongly scattering objects we would expect each component of the S-matrix to be increasingly different and we consider the similarity or lack thereof of each of the S-matrix components to be an indication of the degree of strong multiple scattering of a sample. When reconstructing the S-matrix from a 4D-STEM dataset the automatic choice for choosing the sampling of beams, the set of h vectors, is to match it to the number of pixels within the bright-field disk or aperture function of the STEM probe in the diffraction patterns. For a fine diffraction space sampling of an object that does not exhibit much multiple scattering this sampling of beams might be highly redundant and we might improve our reconstruction by forcing a more sparse sampling of beams and increasing the ratio of experimental measurements to unknown parameters in our reconstruction. On the other hand, very thick and strongly scattering samples might require very high sampling of the diffraction patterns for an accurate reconstruction of the S-matrix. While the latter case can only be solved with better sampling in the diffraction plane, for the former case in this section we outline a strategy for choosing a sparser sampling of the input beams h that involves partitioning of the bright-field disk into separate "tiles". Fig. 3 are the complex values of a subset of S-matrix components for a) 7.3Å, b) 36.5Å and c) 109.5Å thicknesses of an ScAlO 3 crystal. All beams have been multiplied by the conjugate of the phase ramp that appears in Eq. 23, e 2πih·r . For Fig. 3(a) the difference between beams is minimal so a phase object approximation would be appropriate for this thickness. For Fig. 3 b)-c) we see increasing variation between beams as the object becomes thicker. A partitioning system aims to group these S-matrix components by similarity and visual comparison of the S-matrix montages with the Fresnel propagator, P(h) = exp −iλπh 2 t , for free-space of equivalent thickness of the crystal (shown to the right of each subfigure) suggests that a criterion based on phase variation of a Fresnel free-space propagator might be an effective way of doing this. We partition the bright-field disk into annular regions (i − 1)∆φ/λπt < h i < i∆φ/λπt where i ∈ N and ∆φ, the Fresnel propagator phase variance, is a predetermined criterion (we use ∆φ = π/4 in this work). These regions are further divided azimuthally, with an arclength equal to the radius of the inner-most partition, ∆φ/λπt. We represent these partitions with the map τ : {0, ..., B} → {0, ..., B tile } from beams to beam tiles. Partitioning according to his criterion is shown to the right where each different color indicates a separate partition of the bright-field disk for each of the thicknesses in Fig. 3(a)-(c). We note finally that this is an approximate criterion only since thickness of an uncharacterised object can only be guessed at based on the intuition of the microscopist and the Fresnel criterion does not take into account the scattering strength per unit volume of the object. For example we might expect samples containing a high density of heavy (large Z) elements to exhibit greater beam to beam variation of the S-matrix components than materials with only small Z numbers.

Shown in
Partitioning reduces the number of needed measurements by a factor B tile /B.
The reconstructed S-matrix then has the reduced dimension S ∈ C B tile ×N1×N2 , and the beam-dependent beam-tilt e −2πih·r has to be separated from the S-matrix variable to allow the beamaveraging. We define therefore the un-tilted S-matrix S t r,h := S r,h e −2πih·r , which is the new latent-variable in the projected S-matrix problem. The forward operator becomes and the gradient with respect to S t is where we have introduced the set Q b = {n | τ (n) = b ∀ n ∈ {1, ..., B}} of beams that belong to tile b and |Q b | is the cardinality of Q b .

IV. CONCLUSION AND OUTLOOK
We have introduced a new method for S-matrix retrieval, that converges without any regularization for samples which span 4 depths of focus and more, and numerical apertures which are experimentally accessible, and can recover aberration miscalibrations of up to 30 %. We have also introduced a simplified model, projected Smatrix inversion, for the case when the sample is thin enough that not every beam that is measured on the detector has to be included in the model. In future work, we will compare projected S-matrix inversion to mixedstate ptychography and multi-slice ptychography, since both offer alternative methods for moving beyond the simple model of single-mode ptychography.
The S-matrix-retrieval methods developed here could be used for a number of advancements in imaging through and with strongly scattering materials in X-ray and electron microscopy. In combination with adaptive electron optics [68], selective focusing through crystalline materials may become possible in a similar vein to light Partitioning of the S-matrix beams into separate tiles according to the expected degree of departure from the phase object approximation. In each subplot a subset of the complex components of the scattering matrix (color hue is phase and color saturation is amplitude according to inset colorwheel) are shown for a ScAlO3 crystal with (a) 7.3Å, (b) 36.5Å and (c) 109.5Å thickness. With increasing thickness there is increasing variability between the different components. The beam partitioning suggested by a phase variance of the Fresnel freespace propagator of π/4 described in the text is inset on the top right for each of the thicknesses in (a)-(c) where each colour in the bright-field disk corresponds to a seperate B tile optical experiments [69].
The retrieved S-matrix can be used for depth-sectioning which is robust against multiple scattering.
Smatrix-retrieval may also form the basis of inverse multi-slice algorithms for phase-contrast tomography in scanning diffraction microscopes.
The angular decomposition in the S-matrix may be useful for abinitio angular and transverse alignment of different tilt angles for phase-contrast tomography. This approach may be experimentally more feasible than end-to-end tomographic reconstruction algorithms. Finally, one could think about characterizing amorphous materials from their S-matrix. To allow optimal image quality, future refinements of the algorithm could include experimental uncertainties like position errors, and modeling of nuisance parameters like spatial and temporal incoherence, similar to their treatment in ptychographic reconstruction algorithms [70][71][72][73][74].