math.txt

This is intended to be a brief explainer on our implementation of the potential-generation and multislice algorithm. For a full derivation/explanation/references on the math, the reader is encouraged to check out "Advanced Computing in Electron Microscopy" by Earl J. ψ(x,y,z+Δz) = ℱ⁻¹{ P(x,y,Δz) * ℱ{ t(x,y,z) * ψ(x,y,z) } } (https://doi.org/10.1007/978-3-030-33260-0). This test document focused on the relevant bits, and the nitty gritty on how the math is implemented in code. 

POTENTIAL GENERATION:
We begin with generation of the potential. A series of atoms (typically saved by their coordinates, i.e., we have a list of x,y,z coordinates) needs to be broadcast onto a grid of points with a chosen intensity profile around each atom. We will begin discussion in real space, since this is generally more intuitive to more people, but eventually we will build the potential in reciprocal space (or fourier space). 

In 1D: 
atomic_cordinates=[1.1,2,3.6,5.75]
needs to be converted into (for example) a gaussian at each point:
grid pts :    1      2      3      4      5      6      7
intensity: __.-'-._.-'-.______.-'-.__________.-'-.__________
We can do this in real space, by creating a grid of points, then "for each atom, for each grid point, calculate the distance and then the intensity for that grid point". This is slow. 

Instead, we can place a delta function at each atom's nearest grid-point, then convolve this with our potential function (a single gaussian in this case). 
grid pts :    1      2      3      4      5      6      7
delta fun: ____|_____|__________|______________|____________
1x gauss : ______________________.-'-.______________________
convolved: __.-'-._.-'-.______.-'-.__________.-'-.__________

But this means we must "round" the atomic positions to the grid. We can use a finer and finer grid, but fundamentally this "rounding" may not be desired. 

Instead, we make the observation that a delta function in real-space is a sinusoid in reciprocal space:
ℱ{ f(x) } = ∫ f(x) exp( -i 2 π k x ) dx = f(k)
so if f(x)=δₓ (unity at position x, and zero elsewhere)
ℱ{ δₓ } = ∫ δₓ exp( -i 2 π k x ) dx = f(k) = exp( -i 2 π k x )
i.e., this is a sinusoid across k-space with a fixed periodicity "x"
We also observe that fourier transforms are additive:
ℱ{ f(x) + g(x) } = ℱ{ f(x) } + ℱ{ g(x) }

This means we can instead start with a reciprocal-space grid of points, "paint in" sinusoids for each atom, and iFFT back. 
1st atom : .-'`'-._.'`'-._.'`'-._.'`'-._.'`'-._.'`'
2nd atom : /'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'
3rd atom : __..--'''--..__..--'''--..__..--'''--..__
these are summed, and the iFFT will show a series of (sub-pixel-resolution) delta functions. 

Finally, noting that convolution in real-space is simply multiplication in fourier space, we simple multiply the FFT of the gaussian prior to taking the iFFT. And for the case of a more-realistic potential (e.g. Kirkland's Equation C.15), the potential itself can be defined in reciprocal space. 

The benefit of this approach is speed: all atoms can be placed on the grid in parallel, and no "neighbor finding" is required between atoms and grid points. 

This is implemented in code inside src/multislice/potentials.py. The class "Potential" generates the reciprocal-space grid points (variables kxs and kys) via fft.fftfreq from the real-space grid, then "paints in" complex sinusoids (variables expx and expy), and retreives the potential (from the function "kirkland"). 

PROBE GENERATION:
In a STEM, the probe is generated by passing the beam through a round aperature (the "virtual objective aperture" or "VOA"), and lenses perform a fourier transform which yields a focused probe with an "airy disk" intensity profile. We do the same. 

Inside src/multislice/multislice.py, the class Probe generates the reciprocal-space grid points (variables kxs and kys) via fft.fftfreq from the real-space grid, creates a circular mask of intensity (unity inside a given radius, determined by the selected convergence angle, zero elsewhere), and then the real-space probe is taken as the iFFT.

PROBE POSITIONING: 
During potential generation, we were able to apply "sub-pixel" positioning of each atom's gaussian (or whatever potential's intensity profile) by placing a delta function where we wanted the atom (generated in reciprocal space via a sinusoid) and convolving it with the gaussian "kernel function". Probe position is handled in exactly the same manner. The original probe profile is the kernel function, a sinusoid is generated in reciprocal space based on the desired placement of the probe. the two (reciprocal-space) functions are multiplied, and the iFFT yields a spatially-shifted probe. This is done inside src/multislice/multislice.py in the function "create_batched_probes", via what are termed as "phase ramps" (variables kx_shift and ky_shift). 

MULTISLICE: 
The propagation of the probe (as an electron wave) through each layer of the potential is performed as a series of multiplication steps in real or reciprocal space, simply following the equations from Kirkland, Equations 6.66 and friends:
ψ(x,y,z+Δz) = 𝒞{ p(x,y,Δz) , t(x,y,z) * ψ(x,y,z)	# Eq 6.66
where ψ(x,y,z) is the incoming electron wave, and ψ(x,y,z+Δz) is the outgoing electron wave. 𝒞{ } denotes the convolution, which is the multiplication in reciprocal space, meaning this is equivalent to:
ψ(x,y,z+Δz) = ℱ⁻¹{ P(x,y,Δz) * ℱ{ t(x,y,z) * ψ(x,y,z) } }
where P is termed the "propagator function", defined in Eq. 6.65 as:
P(kx,ky,Δz) = exp( -i π λ ( kx² + ky² ) Δz )
and t is termed as the "transmission function", defined in Eq 6.59 as:
t(x,y,z) = exp( i σ ∫ V(x,y,z') dz' )
where V is the potential. Note that our potential calculation previously is actually the *projected* potential, i.e., the potential we have calculated and stored is already ∫ V(x,y,z') dz'. 

This code can be found inside src/multislice/multislice.py in the function Propagate. We calculate variable "P" (in reciprocal space from variables "kx_grid" and "ky_grid"), then loop through layers computing variable "t" (in real-space based on slices taken from the potential). The t(x,y,z) * ψ(x,y,z) term is calculated (variable "array"), FFT'd, multiplied with P, and iFFT'd: ψ(x,y,z+Δz) = ℱ⁻¹{ P(x,y,Δz) * ℱ{ t(x,y,z) * ψ(x,y,z) } }

DEFOCUS: 
Considering the transmission function as:
t(x,y,z) = exp( i σ ∫ V(x,y,z') dz' )
if V(x,y,z) = 0, then t(x,y,z)=1+0j
ψ(x,y,z+Δz) = ℱ⁻¹{ P(x,y,Δz) * ℱ{ t(x,y,z) * ψ(x,y,z) } }
simplfies to:
ψ(x,y,z+Δz) = ℱ⁻¹{ P(x,y,Δz) * ℱ{ ψ(x,y,z) } }
or:
ℱ{ ψ(x,y,z+Δz) } = P(x,y,Δz) * ℱ{ ψ(x,y,z) }
which means I can "back-propagate" by dividing by P (negative defocus, beam waist is "in" the sample)