Add forward model and fitting routine for rotation/magnification/defocus parameter extraction

[Copilot]

Implements forward model to generate synthetic TEM images and JAX-differentiable fitting routine to extract rotation, magnification, and defocus parameters from diffraction patterns.

Implementation

Forward Model (forward_model_collins.py)

• Collins FFT propagation with square aperture input (128×128, 7.8μm pixels)
• Parameters: 200kV, magnification 0.8-1.5×, defocus 10μm-5mm, rotation 0-60°
• ~1 sec/image generation time

Fitting Routine (rotation_magnification_defocus_fitting.py)

• JAX autodiff with Optax Adam optimizer
• Differentiable rotation/magnification via jax.scipy.ndimage.map_coordinates (bilinear interpolation)
• Fresnel propagation for defocus

Results (5-sample validation)

• Rotation: 0.32° ± 0.38° error
• Magnification: 1.55% ± 0.33% error
• Defocus: 29.5% ± 0.2% error

Technical Notes

Initial nearest-neighbor resampling broke gradient flow in JAX. Switching to bilinear interpolation in map_coordinates preserves gradients for rotation and scaling transforms:

# Differentiable rotation with gradient preservation
Y_coords = -sin(θ) * X + cos(θ) * Y + ny/2
X_coords = cos(θ) * X + sin(θ) * Y + nx/2
rotated = map_coordinates(field, [Y_coords, X_coords], order=1)

Also includes Gaussian beam reference implementation (rotation_magnification_defocus_forward.py) and debug utilities.

Files

• forward_model_collins.py - Forward model generator (356 lines)
• rotation_magnification_defocus_fitting.py - Fitting routine (436 lines)
• debug_fitting.py - Validation utilities (126 lines)
• ROTATION_MAGNIFICATION_DEFOCUS_README.md - Usage documentation
• IMPLEMENTATION_SUMMARY.md - Technical details
• forward_model_data/dataset.pkl - 50 synthetic images (6.3 MB)

Original prompt

I need to build a forward model based on the parameters described in n_lens_inversion - I want to actually be able to generate the images that I will need to fit to measure rotation, magnification and defocus as a test for the algorithm pipeline. I would firstly like to use the gaussian beam forward model, and a square aperture as the input, to generate a set of forward images. I think you might need to look through the temgym code and see if there is an easy way to have a lens that can rotate the beam. Once you have this, you can use gaussian beams and the parameters similar to n lens inversion to try and generate a dataset of magnified, rotated and defocused images. Once you are sure that these images seem reasonable (we get reasoanble rotations on the detector, all images land on the detector, are not too magnified, are defocused enough to fit B parameter), and we have a dataset, it would be good to then try write a fitting routing to fit all of these images. Maybe the best solution is to use the collins FFT above with a rotation and magnification part of the fitting routine, and to differentiate it with jax, or maybe the best solution is to differentiate with gaussian beams. I am not sure what is best. I think the problem with gaussians is there needs tobe enough beams, and thus we can only do it on small detectors (128x128 perhaps). Please reflect on this and decide what is best.

Created from VS Code.

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.