Coronagraphs & Multi-Plane Optical Systems¤

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

# # Optimisation imports
import equinox as eqx
import zodiax as zdx
import optax

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

# Visualisation imports
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt
import matplotlib as mpl

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

# 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 wavefront properties
wf_npix = 256  # Number of pixels in the wavefront
diameter = 1.0  # Diameter of the wavefront, meters

# Define our detector properties
psf_npix = 64  # Number of pixels in the PSF
psf_pixel_scale = 50e-3  # 50 mili-arcseconds

# Generate a euclid-like aperture and Zernike
m2_diam = 0.2 # Secondary diameter
spider_width = 0.05 # Spider width
aper, basis = dlu.euclid_like(
    npixels=wf_npix,
    diameter=diameter,
    secondary_diameter=m2_diam,
    spider_width=spider_width,
    zernike_nolls=np.arange(4, 22),
)

# Normalise the basis to nm and generate some random coefficients
basis *= 1e-9
coeffs = 25 * jr.normal(jr.key(0), basis.shape[0])
lamd = 1e-6 / diameter

# Define the optical layers
layers = [
    ("aperture", dl.TransmissiveLayer(aper, normalise=True)),
    ("aberrations", dl.BasisLayer(basis, coeffs)),
    ("to_focal", dl.FFT(pad=4)),
    ("fpm", dl.Lambda()),
    ("from_focal", dl.FFT(crop=4)),
    ("lyot", dl.Lambda()),
]

# Construct the optics object
optics = dl.AngularOpticalSystem(
    wf_npixels=wf_npix,
    diameter=diameter,
    layers=layers,
    psf_npixels=psf_npix,
    psf_pixel_scale=psf_pixel_scale,
)

# Examine the optics object
print(optics)

# Propagate a wavefront through the system and get the intermediate wavefronts
wf, outputs = optics.debug_propagate_mono(1e-6)

# Create a child-class of OpticalLayer, so dLux can recognise it as a layer
class FocalPlaneMask(dl.OpticalLayer):
    """Applies a dynamic circular focal plane mask to the wavefront in FFT coordinates."""
    radius: float  # Focal plane mask radius in lambda / D units
    shift: np.ndarray  # Shift of the focal plane mask, in lambda / D units
    scale: float # Size of lambda / D in radians

    def __init__(self, radius=2.0, scale=1., shift=None):
        self.radius = np.array(radius, float)
        self.scale = np.array(scale, float)
        if shift is None:
            shift = np.zeros(2)
        self.shift = np.array(shift, float)

    def __call__(self, wavefront):
        """
        Applies the focal plane mask to the wavefront. Assumes the wavefront is in 
        FFT coordinates.
        """
        # Get the wavefront coordinates and shift them
        fp_coords = wavefront.coordinates()
        fp_coords = dlu.translate_coords(fp_coords, -self.shift * self.scale)

        # Calculate the focal plane mask with a soft-edge and apply it
        softness = 0.5 * wavefront.pixel_scale
        radius_phys = self.radius * self.scale
        fp_mask = dlu.soft_circle(fp_coords, radius_phys, softness, invert=True)
        wavefront *= fp_mask

        # Return the wavefront with the focal plane mask applied
        return wavefront

# Generate the focal plane mask, targeting a wavelength of 1 micron
fpm = FocalPlaneMask(radius=2, shift=[0, 0], scale=1e-6 / diameter)

# Examine
print(fpm)

