Replies: 4 comments
-
Sounds cool, on the implementation side it can be quite tricky to pass functions dynamically in class ChromaticPhasePlate():
# A class with the desired attributes and behaviours
def apply_chromatic_dependence(chromatic_dependence: Callable) -> ndarray:
# Implementation fitting the desired spec
def chromatic_dependence(*args, **kwargs):
# Some implementation of the desired behaviour with the correct parameters to work
# with ChromaticPhasePlate
# This will through an error if we try and jit through it because Callable objects
# are not Hashable and so cannot be marked as static.
ChromaticPhasePlate().apply_chromatic_dependence(chromatic_dependence(*args, **args)) This just means that it is easiest to use an class ChromaticPhasePlate(equinox.Module):
chromatic_dependence: Callable # Marked as static automatically
phases: ndarray
def __init__(self, chromatic_dependence: Callable, phases: ndarray):
self.chromatic_dependence = chromatic_dependence
self.phases = phases
def __call__(self, phasefront: ndarray, wavelength: float) -> ndarray:
return phasefront + self.chromatic_dependence(wavelength, phases) where # The phase-front (called image) that you are acting on is assumed to exist
# The phases array defining how the wave-plate changes the phase-front is
# assumed to exist and satisfy image.shape == phases.shape
# wavelength is assumed to be defined.
def chromatic_dependence(wavelength, phases):
# Desired implementation
image = ChromaticPhasePlate(chromatic_dependence, phases)(phasefront, wavelength) where we have substantiated and called in one line. |
Beta Was this translation helpful? Give feedback.
-
Another option is to use interpolation over a grid rather than a function (credit to @LouisDesdoigts for rewriting into vmap) class PolychromaticPhaseScreen(eqx.Module):
"""
"""
npix: int = eqx.static_field()
wavels: np.ndarray
phase_grid: np.ndarray
def __init__(self, wavels, phase_grid):
"""
wavels : the wavelengths the grid it sampled at
phase_grid : (npix, npix, wavels.size) the phase values to add
"""
self.wavels = wavels
self.phase_grid = phase_grid
self.npix = int(phase_grid.shape[0])
def __call__(self, params_dict):
"""
"""
# Get relevant parameters
WF = params_dict["Wavefront"]
mapped_interp = jax.vmap(np.interp, in_axes=(None, None, 0))
flat_phases = self.phase_grid.reshape([self.npix**2, len(self.wavels)])
phase_screen = mapped_interp(WF.wavel, self.wavels, flat_phases)
phase_screen = phase_screen.reshape([self.npix, self.npix])
WF = WF.add_phase(phase_screen)
params_dict["Wavefront"] = WF
return params_dict
|
Beta Was this translation helpful? Give feedback.
-
Just to build on what Peter said, I think the chromaticity is rarely going to be arbitrary, but would in general just be a fairly smooth refractive index function which would rescale the same shape phase array. |
Beta Was this translation helpful? Give feedback.
-
Agreed - adding some 'lens' class that is parameterised by a refractive index should be trivial to add. A more general class that can be passed a functional wavelength-phase relationship should able to possible, albeit slightly less trivial. That said having functionality shown by @ataras2 is already very powerful as it allows us to use phase-screens produced by various different libraries, just not optimise the underlying parameters. |
Beta Was this translation helpful? Give feedback.
-
I think it would be cool to add a feature that allows a chromatic phase screen.
That is, a phase screen that allows for some (more complicated) wavelength dependence.
Feed in some kind of function describing the refractive index .vs. Wavelength, and then the phase screen (could be a lens, prism for example) sill adjust the wavefront according to the specific delays impoased at that pixel (depends on the optical system), and at that wavelength (function of material properties).
Beta Was this translation helpful? Give feedback.
All reactions