Make Microscope and related classes nn.Modules
In preparation for #14, in which we'll want to directly estimate probe parameters, and potentially #13 (closed) and #9, we should make it easy to estimate camera parameters, particularly aberrations. Other microscope parameters like the CCD positioning and electron energy might also need some tweaking. However, this should all be user-guided so that we only tweak these values once the potential is well initialized.
Currently, CCD
, Aberrations
, and Microscope
are all Python dataclasses, which is useful for reducing boilerplate but does not result in automatic tracking of parameters. If they were nn.Module
objects we could easily include their parameters into an optimizer. Care must be taken then in all the microscope.py
code to accomodate having Tensor
-valued arguments instead of floats or ints. We should also provide tests that each of these things can be computed properly and that their gradients are computed correctly.
Related #12 (closed) #14
Plan
We can simply make the CCD
, Aberrations
, and Microscope
classes subclass nn.Module
, and implement constructors that set their attributes as nn.Parameter
objects. This means the Microscope
will have a full set of parameters, and the calling code will need to handle whether those are included in any optimizers.