Invert a Scene Blurred by a Different PSF in Each Channel

In this tutorial, we’ll image the gaussians() sample scene with the instrument from Invert a Synthetic Scene with MART, except that each of its four channels blurs the scene by a point-spread function (PSF) of its own. We’ll then invert the images with MART twice: once with an instrument model which includes the PSFs, and once with one which ignores them.

The PSF is an attribute of the instrument, psf. The instrument folds it into its weights with named_arrays.regridding.convolve_weights(), so image() blurs the scene by it, and backproject(), the transpose of the forward model, accounts for it as well. MART uses only those two methods, so it needs nothing else to take the PSFs into account.

import warnings
import dataclasses
import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u
import astropy.visualization
import named_arrays as na
import ctis

Define the same grids and scene as in Invert a Synthetic Scene with MART: a grid of Doppler velocities and positions on the sky, a grid of pixels on the sensor, and the synthetic scene with a background of 1 percent of its maximum.

velocity = na.linspace(-500, 500, axis="wavelength", num=21) * u.km / u.s
wavelength_rest = 171 * u.AA

position_scene = na.Cartesian2dVectorLinearSpace(
    start=-10 * u.arcsec,
    stop=10 * u.arcsec,
    axis=na.Cartesian2dVectorArray("scene_x", "scene_y"),
    num=na.Cartesian2dVectorArray(64 + 1, 64 + 1),
)
position_sensor = na.Cartesian2dVectorArray(
    x=na.arange(0, 128 + 1, axis="sensor_x") * u.pix,
    y=na.arange(0, 64 + 1, axis="sensor_y") * u.pix,
)

coordinates_scene = na.DopplerPositionalVectorArray.from_velocity(
    velocity=velocity,
    wavelength_rest=wavelength_rest,
    position=position_scene,
)
coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity(
    velocity=velocity,
    wavelength_rest=wavelength_rest,
    position=position_sensor,
)

scene = ctis.scenes.gaussians(coordinates_scene)
scene = scene + scene.max() / 100

Define the four dispersion angles and the dispersion of the instrument, also as in Invert a Synthetic Scene with MART.

angle = na.linspace(0, 360, num=4, axis="channel", endpoint=False) * u.deg + 5.64 * u.deg

dispersion = 10 * u.km / u.s
dispersion = dispersion.to(u.AA, equivalencies=u.doppler_optical(wavelength_rest))
dispersion = (dispersion - wavelength_rest) / u.pix

Now define a PSF for each channel: an elliptical Gaussian with its own widths and orientation. A PSF is given as the fraction of the light landing in a pixel which reaches each of the pixels around it, sampled on a grid of offsets centered on the pixel the light lands in. Here the grid is 21 by 21 pixels, along two new axes, psf_x and psf_y, which run along the sensor axes sensor_x and sensor_y. Because the widths and the orientation vary along the channel axis, so does the PSF, and each channel gets its own.

Strictly, the fraction of the light reaching each pixel is the PSF convolved with a pixel twice, once for the pixel the light lands in and once for the pixel it is collected in. A Gaussian sampled at the pixel centers is a good enough stand-in for this tutorial.

offset = na.Cartesian2dVectorArray(
    x=na.arange(-10, 11, axis="psf_x"),
    y=na.arange(-10, 11, axis="psf_y"),
)

# the widths, in pixels, and the orientation of the PSF of each channel
width_major = na.ScalarArray(np.array([2.0, 3.0, 3.5, 2.5]), axes="channel")
width_minor = na.ScalarArray(np.array([0.8, 0.6, 0.9, 0.5]), axes="channel")
orientation = na.ScalarArray(np.array([0, 30, 75, 135]), axes="channel") * u.deg

offset_rotated = na.Cartesian2dRotationMatrixArray(-orientation) @ offset

psf = np.exp(
    -(
        np.square(offset_rotated.x / width_major)
        + np.square(offset_rotated.y / width_minor)
    ) / 2
)

