Modernize lens inversion notebook: JAX+BFGS, image-based fitting, complete inversion pipeline - #9

Draft

DavidLanders95 with Copilot wants to merge 186 commits into

mainfrom

copilot/vscode-mlcbwylf-12h0

Commits on Sep 18, 2025

trying to figure out local curvature change
DavidLanders95
committed
db68c3a

Commits on Sep 19, 2025

About to implement local hessian solution
DavidLanders95
committed
9678453

Commits on Sep 20, 2025

Opl difference for one local gaussian seems to work
DavidLanders95
committed
af4d26d
Added gaussian evaluation routine
DavidLanders95
committed
adfcecb

Commits on Sep 22, 2025

Commits on Sep 23, 2025

Not working potential
DavidLanders95
committed
7bbc136
going to simplify further
DavidLanders95
committed
d59d66f
simplified taylor solution
DavidLanders95
committed
b01096b
simplified taylor solution saved
DavidLanders95
committed
cbb7cac
simplified taylor solution - removed tangent and normal calculation
DavidLanders95
committed
dd1db45
single_ray_single_atom works
DavidLanders95
committed
ea34774
Functional version against handcoded single ray single atom works
DavidLanders95
committed
191aed3
Single atom multiple rays seems to work
DavidLanders95
committed
2b3dc91

Commits on Sep 24, 2025

Rearranged files tidily in neccessary folders
DavidLanders95
committed
a15a629
Aperture working
DavidLanders95
committed
7729158

Commits on Sep 26, 2025

Gaussian Taylor working with simplified design
DavidLanders95
committed
a38490b
Removed unneccessary functions
DavidLanders95
committed
a3b944e
Fixed up evaluate.py
DavidLanders95
committed
b3a40ee

Commits on Sep 27, 2025

Commits on Sep 29, 2025

aberrated lens working again
DavidLanders95
committed
8bc3915

Commits on Oct 7, 2025

Finalising aberrated fit with gaussians.
DavidLanders95
committed
86bbb2c
Aberration of lens with many gaussians beginning to work.
DavidLanders95
committed
6bbb588
Renamed aberrated lens
DavidLanders95
committed
5c64dd2
adjusted and cleaned evaluate.py functions
DavidLanders95
committed
89b564a
Fixed bug in gaussian_taylor where ray.dx was coming out with an extra dimension after vmap
DavidLanders95
committed
e2123f7
Aberrated lens runs with gpu_kernel and is quite fast
DavidLanders95
committed
6bb9bb6
Edited pallas kernel and added single gaussian test
DavidLanders95
committed
6772714
Very important sign change added to gaussian_taylor.
DavidLanders95
committed
b36bb18
Aberrated lens and individual biprism seem to work.
DavidLanders95
committed
f7255fc
Adjusted settings in aperture notebook
DavidLanders95
committed
36b2288
Removed unneccessary notebooks
DavidLanders95
committed
f114e26
Edited sign of single gaussian so it matches forward model now
DavidLanders95
committed
6d3438b
Renamed aberrated lens simple
DavidLanders95
committed
f6e0be3
Created aberrated probe notebook
DavidLanders95
committed
6ecaebb
Added complete Krivanek lens to Gaussian Taylor
DavidLanders95
committed
5d5daf3
Added CollinsPropagator, and FourierTransform Functions - Still need to test them. Will replace unmaintainable GaussianBeamCode
DavidLanders95
committed
d419aa7

Commits on Oct 10, 2025

Commits on Oct 11, 2025

Commits on Oct 13, 2025

Commits on Oct 14, 2025

Added magnetic phase sample and start of notebook
DavidLanders95
committed
2c23217

Commits on Oct 16, 2025

rough gaussian physics solution explained with interactive sliders. About to modify it some more though because the current representation is not ideal
DavidLanders95
committed
aa71885

Commits on Oct 20, 2025

Commits on Oct 22, 2025

Simplified Gaussian1D some more in preparation for .py file
DavidLanders95
committed
4729fd9
Removed comment
DavidLanders95
committed
441b031
Moved gaussian code to .py file and verified it still works.
DavidLanders95
committed
39b3c91

Commits on Oct 23, 2025

Gaussian1D and 2D seem to work.
DavidLanders95
committed
4f3e1f5
Moved explainer files
DavidLanders95
committed
a4b2314
Modified BeamFactory to work with new gaussian implementation
DavidLanders95
committed
4e2ebc5

Commits on Oct 28, 2025

Commits on Oct 29, 2025

Commits on Nov 3, 2025

Resetting again before implementing simpler propagator
DavidLanders95
committed
5b17dc7

Commits on Nov 4, 2025

Commits on Nov 5, 2025

Commits on Nov 6, 2025

Sample Hologram tidied up a bit but works
DavidLanders95
committed
06204d1
Moved files to utils, tidied up sample phase hologram a bit more
DavidLanders95
committed
b24e16a
on-axis distortion seems to work
DavidLanders95
committed
39fcb94
Off axis and on axis aberrated hologram programmes are made, but I have no idea what I am looking at.
DavidLanders95
committed
96f331e
Found parameters to make distortion field
DavidLanders95
committed
118eb42
Lowered simulation size
DavidLanders95
committed
23bb058
Removed old gaussian code
DavidLanders95
committed
26ce3de
Removing old code and simplifying gaussian.py
DavidLanders95
committed
3a2e54b
Tidied up gaussian.py and moved functions to utils.py
DavidLanders95
committed
f973125
Added offset to input plane wave object
DavidLanders95
committed
06ab5d8

