Mutable vs Immutable

[Jordan-Dennis]

Hi all,
I have been doing some fiddling this morning and despite the fact that equinox and jax enforce immutability, it is possible for us to implement a mutable interface. This is easy and instead of returning the modified state via equinox.tree_at we re-assign it to self. There is no computational benefit to doing this, but I thought it was worth mentioning that it was possible. Let me know what you think.

[Jordan-Dennis]

So interestingly this is what update_phasor does, modifying the initial object and returning the pointer to its location in RAM. This is a nice syntax because it means that one can chain together transformations like java streams or r pipelines. For example:

wavefront.update_phasor()
    .do_something()
    .do_another_thing()
    .do_more_things()
    .do_final_thing()

While we can't do this at the moment except for when using update_phasor it is an interesting concept and means that our OpticalSystem could essentially reduce to:

import dLux
wavefront = dLux.GaussianWavefront(parameters)

def gradable_with_respect_to_everything(wavefront, *other_required_params):
    return wavefront.apply_optic(dLux.CircularApperture(radius))
    .apply_propagator(dLux.FresnelDiffractor(distance))
    .apply_optic(dLux.HexagonalMirror(required_params))
    .apply_propagator(dLux.FraunhoferDiffractor(distance))
    .apply_optic(dLux.Coronagraph(hope_I_spelled_it_correct))
    .apply_propagator(dLux.FresnelDiffractor(distance))
    .apply_detector(dLux.JitteredDetector(params))

Now obviously this syntax is almost the exact reverse of our current syntax, but I thought it was worth mentioning as a toy idea.

[benjaminpope]

Neat stuff. I am a bit tempted to say that this isn't very equinoxy - we would want to be careful so that our code resembles ordinary equinox modules so that it plays nicely with other NNs.

I would also say that a stack of methods like above is probably how I would have written this a year ago before equinox came out, to make a really nice functional Jax code, but we probably want to instead do all of the stacking of layers in an equinox __init__ like normal and take this out of the user's hands at the point of compute.

[Jordan-Dennis]

You were correct, this isn't very equinoxy. @LouisDesdoigts and I discussed this in our meeting on Friday and retested and it didn't work. Not sure what I did the first time I tested it, probably initialised it with the same value I changed it to. I am still not a fan of the stacked re-assignments, e.g.

wavefront = wavefront.update_phasor()
wavefront = wavefront.add_phase()
ect.

but I guess it cannot be helped.