Propagation¤

def transfer_fn(coords, npixels, wavelength, pscale, distance):
    """Evaluate the active Fourier-domain transfer kernel used for the current propagation model."""
    del npixels, pscale
    rho_sq = (coords**2).sum(0)
    return _fftshift(jnp.exp(-1.0j * jnp.pi * wavelength * distance * rho_sq))

def transfer(wf, distance, pad=2):
    """Build a transfer kernel for propagation from one plane to another."""
    npix = pad * wf.npixels
    diam = pad * wf.diameter
    freqs = jnp.fft.fftshift(jnp.fft.fftfreq(npix, diam / npix))
    coords = jnp.array(jnp.meshgrid(freqs, freqs))
    return transfer_fn(
        coords, wf.npixels, wf.wavelength, pad * wf.pixel_scale, distance
    )

def plane_to_plane(wf, distance, pad=2):
    """Propagate a wavefront between planes using the current transfer kernel."""
    fft_wf = _fft(wf.phasor, pad=pad)
    tf = transfer(wf, distance, pad=pad)
    phasor = dlu.resize(_ifft(fft_wf * tf), wf.npixels)
    return wf.set(["amplitude", "phase"], [jnp.abs(phasor), jnp.angle(phasor)])