Normalize each PSF so that it redistributes the light landing in a pixel without adding or removing any.

psf = psf / psf.sum(("psf_x", "psf_y"))

Display the PSF of each channel.

fig, ax = na.plt.subplots(
    axis_cols="channel",
    ncols=4,
    sharex=True,
    sharey=True,
    figsize=(9, 2.8),
    constrained_layout=True,
)
na.plt.pcolormesh(
    offset.x,
    offset.y,
    C=psf,
    ax=ax,
)
na.plt.set_title("PSF " + (angle.to_string_array("%03d")), ax=ax)
na.plt.set_xlabel("$x$ offset (pix)", ax=ax)
na.plt.set_aspect("equal", ax=ax)
ax.ndarray[0].set_ylabel("$y$ offset (pix)");

Create the instrument, passing it the PSFs. Its axis_psf_xy attribute says which axes of the PSF run along the two sensor axes, and it defaults to ("psf_x", "psf_y"), the names used above.

instrument = ctis.instruments.IdealInstrument(
    area_effective=1 * u.cm ** 2,
    timedelta_exposure=20 * u.s,
    plate_scale=.4 * u.arcsec / u.pix,
    dispersion=dispersion,
    angle=angle,
    wavelength_ref=wavelength_rest,
    position_ref=na.Cartesian2dVectorArray(64, 32) * u.pix,
    coordinates_scene=coordinates_scene,
    coordinates_sensor=coordinates_sensor,
    channel="dispersion angle = " + angle.to_string_array("%03d"),
    axis_channel="channel",
    axis_wavelength="wavelength",
    axis_scene_xy=("scene_x", "scene_y"),
    axis_sensor_xy=("sensor_x", "sensor_y"),
    psf=psf,
)

For comparison, define an instrument which is the same in every way except that it has no PSF.

instrument_no_psf = dataclasses.replace(instrument, psf=None)

Image the scene with the instrument. These are the observed images, which include both the PSFs and photon shot noise.

images = instrument.image(scene)

Display the images each channel would record without a PSF above the images it does record.

images_no_psf = instrument_no_psf.image(scene)

with astropy.visualization.quantity_support():
    fig, axs = plt.subplots(
        nrows=2,
        ncols=4,
        sharex=True,
        sharey=True,
        figsize=(10, 3.1),
        constrained_layout=True,
    )
    vmax = images.outputs.value.max().ndarray
    title = angle.to_string_array("%03d")
    for c in range(instrument.num_channel):
        index = dict(channel=c)
        for row, image in enumerate((images_no_psf, images)):
            axs[row, c].pcolormesh(
                image.inputs.position.x.ndarray.value,
                image.inputs.position.y.ndarray.value,
                image.outputs[index].value.ndarray_aligned(("sensor_y", "sensor_x")),
                cmap="gray",
                vmin=0,
                vmax=vmax,
            )
            axs[row, c].set_aspect("equal")
        axs[0, c].set_title(title[index].ndarray)
        axs[1, c].set_xlabel("sensor $x$ (pix)")
    axs[0, 0].set_ylabel("without PSF")
    axs[1, 0].set_ylabel("observed")
    fig.supylabel("sensor $y$ (pix)")

Invert the observed images with MART twice: once with the instrument which models the PSFs, and once with the one which ignores them. Save the intermediate solutions of both, so that we can look at them iteration by iteration.

inversion = ctis.inverters.MartInverter(
    instrument=instrument,
    intermediate=True,
)(images)

inversion_no_psf = ctis.inverters.MartInverter(
    instrument=instrument_no_psf,
    intermediate=True,
)(images)

Plot the mean \(\langle \chi^2 \rangle\) of the channels at each iteration of both inversions. With the PSFs modeled, the predicted images can match the observed ones to within the noise, and \(\langle \chi^2 \rangle\) falls below one. Without them, no scene blurred only by the instrument without PSFs can reproduce the blurred images, and \(\langle \chi^2 \rangle\) stalls well above one.

