multislice simulator

[mjo22]

This package may be great for implementing this: https://github.com/patrick-kidger/diffrax

[geoffwoollard]

Some resources on how to implement https://abtem.github.io/doc/user_guide/walkthrough/multislice.html https://www.ccp4.ac.uk/ccp4-ed/projects/multislice/ https://hamishgbrown.github.io/py_multislice/pyms/py_multislice.html#pyms.py_multislice.multislice

[mjo22]

Some thoughts on implementation:

1. See the function _iterate here for some control flow for an iterative solver in JAX, and the function loop here for numerical integration control flow. Patrick has implemented a while loop that can be backpropagated in the equinox.internal submodule. The while loop in core JAX only supports forward-mode autodiff.
2. I imagine multislice fitting in as an AbstractPotentialIntegrator, which takes in an abstract interface for an atomistic representation of a potential (which does not yet exist!).
3. We will need to generalize the cryojax imaging pipeline. Right now, we go from 3D potential -> phase shifts -> squared wavefunction but we will need to support potential -> wavefunction -> squared wavefunction. I expect this generalization will involve work in the Instrument and AbstractSpecimen, and creating a new AbstractOptics model for a complex-exponential CTF. Should be a good opportunity to improve upon the code.

[geoffwoollard]

Re 2, I drafted some mulislice code that slices through a discretized potential, and produces the exit wave.

It takes in the potential, slice thickness, and beam wavelength (lambda), and produces the exit wave.

Can you draft something that I can put the core multislice code into?

[mjo22]

Starting to implement handling passing an exit wave through the imaging pipeline. A few comments

1. One issue in this is that the ImageConfig needs to store a full frequency plane, rather than the half plane needed for the fourier transform of a real-valued function. For now, we’ve just computed more frequency coordinate systems upon ImageConfig.__init__, but this is too cumbersome I think. Instead, I wonder if all coordinate systems in pixels can be computed at run-time and stored in the ImageConfig as python properties (using @property or @cached_property). I bet XLA will know to compute the meshes only once.

2. Okay also, it would be good to have a way that the user can externally check if a particular integrator returns an exit wave or a phase contrast. This can either be done through a ClassVar, or by making abstract interfaces that subclass AbstractPotentialIntegrator. Perhaps names like AbstractMultiSliceIntegrator and AbstractProjectionApproximation would be appropriate.

3. The current treatment of the solvent won’t be compatible with passing an exit wave through the pipeline, since the solvent model just randomly samples a phase contrast field. We may need to generalize this to turn the phase contrast into an exit wave, or just add a run-time error that deems them incompatible. In general, the solvent model needs work. Also, the AssemblyPipeline relies on linearity and is fundamentally incompatible with multislice, so this will need a run-time error.

[geoffwoollard]

FYI the multislice treatment (for material science applications, including using density functional theory for the potential) in https://abtem.readthedocs.io/en/latest/intro.html follows Eq 6 in Vulović, M., Ravelli, R. B. G., van Vliet, L. J., Koster, A. J., Lazić, I., Lücken, U., … Rieger, B. (2013). Image formation modeling in cryo-electron microscopy. Journal of Structural Biology, 183(1), 19–32. http://doi.org/10.1016/j.jsb.2013.05.008

[mjo22]

Added an initial commit in the multislice-updated branch

[geoffwoollard]

Recent TODOs

• [ ] Wave transfer function: functional form
• [ ] Amplitude contrast: bit of an open question. what to do. should agree to projection of one slice in a limit? or do we need separate amp const theory (is mutlislice pure phase)
• [ ] envelope for transfer function. https://github.com/mjo22/cryojax/blob/multislice-updated/src/cryojax/simulator/_transfer_theory/contrast_transfer_theory.py
• [x] translational phase shifts: can apply to wave, or have to do after ||^2. https://github.com/mjo22/cryojax/blob/multislice-updated/src/cryojax/simulator/_scattering_theory/base_scattering_theory.py#L82
• [ ] tests
  • [ ] one slice should match projection assumption
  • [ ] stacking objects and they are different that projecting
  • [ ] residuals compared to projection is close (including with ctf/wtf) - use imaging pipeline for projection comparison
  • [ ] compare to high energy approx: change scattering theory and transfer function

[mjo22]

Translations should be taken care of!