Add PSF model based on Zernike coefficients

docs/changes/3056.feature.rst

@@ -0,0 +1,5 @@
+Added ``ZernikePSFModel``, a new PSF model, which reconstructs the
+optical wavefront from per-telescope Zernike coefficients (Noll
+indexing) and computes the point spread function via Fraunhofer
+diffraction, polychromatically averaged over a configurable wavelength
+range and weighted by a Cherenkov-like spectral index.

docs/conf.py

@@ -154,13 +154,10 @@ def add_reference_type(prefix, objs):
         "DTypeLike",
         # astropy, new errors in 8.0, see https://github.com/astropy/astropy/issues/19933
         "Attribute",
-        "r.BaseDifferential",
-        "r.BaseRepresentation",
-        "r.BaseRepresentationOrDifferential",
-        "RepresentationMapping",
         "BaseDifferential",
         "BaseRepresentation",
         "BaseRepresentationOrDifferential",
+        "RepresentationMapping",
     ],
 )
 nitpick_ignore += add_reference_type(

examples/tutorials/psf_model.py

@@ -7,7 +7,7 @@
 import astropy.units as u
 import numpy as np
 from itertools import product
-from ctapipe.instrument.optics import ComaPSFModel
+from ctapipe.instrument.optics import ComaPSFModel, ZernikePSFModel
 from ctapipe.instrument import SubarrayDescription
 import matplotlib.pyplot as plt
 
@@ -26,8 +26,8 @@
 
 
 ######################################################################
-# This sets up the PSF model describing pure coma aberrations PSF
-# effect for the LSTs. The parameters are taken from
+# This sets up the PSF models describing PSF effect for the LSTs.
+# The parameters for Coma PSF model are taken from
 # :cite:p:`startracker`, which was original given in polar coordinates
 # in the camera frame. We here manually convert the parameters using
 # the plate scale of LSTs to get the parameters in the TelescopeFrame.
@@ -41,7 +41,7 @@
 
 lst1 = subarray.select_subarray([1])
 
-psf_model = ComaPSFModel(
+coma_psf_model = ComaPSFModel(
     subarray=lst1,
     asymmetry_max=0.49244797,
     asymmetry_decay_rate=9.23573115 / lst_plate_scale_deg,
@@ -55,6 +55,25 @@
     polar_scale_offset=0.02037972 * lst_plate_scale_deg,
 )
 
+zernike_psf_model = ZernikePSFModel(
+    subarray=lst1,
+    pupil_size=512,
+    psf_extent=[("type", "*", 0.5 * u.deg)],
+    pupil_diameter_fraction=0.12,
+    pupil_edge_softness=0.08,
+    focal_plane_smoothing_sigma_pix=3.0,
+    wavelength_min=300e-9 * u.m,
+    wavelength_max=600e-9 * u.m,
+    wavelength_samples=30,
+    cherenkov_spectrum_index=2.0,
+    z4=[("type", "*", 1.825e-07 * u.m)],
+    z5=[("type", "*", 0.0 * u.m)],
+    z6=[("type", "*", 0.0 * u.m)],
+    z11=[("type", "*", 4.467e-08 * u.m)],
+    coma_radial_growth=[("type", "*", 1.919e-07 * u.m / u.deg)],
+    z5_theta2=[("type", "*", 3.501e-08 * u.m / u.deg**2)],
+    z6_theta2=[("type", "*", 3.501e-08 * u.m / u.deg**2)],
+)
 
 ######################################################################
 # calculate PSF at different positions in the field of view
@@ -74,8 +93,20 @@
 centers_y = 0.5 * (edges_y[:-1] + edges_y[1:])
 x, y = np.meshgrid(centers_x, centers_y)
 
-psf_center = psf_model.pdf(tel_id=1, lon=x, lat=y, lon0=0.0 * u.deg, lat0=0.0 * u.deg)
-psf_border = psf_model.pdf(
+psf_center_coma = coma_psf_model.pdf(
+    tel_id=1, lon=x, lat=y, lon0=0.0 * u.deg, lat0=0.0 * u.deg
+)
+psf_center_zernike = zernike_psf_model.pdf(
+    tel_id=1, lon=x, lat=y, lon0=0.0 * u.deg, lat0=0.0 * u.deg
+)
+psf_border_coma = coma_psf_model.pdf(
+    tel_id=1,
+    lon=x + 1 * lst_plate_scale_deg * u.deg,
+    lat=y + 1 * lst_plate_scale_deg * u.deg,
+    lon0=lon0,
+    lat0=lat0,
+)
+psf_border_zernike = zernike_psf_model.pdf(
     tel_id=1,
     lon=x + 1 * lst_plate_scale_deg * u.deg,
     lat=y + 1 * lst_plate_scale_deg * u.deg,
@@ -89,30 +120,46 @@
 # ----------------
 #
 
-fig, (ax1, ax2) = plt.subplots(1, 2, layout="constrained", figsize=(8, 4))
-
-ax1.pcolormesh(
-    edges_x.to_value(u.deg),
-    edges_y.to_value(u.deg),
-    psf_center,
-    cmap="inferno",
+fig, axes = plt.subplots(
+    2,
+    2,
+    layout="constrained",
+    figsize=(10, 8),
 )
 
-ax2.pcolormesh(
-    edges_x.to_value(u.deg),
-    edges_y.to_value(u.deg),
-    psf_border,
-    cmap="inferno",
-)
+plots = [
+    (axes[0, 0], psf_center_coma, "Coma PSF at (0°, 0°)"),
+    (axes[0, 1], psf_center_zernike, "Zernike PSF at (0°, 0°)"),
+    (
+        axes[1, 0],
+        psf_border_coma,
+        f"Coma PSF at ({lon0.to_value(u.deg):.2f}°, {lat0.to_value(u.deg):.2f}°)",
+    ),
+    (
+        axes[1, 1],
+        psf_border_zernike,
+        f"Zernike PSF at ({lon0.to_value(u.deg):.2f}°, {lat0.to_value(u.deg):.2f}°)",
+    ),
+]
+
+for ax, psf, title in plots:
+    plot_vmax = max(float(np.percentile(psf, 99.5)), 1e-12)
+    ax.pcolormesh(
+        edges_x.to_value(u.deg),
+        edges_y.to_value(u.deg),
+        psf,
+        cmap="inferno",
+        vmin=0.0,
+        vmax=plot_vmax,
+        shading="auto",
+    )
+    ax.set(
+        aspect=1,
+        title=title,
+        xlabel="lon [deg]",
+        ylabel="lat [deg]",
+    )
 
-ax1.set(
-    aspect=1,
-    title="PSF at (0°, 0°)",
-)
-ax2.set(
-    aspect=1,
-    title=f"PSF at ({lon0.to_value(u.deg):.2f}°, {lat0.to_value(u.deg):.2f}°)",
-)
 plt.show()
 
 
@@ -130,22 +177,57 @@
 centers_y_stack = 0.5 * (edges_y_stack[:-1] + edges_y_stack[1:])
 x_stack, y_stack = np.meshgrid(centers_x_stack, centers_y_stack)
 
-psf_stacked = np.zeros(x_stack.shape)
+psf_stacked_coma = np.zeros(x_stack.shape)
+psf_stacked_zernike = np.zeros(x_stack.shape)
 for source_lon, source_lat in product(lons, lats):
-    psf_stacked += psf_model.pdf(
+    psf_stacked_coma += coma_psf_model.pdf(
+        tel_id=1,
+        lon=x_stack,
+        lat=y_stack,
+        lon0=source_lon,
+        lat0=source_lat,
+    )
+    psf_stacked_zernike += zernike_psf_model.pdf(
         tel_id=1,
         lon=x_stack,
         lat=y_stack,
         lon0=source_lon,
         lat0=source_lat,
     )
 
-fig_stack, ax_stack = plt.subplots(1, 1, layout="constrained", figsize=(6, 5))
-mesh = ax_stack.pcolormesh(
+fig_stack, axes = plt.subplots(
+    1,
+    2,
+    layout="constrained",
+    figsize=(10, 5),
+)
+
+axes[0].pcolormesh(
     edges_x_stack.to_value(u.deg),
     edges_y_stack.to_value(u.deg),
-    psf_stacked,
+    psf_stacked_coma,
     cmap="inferno",
+    shading="auto",
+    vmin=0.0,
+    vmax=np.percentile(psf_stacked_coma, 99.5),
 )
-ax_stack.set(aspect=1, title="Stacked PSF over source-position grid")
+axes[0].set(
+    aspect=1,
+    title="Stacked Coma PSF",
+)
+
+axes[1].pcolormesh(
+    edges_x_stack.to_value(u.deg),
+    edges_y_stack.to_value(u.deg),
+    psf_stacked_zernike,
+    cmap="inferno",
+    shading="auto",
+    vmin=0.0,
+    vmax=np.percentile(psf_stacked_zernike, 99.5),
+)
+axes[1].set(
+    aspect=1,
+    title="Stacked Zernike PSF",
+)
+
 plt.show()

pyproject.toml

@@ -60,6 +60,7 @@ all = [
     "iminuit >=2",
     "matplotlib ~=3.0",
     "pyirf ~=0.14.0",
+    "zernike",
 ]
 
 tests = [