search

chromatix-team / chromatix

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

Refactoring of phasemasks #102

Closed GJBoth closed 1 year ago

GJBoth commented 1 year ago

This PR modifies the phasemasks to take a grid and pupil radius as input, rather than shape / spacing or the Field. This makes things a lot simpler. Main reason for not having the Field as input is that that would mix the abstractions; right now every function which takes in a field returns a field, and doing so would break this.

This PR also makes the distinction between functional and other things stronger: functional only contains functions which take in Fields and return Fields.

To achieve this I've moved the seidel and zernike aberrations to the utils (plus small functions like wrap_phase. The name and ordering for now is fluid, will fix this up in the next PR which will clean up utils and ops.

Essentially, this PR is a start to the following structure

  1. Functional - only and only functions which transform Fields (and sources).
  2. Elements - Functional, but in a stateful class to make them into layers.
  3. Systems - Microscopes, 4f systems, etc.
  4. Ops and Utils - to be cleaned, but more or less everything else.
diptodip commented 1 year ago

Closing based on our discussion: we'll keep the structural change (initializers.py) but not the changes to the phase mask initializer function signatures.