chromatix-team / chromatix

Differentiable wave optics using JAX! Documentation can be found at https://chromatix.readthedocs.io
MIT License
75 stars 13 forks source link

Grid for phase masks etc. #30

Closed schneidermc closed 5 months ago

schneidermc commented 1 year ago

Each function for each different kind of phase mask calculates its own grid. The grid should rather be passed to the function, e.g. from the Lightfield that is put in. That would remove duplicated code and make it easier to write a new function / phase mask.

diptodip commented 1 year ago

I think the simpler change is to have the grid in Field use the grid utility. Right now, the PhaseMask elements can use either jax.nn.initializers functions or our phase mask initialization functions. If we always ask for Fields, we lose that and it becomes a bit cumbersome to tweak the phase masks interactively.