Commits on Nov 7, 2025

made elaborate sample notebook, but not working yet
DavidLanders95
committed
fd7f182
Created aberrated Sample notebook and it seems to work ok.
DavidLanders95
committed
6728bda
changes to get model working for malika
DavidLanders95
committed
e1c8b5b

Commits on Nov 13, 2025

hard to get chosen fringe spacing to line up with model
DavidLanders95
committed
3099be9
Down to 6nm fringe spacing and with the equations to automate it
DavidLanders95
committed
44ef710
Updated lens_biprism with some titles.
DavidLanders95
committed
8653238
Added in overlap factor instead of num_rays argument to square wave
DavidLanders95
committed
f2b1ded
added overlap factor to rectangular_input_wave
DavidLanders95
committed
d96fb29
Working with projector lens turned off
DavidLanders95
committed
54de327
Readded in projector lens
DavidLanders95
committed
464a697

Commits on Nov 17, 2025

Commits on Nov 20, 2025

Added unit option to gaussian beam
DavidLanders95
committed
f0d8ace

Commits on Nov 21, 2025

Commits on Nov 23, 2025

Ray method still won't converge with trying some new parameters
DavidLanders95
committed
65681b5
Added probe input wavefunction
DavidLanders95
committed
99a9aee

Commits on Dec 1, 2025

Commits on Dec 5, 2025

Ran tests, simulation does seem to work for small non-linear phase plates
DavidLanders95
committed
44141c3

Commits on Dec 8, 2025

Commits on Dec 10, 2025

Can get ball park MIP measurements. Still off by a volt though.
DavidLanders95
committed
3b58c68

Commits on Dec 15, 2025

Removing buggy phase of square
DavidLanders95
committed
fd3bea6

Commits on Dec 22, 2025

Managed to get a simulation working - Fringe spacing of utmost importance. I need to find a way to non dimensionalise everything.
DavidLanders95
committed
ccb367f

Commits on Feb 4, 2026

Commits on Feb 5, 2026

SSIM with single lens and lbfgs seems to work ok.
DavidLanders95
committed
7ab7e63

Commits on Feb 6, 2026

Did quick test with simpsons rule, I am on the right track with gaussians sadly. The problem is just crazy.
DavidLanders95
committed
a5c0421
Checkpoint from VS Code for cloud agent session
DavidLanders95
committed
546ee4a
Create simplified two-lens notebook with Collins FFT and Optuna integration

Copilot
and
DavidLanders95
committed
85b29b4
Add test script and documentation for simplified two-lens notebook

Copilot
and
DavidLanders95
committed
9f42abb
Fix code style issues in test script

Copilot
and
DavidLanders95
committed
74e2fad
Add comprehensive implementation summary document

Copilot
and
DavidLanders95
committed
4ddd459

Commits on Feb 7, 2026

Checkpoint from VS Code for cloud agent session
DavidLanders95
committed
26bfc8e
Completely rewrite two_lenses_simplified.ipynb with JAX and Optuna

Copilot
and
DavidLanders95
committed
633f330
Add comprehensive documentation and improve optimization bounds

Copilot
and
DavidLanders95
committed
e6518ba
Update README_simplified.md with comprehensive documentation

Copilot
and
DavidLanders95
committed
b4fef35
Add comprehensive test suite for simplified notebook

Copilot
and
DavidLanders95
committed
9a23193
Add comprehensive summary document for notebook cleanup

Copilot
and
DavidLanders95
committed
bfe8f70
Replace Optuna with JAX BFGS and add full image generation pipeline

Copilot
and
DavidLanders95
committed
3799642
Update README to reflect BFGS optimization and image generation approach

Copilot
and
DavidLanders95
committed
beadb14
Add BFGS implementation summary document

Copilot
and
DavidLanders95
committed
05259df

Commits on Feb 9, 2026

Commits on Feb 10, 2026

Added electron magnetic components with real lens parameters
DavidLanders95
committed
a8aabe0
Added computation of new constants for lens functions
DavidLanders95
committed
acc03fb
Tried to jit evaluate_jax_scan, but it was a failure
DavidLanders95
committed
9135423
Added electromagnetic lens to gaussian also
DavidLanders95
committed
cf920e3
Added rotation matrix to transfer_matrices
DavidLanders95
committed
99e44ad
added test for constant component
DavidLanders95
committed
6b8f3b3
tried to get gaussians working for forward model, but no such luck
DavidLanders95
committed
5467ad8
Made fitting of larger As possible
DavidLanders95
committed
d52b375
improved lens_inversion
DavidLanders95
committed
8c99773

Commits on Feb 11, 2026

Can fit 5 lenses with no noise and regularisation
DavidLanders95
committed
810de99
Can work somewhat with noise.
DavidLanders95
committed
66db8c0
tried to simplify lens_inversion a little
DavidLanders95
committed
1340d87

Commits on Feb 12, 2026

Cleaned up lens inversion
DavidLanders95
committed
a768cc1
readded reference voltage
DavidLanders95
committed
aeb1a4a
tested tie method to measure phase at aperture plane and removed copilot slop
DavidLanders95
committed
2c45e48