Fresnel Propagation with abcdLux¤

# Basic imports
import jax.numpy as np
import jax.random as jr

# Optimisation imports
import zodiax as zdx
import optimistix as optx

# dLux imports
import dLux as dl
import dLux.utils as dlu

# Visualisation imports
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import CenteredNorm

%matplotlib inline
plt.rcParams['image.cmap'] = 'inferno'
plt.rcParams["font.family"] = "serif"
plt.rcParams["image.origin"] = 'lower'
plt.rcParams['figure.dpi'] = 120

# Nan friendly colormapping
inferno = mpl.colormaps["inferno"]
seismic = mpl.colormaps["seismic"]
inferno.set_bad("k", 0.5)
seismic.set_bad("k", 0.5)

# Define our properties
wf_npix = 512  # Number of pixels in the wavefront
diameter = 6.6  # Aperture diameter
focal_length = 130  # Effective focal length of the internal optics
defocus = 30e-3  # Defocus of the focal plane

# Propagate to the science plane using an MFT Fresnel propagator
psf_npix = 256  # Number of pixels in the PSF
psf_pixel_scale = 10e-6  # 10 microns

# Propagate to focus using an MFT Fresnel propagator
to_focal = dl.MFTPropagator(
    [
        ("ThinLens", dl.ABCDConjugatePlane(focal_length)),  # Pupil -> Focal
        ("FreeSpace", dl.ABCDFreeSpace(-defocus)),  # Focal -> Defocused Detector
    ],
    dl.CoordSpec(n=psf_npix, d=psf_pixel_scale),
)

# Generate other layers
aper = dlu.jwst_like(npixels=wf_npix, oversample=5)
aperture = dl.TransmissiveLayer(aper, normalise=True)
aberrations = dl.FourierBasis(wf_npix, n_modes=128)

# Define the optical layers
layers = [
    ("aperture", aperture),
    ("aberrations", aberrations),
    ("to_focal", to_focal),
]

# Construct the optics object
optics = dl.LayeredOpticalSystem(
    wf_npixels=wf_npix,
    diameter=diameter,
    layers=layers,
)

# Examine the optics object
print(optics)

# Define some random coefficients for the Fourier basis
coeffs = 2.0e-9 * jr.normal(jr.key(0), optics.coefficients.shape)
optics = optics.set(coefficients=coeffs)

# Model the positive and negative defocused PSFs
psf_pos = optics.propagate(1e-6)
psf_neg = optics.multiply(distance=-1).propagate(1e-6)

# Generate some data
flux = 1e6
read_std = 10
n_ints = 1000
ints = flux * np.ones(n_ints)[:, None, None]

#
keys = jr.split(jr.key(0), 3)
photons_pos = jr.poisson(keys[0], ints * psf_pos[None, :, :])
photons_neg = jr.poisson(keys[1], ints * psf_neg[None, :, :])
read_noise = read_std * jr.normal(keys[2], (2, *photons_pos.shape))

#
electrons_pos = photons_pos + read_noise[0]
electrons_neg = photons_neg + read_noise[1]

#
mean_pos = np.mean(electrons_pos, axis=0)
mean_neg = np.mean(electrons_neg, axis=0)
err_pos = np.sqrt(np.var(electrons_pos, axis=0) / n_ints)
err_neg = np.sqrt(np.var(electrons_neg, axis=0) / n_ints)

#
data = np.array([mean_pos, mean_neg])
error = np.array([err_pos, err_neg])

def eval_wfs_psfs(optics):
    psf_pos = optics.propagate(1e-6)
    psf_neg = optics.multiply(distance=-1).propagate(1e-6)
    return np.array([psf_pos, psf_neg])

def eval_wfs_data(params, optics):
    optics = optics.set(coefficients=1e-9 * params["coeffs"])
    return 1e6 * eval_wfs_psfs(optics)

def z_score_fn(params, optics, data, error):
    pred = eval_wfs_data(params, optics)
    return zdx.z_score(pred, data, error)

opt_fn = lambda params, args: np.mean(z_score_fn(params, *args)**2)

# Re-initialise our guess of parameters
keys = jr.split(jr.key(1), 3)
params = {"coeffs": np.zeros_like(optics.coefficients)}

# Apply the optimiser
args = (optics, data, error)
solver = optx.BestSoFarMinimiser(optx.LBFGS(rtol=1e-3, atol=1e-3))
sol = optx.minimise(opt_fn, solver, params, args, max_steps=512, throw=False)

# Check the final reduced chi-squared statistic
chi2r = zdx.chi2r(eval_wfs_data(sol.value, optics), data, error, zdx.ddof(params, data))
print("Final reduced chi-squared:", chi2r)
print("Steps:", int(sol.stats["num_steps"]))