Repository files navigation

Differentiable Optical Models asParameterised 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 asparametric neural networks.

∂Lux is built inZodiax, which is an open-source object-orientedJax framework built as an extension ofEquinox 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 enablesfast phase retrieval, image deconvolution, andhardware design in high dimensions. Because ∂Lux models are fully differentiable, you canoptimize them by gradient descent over millions of parameters; or useHamiltonian 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 inZodiax, which gives it a deep range of capabilities from bothJax andEquinox:

• 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 inZodiax, please go through thisZodiax 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 toemail me or contact me ontwitter! More details about contributing can be found in ourcontributing guide.

We have a multitude of publications in the pipeline using dLux, some built from our tutorials. To start we would recommend looking atthis invited talk on ∂Lux which gives a good overview and has an attached recording of it being presented! We also havethis poster!