search

chromatix-team / chromatix

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

nn.compact vs setup for holoscope #115

Open schneidermc opened 6 months ago

schneidermc commented 6 months ago

In the holoscope example, trying to get the psf of the Holoscope class via psf = holoscope.apply({"params": params, "state": state}, 0, method=holoscope.psf, rngs={"noise": key}).intensity.squeeze() gives the following error: AssignSubModuleError: Submodule ObjectivePointSource must be defined in 'setup()' or in a method wrapped in '@compact'

diptodip commented 6 months ago

Can you try with method=Holoscope.psf? But either way the Holoscope class in that example could definitely be improved if we make the psf method just use functions instead of elements. I'll get on that.