PRISM-style plane-wave simulation ptychography
PRISM
In Ophus (2017)1, Colin Ophus introduced PRISM, which is a method for accelerating multislice simulations. The key observation is that the multislice propagation is linear (the "S matrix" links probe Fourier coefficients to complex CBED values). Also, since beams are shifted using the Fourier shift theorem, the Fourier basis vectors for the probe wavefunction can be mapped to a Fourier basis of complex CBED patterns, and the shift can be applied after multislice simulation. This is tremendous since it means that CBEDs can be simulated at many locations with very little computation after the "plane-wave" (Fourier basis functions) probes are simulated, by simply computing a new complex-weighted sum of plane-wave simulations which serve as a non-orthonormal basis of simulated diffraction patterns.
In PRISM, this scheme is implemented by first approximating the probe wavefunction with a small collection of frequencies, determined by laying down a coarse grid in the Fourier domain. This reduces the flexibility in the probes that can be represented, which may lead to aliasing when approximating a particular probe. This effect is not considered in the original PRISM paper.
An additional aspect of this approach is that although the object function is not a simple shift-invariant convolution it still tends to map an impulse (plane wave) to a localized CBED pattern, especially for thinner samples. That is to say the plane-wave CBED patterns look like impulses at the plane's frequency location that are slightly broadened, with a fringe due to the complexities of the scattering process. The limited extent of this broadening limits the sparsity of the grid that can be used in PRISM, since the broadened CBEDs must overlap in order to obtain contrast between beam locations. We can interpolate plane-wave responses in order to maintain contrast even with high frequency spacings but that is a subject for a future issue.
Plan
We will implement an nn.Module class called SparseProbeMultislice. It will use a sparse grid in the Fourier domain to represent an unshifted probe. In that class's basis_patterns() method, each sparse grid point will be mapped to a plane wave and simulated using a given potential, and the CBED patterns will be returned as a 4D complex tensor. These basis diffraction patterns can then be combined with a phase shift in order to simulate any given complex CBED pattern. We will provide an example of how to use this in the form of a Jupyter notebook or example script.