chi2 = inversion.mean_chi_squared.mean(instrument.axis_channel)
chi2_no_psf = inversion_no_psf.mean_chi_squared.mean(instrument.axis_channel)

fig, ax = plt.subplots(constrained_layout=True)
na.plt.plot(
    inversion.iteration,
    chi2,
    ax=ax,
    label="modeling the PSFs",
)
na.plt.plot(
    inversion_no_psf.iteration,
    chi2_no_psf,
    ax=ax,
    label="ignoring the PSFs",
)
ax.axhline(1, color="black", linestyle="dashed")
ax.set_xlabel("iteration")
ax.set_ylabel(r"$\langle \chi^2 \rangle$")
ax.set_yscale("log")
ax.legend();

Once \(\langle \chi^2 \rangle\) falls below one, the reconstruction matches the observed images more closely than the noise in them warrants, and further iterations fit the noise rather than the scene. So stop the inversion which models the PSFs at the first iteration where \(\langle \chi^2 \rangle\) reaches one, the discrepancy principle, by running MART again for that many iterations.

reached = chi2.ndarray <= 1
num_iteration = int(np.argmax(reached)) + 1 if reached.any() else inversion.num_iteration

# MART warns when it stops at its iteration limit, which is the point here
with warnings.catch_warnings():
    warnings.simplefilter("ignore", UserWarning)
    inversion = ctis.inverters.MartInverter(
        instrument=instrument,
        num_iteration=num_iteration,
    )(images)

num_iteration

16

Display the original scene next to the reconstructions which ignore and which model the PSFs.

with astropy.visualization.quantity_support():
    fig, axs = plt.subplots(
        ncols=4,
        gridspec_kw=dict(width_ratios=[.3, .3, .3, .1]),
        sharey=False,
        constrained_layout=True,
        figsize=(12, 4.2),
    )
    kwargs = dict(
        axis_wavelength="wavelength",
        vmin=0,
        vmax=scene.outputs.max(),
    )
    colorbar = na.plt.rgbmesh(C=scene, ax=axs[0], **kwargs)
    na.plt.rgbmesh(C=inversion_no_psf.solution, ax=axs[1], **kwargs)
    na.plt.rgbmesh(C=inversion.solution, ax=axs[2], **kwargs)
    na.plt.pcolormesh(
        C=colorbar,
        axis_rgb="wavelength",
        ax=axs[3],
    )
    axs[0].set_title("original")
    axs[1].set_title("ignoring the PSFs")
    axs[2].set_title("modeling the PSFs")
    for ax in axs[:3]:
        ax.set_aspect("equal")
        ax.set_xlabel(f"scene $x$ ({ax.get_xlabel()})")
    axs[0].set_ylabel(f"scene $y$ ({axs[0].get_ylabel()})")
    axs[1].set_ylabel("")
    axs[2].set_ylabel("")
    axs[3].xaxis.set_ticks_position("top")
    axs[3].xaxis.set_label_position("top")
    axs[3].yaxis.tick_right()
    axs[3].yaxis.set_label_position("right")

Compare the true and the reconstructed radiance, median (Doppler shift), and interquartile range (Doppler width) of every pixel in the scene, first for the inversion which ignores the PSFs

inversion_no_psf.plot_moments(scene, axis="wavelength");

and then for the one which models them.

inversion.plot_moments(scene, axis="wavelength");

Modeling the PSFs recovers all three moments more faithfully: the reconstruction is sharper, so its radiance follows the true radiance more closely, and the Doppler shift and width of each pixel are no longer mixed with those of its neighbors.

The PSF may vary along other axes as well. A PSF which varies along wavelength gives each wavelength its own PSF, and one which varies along sensor_x and sensor_y varies across the field.