Achieving modularity in cryojax

[mjo22]

With the new factors I've been thinking a lot of how to really achieve modularity in cryojax.

The most notable additions are the FourierOperator and RealOperator classes. With this, I've started to expand on a paradigm that I already had loose implementations of: cryojax.simulator models should take functions as input.

For example, consider the new, refactored, CTFOptics class. This used to be a function which computed the CTF when called. Now, it is more so what I would call an optics model. When it is called, it operates on the image somehow. This involves applying the CTF function---the new CTF class, a subclass of a FourierOperator---and also calling the envelope function, also a FourierOperator.

Note two cool aspects of this.

1) This allows us to interchangeably use phenomenological and physical models, both of which are very common in cryo-EM. In one case, an envelope function could be a gaussian, parameterized by a phenomenological beta factor. In another case, we could create an envelope function that is a gaussian whose parameters model chromatic aberration. Further, this even allows us to interchangeably use analytical and empirical models! There is an operator called Empirical, which basically just stores a buffer. 2) These functions---namely these FourierOperator and RealOperator classes---can be composed before ever being called with multiplication * and addition + operators, using magic methods __mul__ and __add__. This makes for some cool behavior. For example, see the cryojax.inference.distributions.FourierIndependentGaussian.__init__.

Overall, continuing on this paradigm I think a) should give users a ton of control when scripting and b) continue to make it so us developers need not have way too many subclasses! Expanding on a), note that FourierOperator and RealOperator are public API! A user should be able to subclass these and pass their own models into cryojax, without needing to make a PR. Expanding on b), note that there is no need to have two different CTFOptics subclasses for the different envelopes. The user chooses which envelope model to pass, and even whether or not to pass one at all.

For anyone interested in learning more about and expanding on this paradigm (@geoffwoollard? @ehthiede ?), here is one refactor that should be done:

• Right now, we have a different Pose object for each kind of rotation---euler, quaternion, and matrix. However, a Pose is a rotation and translation.
• Really, we should have one Pose object. This Pose should somehow have a Rotation object, which can be Euler, Quaternion, or Matrix.
• If we want to go crazy, this Rotation can adhere to similar principles as the FourierOperator class. It can use the __matmul__ magic method to compose rotations of different types.

Then, similar to the CTFOptics class, the user will similarly be passing an operator to the Pose class, rather than choosing a Pose class with particular parameters!

For more examples of how I use FourierOperators, see

• GaussianDetector.variance and GaussianIce.variance
• Exposure.dose and Exposure.radiation (a new attempt to give the Exposure model a vision, eventually giving functionality for something like here: https://elifesciences.org/articles/6980)

[geoffwoollard]

I'll finish some PRs first, but I'd be interested in generalizing pose to have a handful of back end encodings

[mjo22]

Awesome. Note we already have a QuaternionPose, EulerPose, and MatrixPose, but I think it would be more powerful to just have one class called Pose and to pass a Rotation operator to it, which could be these different representations.