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

Wavefronts with Variable Pixel Scales. #204

Closed Jordan-Dennis closed 1 year ago

Jordan-Dennis commented 1 year ago

Hi all, The (not so) new property based implementation of the Wavefront class is great but I have encountered a small inconvinience trying to upgrade the GaussianWavefront to match. The problem is pixel_positions accesses pixel_scale as a parameter. Via the inheritance mechanisms of equinox pixel_scale must be kept as a parameter so I cannot, transfer it to a property. This implies that I will require a get_pixel_scale method and then re-implement pixel_positions accessing the pixel_scale parameter via the get_pixel_scale method. This seems like a fair amount of extra work. It may be possible to hack a solution using a lambda function closure, but I am not sure if this is a good idea. Regards Jordan.

Jordan-Dennis commented 1 year ago

Here is a minimal working example of the problem and proposed hack, which does not work,

import equinox as eqx
class A(eqx.Module):
    a: int
    def __init__(self, a: int) -> object:
        self.a = a

class B(A):
    _a: int
    _b: int
    _c: bool 
    def __init__(self, _a: int, _b: int, _c: bool) -> object:
        self._a: int = _a
        self._b: int = _b
        self._c: bool = _c
        self.a = property(lambda: self._a if self._c else self._b)

B(1, 2, True).a # A property object
B(1, 2, True).a.fget # The lambda function (needs to be evaluated).

This shows that the hack is not available. For now I will simply override the required methods and add the additional getter. Regards Jordan

Jordan-Dennis commented 1 year ago

This lead me to have an interesting thought. Using meta-classes it should be possible to make equinox automatically turn variables into properties (the setter using tree_at). This would give us back the traditional python syntax class.parameter = new_value. However, because of immutability you would end up with variable: object = class.parameter = new_value, which would probably cause a lot of headaches.

LouisDesdoigts commented 1 year ago

Sorry I don't really follow, why do you need to access a modified versoin of the pixel scale for the GaussianWavefront?

Jordan-Dennis commented 1 year ago

The pixel_scale changes if the wavefront is angular or not. It is an implementation detail really ...

LouisDesdoigts commented 1 year ago

I mean it only changes the units. Can't we just have a BaseGaussianWavefront and then and CartesianGaussianWavefront and AngularGaussianWavefront that implements the function to output in the right units? This would match the way Wavefronts behaves current

Jordan-Dennis commented 1 year ago

I do not believe that would work since the Wavefront becomes angular within the waist of the beam (if I am recalling correctly).

LouisDesdoigts commented 1 year ago

Closing this for now, until we try to re-implement this functionality.