Add simplified two-lens notebook with Collins FFT and automatic ABCD computation

[Copilot]

Simplifies two-lens optical system analysis by replacing manual ABCD matrix computation with automatic differentiation and adding Collins FFT diffraction modeling with Bayesian optimization.

Implementation

• two_lenses_simplified.ipynb: Streamlined workflow (28KB vs 616KB original)

  • Forward model using Lens + Detector components
  • ABCD transfer matrix via solve_model() with JAX autodiff
  • Collins FFT: H(f) = exp(-iπλ(B/A)(fx² + fy²)) using extracted ABCD parameters
  • Magnification-aware zooming with jax.image.resize
  • Optuna integration for parameter optimization

• Grid configuration:

  • Input: 5μm × 512px with 1μm circular aperture + padding
  • Output: 10mm × 256px detector plane

• Test suite: test_two_lenses_simplified.py validates all components

Example Usage

# Build system and extract ABCD matrix automatically
abcd = get_abcd_matrix(z1=0.0, z2=0.1, z3=0.5, f1=0.05, f2=0.15)
A_x, B_x = abcd[0, 0], abcd[0, 2]  # Magnification and defocus

# Propagate via Collins FFT
output_field = collins_fft_propagation(input_field, INPUT_SIZE, A_x, B_x, WAVELENGTH)

# Optimize with Optuna
study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=100)

Files Added

• examples/lens_inversion/two_lenses_simplified.ipynb
• examples/lens_inversion/test_two_lenses_simplified.py
• examples/lens_inversion/README_simplified.md
• examples/lens_inversion/IMPLEMENTATION_SUMMARY.md

Original prompt

Ok, what I would like you to work on while I am away is the following:

Redo the two lens notebook to simplify it.

Make a forward model using the 2 lenses in jaxgym, and use the detector component to make an input and output grid. The input grid should be larger than the 1 um aperture to make sure there are enough padding to avoid effects. The output grid should be 10mm x 10mm and also perhaps 256x256 in pixel sampling.

Differentiate through this forward model using a single input ray, and using jax.jacobian to get the 5x5 ABCD transfer matrix of the system.

Input the A and B values of this transfer matrix into the collins FFT model to generate a diffracted aperture.

Use jax.scipy to then zoom the solution onto the output grid. Then we can compare the two for a loss function.

Find a way to implement the beginning of the bayesian optimisation step with optuna.

Are there any clarifying questions?

Created from VS Code.

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.