# Create a child-class of OpticalLayer, so dLux can recognise it as a layer
class LyotMask(dl.OpticalLayer):
    """Applies a dynamic circular Lyot mask to the wavefront in FFT coordinates."""
    undersize: float  # Lyot mask undersize, ie 0.9 = 90% of the pupil radius
    shift: np.ndarray  # Shift of the Lyot mask, in physical units (i.e. meters)
    aperture: np.ndarray  # The static aperture to apply in the Lyot plane, including spiders

    def __init__(self, undersize=1., shift=None, aperture=None):
        self.undersize = np.array(undersize, float)
        if shift is None:
            shift = np.zeros(2)
        self.shift = np.array(shift, float)
        if aperture is None:
            aperture = np.array(1.) # empty mask if not provided
        self.aperture = np.array(aperture, float)

    def __call__(self, wavefront):
        """Applies the Lyot stop to the wavefront."""

        # Get the wavefront coordinates and shift them,
        coords = wavefront.coordinates()
        coords = dlu.translate_coords(coords, -self.shift)

        # Calculate the Lyot mask with a soft-edge and apply it
        softness = 0.5 * wavefront.pixel_scale
        radius_phys = wavefront.diameter * self.undersize / 2
        lyot_stop = dlu.soft_circle(coords, radius_phys, softness)
        wavefront *= lyot_stop * self.aperture

        # Return the wavefront with the Lyot mask applied
        return wavefront

lyot = LyotMask(undersize=0.9, shift=[0, 0], aperture=aper[::-1, ::-1])

# Examine
print(lyot)

# Insert our optical layers
coron = optics.set(fpm=fpm)
coron = coron.set(lyot=lyot)

# Examine the optics object again to see the new layers
coron

# Remove the wavefront error
coron = coron.multiply(coefficients=0.)

# Propagate a wavefront through the system and get the intermediate wavefronts
wf, outputs = coron.debug_propagate_mono(1e-6)

wavels = 1e-6 * np.linspace(0.9, 1., 3)
lamd = wavels.mean() / diameter
offsets = np.linspace(0, 4, 5)

psfs = []
for oset in offsets:
    offset = lamd * np.array([-oset, 0])
    psf = coron.propagate(wavels, offset)
    psfs.append(psf)
psfs = np.array(psfs)

# Focal plane coordinates
coords = dlu.pixel_coords(5 * optics.psf_npixels, optics.fov)
coords_lamd = dlu.arcsec2rad(coords) / lamd

# Annular mask to suppress starlight within
inner = dlu.circle(coords_lamd, 2, invert=True).astype(bool)
outer = dlu.circle(coords_lamd, 5).astype(bool)
annular_mask = dlu.downsample(outer & inner, 5)

# Companion PSF samples
positions = lamd * np.array([[2.5, 0], [3.5, 0], [4.5, 0]])

# Define the loss function
@eqx.filter_jit
@eqx.filter_value_and_grad
def loss_fn(params, coron):
    """Simple Loss function"""
    # fpm size, lyot size
    coron = coron.set(params)

    # Get the primary PSF throughput in the annular region
    primary = coron.propagate(wavels)
    primary_throughput = np.sum(annular_mask * primary)

    # Propagate companion PSFs
    companions = vmap(lambda pos: coron.propagate(wavels, pos))(positions)
    companion_throughput = companions.sum(axis=(1, 2)).mean(0)

    # Loss metric
    return np.log10(primary_throughput) - np.log10(companion_throughput)

# Set the initial parameters
params = {
    "fpm.radius": np.array(1.5),
    "lyot.undersize": np.array(0.95),
}

# Set up the optimisers for each parameter
optimisers = {
    "fpm.radius": optax.sgd(5e-2, momentum=0.5),
    "lyot.undersize": optax.sgd(1e-2, momentum=0.5),
}
optim, state = zdx.map_optimisers(params, optimisers)

# Now we can run a simple optimisation loop
epochs = 25
pbar = tqdm(range(epochs))
losses, params_history = [], []
for step in pbar:
    loss, grads = loss_fn(params, coron)
    updates, state = optim.update(grads, state)
    params = optax.apply_updates(params, updates)
    pbar.set_postfix(log_loss=f"{loss:.4f}")
    losses.append(loss)
    params_history.append(params)
losses = np.array(losses)

# Propagate some wavelengths thought the system
offsets = np.linspace(0, 4, 5)

psfs = []
for oset in offsets:
    offset = lamd * np.array([-oset, 0])
    psf = coron.set(params).propagate(wavels, offset)
    psfs.append(psf)
psfs = np.array(psfs)

