Differentiable Optical Models as Parameterised Neural Networks in Jax using Zodiax
Contributors: Louis Desdoigts, Jordan Dennis, Adam Taras, Max Charles, Connor Langford, Benjamin Pope, Peter Tuthill
∂Lux is an open-source differentiable optical modelling framework harnessing the structural isomorphism between optical systems and neural networks, giving forwards models of optical systems as parametric neural networks.
∂Lux is built in Zodiax, which is an open-source object-oriented Jax framework built as an extension of Equinox for scientific programming. This framework allows for the creation of complex optical systems involving many planes, phase and amplitude screens in each, and propagates between them in the Fraunhofer or Fresnel regimes. This enables fast phase retrieval, image deconvolution, and hardware design in high dimensions. Because ∂Lux models are fully differentiable, you can optimize them by gradient descent over millions of parameters; or use Hamiltonian Monte Carlo to accelerate MCMC sampling. Our code is fully open-source under a 3-clause BSD license, and we encourage you to use it and build on it to solve problems in astronomy and beyond.
The ∂Lux framework is built in Zodiax, which gives it a deep range of capabilities from both Jax and Equinox:
Accelerated Numpy: a Numpy-like API that can run on GPU and TPU
Automatic Differentiation: Allows for optimisation and inference in extremely high-dimensional spaces
Just-In-Time Compilation: Compiles code into XLA at runtime and optimising execution across hardware
Automatic Vectorisation: Allows for simple parallelism across hardware and asynchronous execution
For an overview of these capabilities and different optimisation methods in Zodiax, please go through this Zodiax Tutorial.
Documentation: https://louisdesdoigts.github.io/dLux/
Requires: Python 3.10+, Jax 0.4.13+, Zodiax 0.4+
Installation: pip install dLux
If you want to run the tutorials locally, you can install the 'extra' dependencies like so: pip install 'dLux[extras]'
We are always looking to collaborate and further develop this software! We have focused on flexibility and ease of development, so if you have a project you want to use ∂Lux for, but it currently does not have the required capabilities, have general questions, thoughts or ideas, don't hesitate to email me or contact me on twitter! More details about contributing can be found in our contributing guide.
We have a multitude of publications in the pipeline using dLux, some built from our tutorials. To start we would recommend looking at this invited talk on ∂Lux which gives a good overview and has an attached recording of it being presented! We also have this poster!