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
43 stars 6 forks source link

Investigate the creation of `UnifiedLayers` that can operate on both `Wavefronts` and `Images` #240

Closed LouisDesdoigts closed 9 months ago

LouisDesdoigts commented 11 months ago

An example would be a unified Rotate layer that is a subclass of both OpticalLayer and DetectorLayer for type checking. ie


class Rotate(dLux.optics.OpticalLayer, dLux.detector.DetectorLayer):
    angle: float
    order: int

     def rotate_wavefront(self, wavefront):
        ...

     def rotate_image(self, image):
        ...

    def __call__(self, input):
        if isinstance(input, dLux.wavefronts.Wavefront):
            return self.rotate_wavefront(input)
        else:
            return self.rotate_image(input)

Other candidate layers are Flip and Resize.

Some details of implementation are probably sub-optimal, such as requiring the complex parameter for the RotateWavefront but not for the RotateDetector. The benefit would be a unified and simplified API, which is reasonably desirable.

Related to #239

LouisDesdoigts commented 9 months ago

Fixed with #246