Basic refractive optic class

[LouisDesdoigts]

class RefractiveOptic(dl.layers.optical_layers.OpticalLayer):
    """
    A refractive optical element that applies a wavelength-dependent phase
    based on a thickness map and a refractive index dispersion curve.
    """

    path_length: jnp.ndarray
    r_index: ipx.Interpolator1D

    def __init__(self, path_length, wavels, r_index, kind="linear", extrap=jnp.nan):
        super().__init__()

        self.path_length = jnp.array(path_length, float)
        self.r_index = ipx.Interpolator1D(
            x=jnp.array(wavels, float),
            f=jnp.array(r_index, float),
            kind=kind,
            extrap=extrap,
        )

    def apply(self, wavefront):
        """Apply wavelength-dependent phase to the input wavefront."""
        r_index = self.r_index(wavefront.wavelength) - 1.0
        return wavefront.add_opd(self.path_length * r_index)```