class SoummerFPM(dl.OpticalLayer):
    """
    Implements the Soummer et al. (2007) algorithm to apply a high-res focal plane
    mask without large FFTs.
    """
    lamd: float  # The size of the an optical fringe
    sigma: float  # The standard deviation of the Gaussian mask in lambda / D units
    power: float  # The power to raise the Gaussian mask to
    npixels: int  # The number of pixels in the focal plane mask
    pixel_scale: float  # The pixel scale of the focal plane mask, in lamda / D units

    def __init__(self, npixels, pixel_scale, lamd, sigma=1., power=5.):
        self.npixels = int(npixels)
        self.pixel_scale = np.array(pixel_scale, float)
        self.lamd = np.array(lamd, float)
        self.sigma = np.array(sigma, float)
        self.power = np.array(power, float)

    @property
    def mask(self):
        """Makes the gaussian mask"""
        radius = self.npixels * self.pixel_scale / (2 * self.lamd)
        mask = dlu.gaussian((0., 0.), self.sigma, self.npixels, extent=radius)
        return (1 - mask / mask.max()) ** self.power

    def __call__(self, wavefront):
        """Applies the focal plane mask to the wavefront."""
        # MFT propagate to the focal plane and apply the mask
        npixels = self.mask.shape[0]
        fp_wf = wavefront.propagate(npixels, self.pixel_scale)
        fp_wf = fp_wf.multiply(phasor=1 - self.mask)

        # Inverse MFT propagate back to the pupil plane
        npixels = wavefront.npixels
        pixel_scale = wavefront.pixel_scale
        pup_wf = fp_wf.propagate(npixels, pixel_scale, inverse=True)

        # Subtract wavefront from the original the get the applied mask
        return wavefront - pup_wf

# Construct our Soummer Focal plane mask layer
fpm = SoummerFPM(npixels=128, pixel_scale=lamd/5, lamd=lamd, sigma=1.)

# Examine the new layer
print(fpm)

# Define the optical layers
layers = [
    ("aperture", dl.TransmissiveLayer(aper, normalise=True)),
    ("aberrations", dl.BasisLayer(basis)),
    ("fpm", fpm),
    ("lyot", LyotMask(undersize=0.9, aperture=aper)),
]

# Construct the optics object
coron = dl.AngularOpticalSystem(
    wf_npixels=wf_npix,
    diameter=diameter,
    layers=layers,
    psf_npixels=psf_npix,
    psf_pixel_scale=psf_pixel_scale,
)

# Examine the optics object again to see the new layers
print(coron)

wf, outputs = coron.debug_propagate_mono(1e-6)

# Propagate some wavelengths thought the system
offsets = np.linspace(0, 4, 5)

psfs = []
for oset in offsets:
    offset = lamd * np.array([oset, 0])
    psf = coron.propagate(wavels, offset)
    psfs.append(psf)
psfs = np.array(psfs)

# Set the initial parameters
params = {
    "fpm.sigma": np.array(1.5),
    "lyot.undersize": np.array(0.9),
}

# Set up the optimisers for each parameter
optimisers = {
    "fpm.sigma": optax.sgd(5e-2, momentum=0.5),
    "lyot.undersize": optax.sgd(2e-3, momentum=0.5),
}
optim, state = zdx.map_optimisers(params, optimisers)

# Now we can run a simple optimisation loop
epochs = 100
pbar = tqdm(range(epochs))
losses, params_history = [], []
for step in pbar:
    loss, grads = loss_fn(params, coron)
    updates, state = optim.update(grads, state)
    params = optax.apply_updates(params, updates)
    pbar.set_postfix(log_loss=f"{loss:.4f}")
    losses.append(loss)
    params_history.append(params)
losses = np.array(losses)

# Propagate some wavelengths thought the system
offsets = np.linspace(0, 4, 5)

psfs = []
for oset in offsets:
    offset = lamd * np.array([oset, 0])
    psf = coron.set(params).propagate(wavels, offset)
    psfs.append(psf)
psfs = np.array(psfs)