LouisDesdoigts / dLux

Differentiable optical models as parameterised neural networks in Jax using Zodiax
https://louisdesdoigts.github.io/dLux/
BSD 3-Clause "New" or "Revised" License
49 stars 6 forks source link

Models #226

Closed LouisDesdoigts closed 1 year ago

LouisDesdoigts commented 1 year ago

Major Update!


This update is a major overhaul of the backed of dLux in preparation of submission to JOSS, and the 1.0 release! This update expands the capabilities and flexibility of the software, all while removing around 5'000 line of code. Lets go through what these changes entail.

High Level

This update separates all of the classes from core.py into their own individual scripts, that means we now have quite a new scripts: optical_layers.py, instruments.py, images.py, detector_layers.py, along with some renamed scripts:optics.py->optical_layers.py,detectors.py->detector_layers.py,spectrums.py->spectra.py`.

The primary aim of these changes is to have a more clear and logical division of the different classes in order to make building your own dLux classes easier. The resulting package is split into a series of different types of classes, lets take a look at them:

Wavefronts and classes that modify them

There are three main types of classes: Wavefronts (wavefronts.py) which represent the state of some monochromatic wavefront. OpticalLayers perform transformations on Wavefronts and Optics (optics.py) classes which hold a series of OpticalLayers in order to model some optical system.

The OpticalLayers classes are split up into four different scripts:

Images and classes that modify them

The dLux module also contains a series of classes that modify Images (images.py) which represent the state of some psf as it is transformed through a detector. The structure matches that of the Wavefront classes, with DetectorLayers (detector_layers.py) performing transformations on Images and Detectors (detectors.py) holding a series of DetectorLayers in order to model some detector.

Sources and Spectra

The dLux module also contains a series of classes that represent sources and their spectra. The Source classes (sources.py) represent some parametric source, while the Spectrum classes (spectra.py) represent the spectrum of the source.

Instruments and Observations

The Instrument (instruments.py) is designed to coherently model the interaction between these different components and the Observation (observations.py) classes allow for fine-grained control over the modelling of the Instrument class. An example of this is to allow for dithering patterns modeled, or for observing the same source using different instrumental filters.

Utils

This simply provides a series of useful functions that are used throughout the rest of the software.


Wavefronts

Now that we have that out of the way, lets take a look at the major changes. First and foremost is a re-think of how the Wavefronts classes should be interacted with. Previously, these were considered 'private' and generally shouldn't be interacted with by the user. Now this class has been re-worked and is now part of the 'public' API in order to allow users to create their own fast and efficient instrument-specific classes.

OpticalLayers

This has also facilitated and change in the way that OpticalLayers behave, made to be more flexible and general. While the Apertures, Aberrations and Propagators remain largely unchanged (beyond merging the classes into a smaller number of more flexible classes), the 'base' OpticalLayers are restructured into four primary classes: Optic, PhaseOptic, BasisOptic, PhaseBasisOptic. The Optic and PhaseOptic classes both hold an array of transmission values, a boolean normalise value, and either an opd or phase array. The BasisOptic and PhaseBasisOptic similarly have a transmission and normalise attribute, but instead of the opd or phase they have basis and coefficients. This allows for a set of basis vectors and coefficients to be specific that are generated dynamically at run time for things like optical aberrations.

Optics

The two above changes has facilitated a re-think in the way we specify optical system in dLux. Previously we had the single Optics class, which has now been re-named to LayeredOptics, which behaves similarly in that it takes in a list of layers, with the difference that we now no longer need to specify the npixels and diameter properties of the wavefront via the CreateWavefront layer (deprecated), as those two parameters are now properties of the LayeredOptics classes themselves.

There are also now three new classes that provide a simplified interface when working with basic single conjugate plane optical systems: AngularOptics, CartesianOptics and FlexibleOptics. Please refer to the updated docs for details on these classes.

Detectors, DetectorLayers and Images

All the classes that relate to the interaction with detectors have been re-worked to now have a symmetrical class structure to Optics, OpticalLayers and Wavefronts. However, the functionality is basically the same so we will not go into detail here.

Utils

The utils functions have been simplified, with some of the code being moved into external repos as they are not specific to dLux.


Building your own dLux class

dLux is designed to be a framework for differentiable optical modelling, in pursuit of this goal each class type now has a Base version, ie BaseOptics, BaseInstrument etc that can be used to create classes that will integrate seamlessly into the rest of the dLux framework. That means you create your own custom optics class that does whatever you want and all dLux sources, instruments etc will recognise and work with these classes! For details on this please refer to the API/Optics section of the docs.


These are the primary changes in this version, however there are lots of small back-end tweaks in terms of where bits of code live and how they interact with each other. I will not specify all of these changes here are most of them will not be relevant to users, and this merge is currently +10k -15k lines of code, so specificity each individual change is kind of beyond the pale at this point. This version can however be considered 'stable' as the package is now very flexible and modular and soon to be released as the 1.0 version!