diff --git a/docs/changes/3056.feature.rst b/docs/changes/3056.feature.rst new file mode 100644 index 00000000000..e1445687864 --- /dev/null +++ b/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. diff --git a/docs/conf.py b/docs/conf.py index 483558ef312..5d0678e56bf 100644 --- a/docs/conf.py +++ b/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( diff --git a/examples/tutorials/psf_model.py b/examples/tutorials/psf_model.py index 74235227fe6..31b4c6b4005 100644 --- a/examples/tutorials/psf_model.py +++ b/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,9 +177,17 @@ 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, @@ -140,12 +195,39 @@ 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() diff --git a/pyproject.toml b/pyproject.toml index 2b70f680079..7905984f605 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ all = [ "iminuit >=2", "matplotlib ~=3.0", "pyirf ~=0.14.0", + "zernike", ] tests = [ diff --git a/src/ctapipe/instrument/optics.py b/src/ctapipe/instrument/optics.py index 0b8ed1b1acc..54bb09e73cd 100644 --- a/src/ctapipe/instrument/optics.py +++ b/src/ctapipe/instrument/optics.py @@ -5,16 +5,32 @@ import logging from abc import abstractmethod from enum import Enum, StrEnum, auto, unique +from functools import cached_property from pathlib import Path import astropy.units as u import numpy as np from astropy.table import QTable, Table +from numpy.fft import fft2, fftshift +from scipy.ndimage import gaussian_filter, map_coordinates from scipy.stats import laplace, laplace_asymmetric +try: + from zernike import RZern +except ModuleNotFoundError: + RZern = None + from ..coordinates import TelescopeFrame from ..core import TelescopeComponent -from ..core.traits import FloatTelescopeParameter +from ..core.traits import ( + AstroQuantity, + Float, + FloatTelescopeParameter, + Int, + IntTelescopeParameter, + TelescopeParameter, +) +from ..exceptions import OptionalDependencyMissing from ..utils import get_table_dataset from ..utils.quantities import all_to_value from .warnings import warn_from_name @@ -26,6 +42,7 @@ "FocalLengthKind", "PSFModel", "ComaPSFModel", + "ZernikePSFModel", ] @@ -627,3 +644,393 @@ def pdf(self, tel_id, lon, lat, lon0, lat0): pdf = radial_pdf * polar_pdf * inv_r return pdf + + +class ZernikePSFModel(PSFModel): + r"""PSF model based on wavefront reconstruction using Zernike wavefront coefficients. + + This model reconstructs the optical wavefront from a set of Zernike + polynomial coefficients and computes the point spread function (PSF) + by Fourier propagation through the telescope pupil. The resulting PSF + naturally includes diffraction and wavefront aberrations and is + evaluated for arbitrary field positions by allowing selected Zernike + coefficients to vary with the source position in the focal plane. + + The model includes: + + - `Zernike polynomials in Noll `__ + indexing to describe the optical path + difference (OPD) across the telescope pupil. This indexing convention + maps the standard radial and azimuthal modes to a 1D sequence, matching + industry standards such as Ansys Zemax. + - Scalar Fourier optics to propagate the complex pupil field into the + focal plane. + - Polychromatic averaging over the Cherenkov emission spectrum using a + configurable wavelength range and spectral weighting. + - Optional Gaussian smoothing to approximate detector and residual + instrumental broadening not explicitly included in the wavefront + model. + + In the current parameterization, field-dependent aberrations are + represented by linear coma terms and quadratic astigmatism terms, + providing a compact phenomenological description of off-axis optical + degradation while retaining a physically motivated wavefront model. + """ + + # Universal model performance parameters + pupil_diameter_fraction = Float( + default_value=0.12, + help=( + "Diameter of the telescope pupil as a fraction of the FFT grid size. " + "Smaller values increase focal-plane sampling at the expense of " + "undersampling the pupil, while larger values improve pupil sampling " + "but reduce the field of view and sampling resolution of the computed PSF." + ), + ).tag(config=True) + + noll_max = 11 # highest Noll index + + wavelength_samples = Int( + default_value=20, + help="Number of wavelength samples for polychromatic averaging", + ).tag(config=True) + + # Universal physical constants + cherenkov_spectrum_index = Float( + default_value=2.0, + help="Power-law index for Cherenkov spectrum weighting (dN/dλ ∝ λ^-index)", + ).tag(config=True) + + wavelength_min = AstroQuantity( + default_value=300e-9 * u.m, + physical_type=u.physical.length, + help="Minimum wavelength for polychromatic averaging", + ).tag(config=True) + + wavelength_max = AstroQuantity( + default_value=600e-9 * u.m, + physical_type=u.physical.length, + help="Maximum wavelength for polychromatic averaging", + ).tag(config=True) + + # Per-telescope optical parameters + pupil_size = IntTelescopeParameter( + default_value=512, + help=( + "Number of samples across the FFT grid used to discretize the pupil. " + "Larger values improve numerical accuracy and PSF sampling at the " + "expense of increased memory usage and computation time." + ), + ).tag(config=True) + + pupil_edge_softness = FloatTelescopeParameter( + default_value=0.08, + help=( + "Width of the sigmoid taper applied to the pupil edge in normalized " + "pupil-radius units. Larger values suppress diffraction ringing but " + "slightly blur the effective aperture." + ), + ).tag(config=True) + + focal_plane_smoothing_sigma_pix = FloatTelescopeParameter( + default_value=3.0, + help="Gaussian smoothing sigma applied to PSF intensity.", + ).tag(config=True) + + psf_extent = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.angle), + default_value=0.5 * u.deg, + help=( + "Angular width of the square grid used to represent a single " + "point source's PSF, centered on the source position. Must be " + "large enough to contain the full extent of the PSF (including " + "aberration tails) or normalization will be biased low; " + "not the telescope's camera field of view." + ), + ).tag(config=True) + + z2 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Tilt X", + ).tag(config=True) + z3 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Tilt Y", + ).tag(config=True) + z4 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=1.825e-07 * u.m, + help="Defocus", + ).tag(config=True) + z5 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Astigmatism 45°", + ).tag(config=True) + z6 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Astigmatism 0°", + ).tag(config=True) + z7 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Vertical Coma", + ).tag(config=True) + z8 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Horizontal Coma", + ).tag(config=True) + z9 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Vertical trefoil", + ).tag(config=True) + z10 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=0.0 * u.m, + help="Horizontal trefoil", + ).tag(config=True) + z11 = TelescopeParameter( + trait=AstroQuantity(physical_type=u.physical.length), + default_value=3.467e-08 * u.m, + help="Spherical", + ).tag(config=True) + + coma_radial_growth = TelescopeParameter( + trait=AstroQuantity(physical_type=(u.m / u.deg).physical_type), + default_value=1.919e-07 * u.m / u.deg, + help="Radial coma growth", + ).tag(config=True) + + z5_theta2 = TelescopeParameter( + trait=AstroQuantity(physical_type=(u.m / u.deg**2).physical_type), + default_value=3.501e-08 * u.m / u.deg**2, + help="Quadratic astigmatism growth", + ).tag(config=True) + + z6_theta2 = TelescopeParameter( + trait=AstroQuantity(physical_type=(u.m / u.deg**2).physical_type), + default_value=3.501e-08 * u.m / u.deg**2, + help="Quadratic astigmatism growth", + ).tag(config=True) + + def __init__(self, subarray, config=None, parent=None, **kwargs): + """Initialize the ZernikePSFModel component and check for missing optional dependency.""" + if RZern is None: + raise OptionalDependencyMissing("zernike") + super().__init__( + subarray=subarray, + config=config, + parent=parent, + **kwargs, + ) + + @cached_property + def _radial_order(self): + n = 0 + while (n + 1) * (n + 2) // 2 < self.noll_max: + n += 1 + return max(1, n) + + def _zernike_grid(self, tel_id): + if not hasattr(self, "_zernike_grid_cache"): + self._zernike_grid_cache = {} + if tel_id in self._zernike_grid_cache: + return self._zernike_grid_cache[tel_id] + n = self.pupil_size.tel[tel_id] + frac = self.pupil_diameter_fraction + if not (0 < frac <= 1): + raise ValueError("pupil_diameter_fraction must be in (0, 1]") + + coord_limit = 1.0 / frac + x = np.linspace(-coord_limit, coord_limit, n) + y = np.linspace(-coord_limit, coord_limit, n) + xx, yy = np.meshgrid(x, y) + + rr = np.sqrt(xx**2 + yy**2) + mask = rr <= 1 + + edge = max(self.pupil_edge_softness.tel[tel_id], 1e-6) + aperture = 1.0 / (1.0 + np.exp(np.clip((rr - 1.0) / edge, -60.0, 60.0))) + rz = RZern(self._radial_order) + rz.make_cart_grid(xx, yy) + self._zernike_grid_cache[tel_id] = (rz, mask, aperture) + + return rz, mask, aperture + + def _coeff_vector(self, tel_id, lon0_deg, lat0_deg): + """ + Build the Noll coefficient vector [m] for a given telescope and + field-of-view offset. Internally works with plain floats in fixed + units (m for OPD amplitudes, deg for angles) after unwrapping the + Quantity-valued traits, since the downstream Zernike/FFT machinery + is unit-agnostic. + """ + rz, _, _ = self._zernike_grid(tel_id) + coeff = np.zeros(rz.nk) + + theta2 = lon0_deg**2 + lat0_deg**2 + theta = np.sqrt(theta2) + if theta > 0: + ux = lon0_deg / theta + uy = lat0_deg / theta + else: + ux = 0.0 + uy = 0.0 + + coma_radial_growth_deg = self.coma_radial_growth.tel[tel_id].to_value( + u.m / u.deg + ) + z5_theta2_m_per_deg2 = self.z5_theta2.tel[tel_id].to_value(u.m / u.deg**2) + z6_theta2_m_per_deg2 = self.z6_theta2.tel[tel_id].to_value(u.m / u.deg**2) + + # Project astigmatism rotation (2*phi) + cos_2phi = ux**2 - uy**2 + sin_2phi = 2.0 * ux * uy + + coma_radial = coma_radial_growth_deg * theta + coma_x = coma_radial * ux + coma_y = coma_radial * uy + noll_coeffs = [ + 0.0, + self.z2.tel[tel_id].to_value(u.m), + self.z3.tel[tel_id].to_value(u.m), + self.z4.tel[tel_id].to_value(u.m), + self.z5.tel[tel_id].to_value(u.m) + + z5_theta2_m_per_deg2 * theta2 * sin_2phi, + self.z6.tel[tel_id].to_value(u.m) + + z6_theta2_m_per_deg2 * theta2 * cos_2phi, + self.z7.tel[tel_id].to_value(u.m) + coma_y, + self.z8.tel[tel_id].to_value(u.m) + coma_x, + self.z9.tel[tel_id].to_value(u.m), + self.z10.tel[tel_id].to_value(u.m), + self.z11.tel[tel_id].to_value(u.m), + ] + + upper = min(len(noll_coeffs), self.noll_max, rz.nk) + coeff[1:upper] = noll_coeffs[1:upper] + return coeff + + def _build_psf(self, tel_id, lon0_deg, lat0_deg): + rz, mask, aperture = self._zernike_grid(tel_id) + coeff = self._coeff_vector(tel_id, lon0_deg, lat0_deg) + wavefront = rz.eval_grid(coeff, matrix=True) + + lam_min = self.wavelength_min.to_value(u.m) + lam_max = self.wavelength_max.to_value(u.m) + if self.wavelength_samples < 1: + raise ValueError("wavelength_samples must be >= 1") + if lam_min <= 0 or lam_max <= 0 or lam_max < lam_min: + raise ValueError("Invalid wavelength_min/wavelength_max") + + wavelengths = np.linspace(lam_min, lam_max, self.wavelength_samples) + weights = wavelengths ** (-self.cherenkov_spectrum_index) + weights /= np.sum(weights) + + # Reference wavelength (mean of the range) + lambda_ref = np.mean(wavelengths) + + intensity = np.zeros_like(wavefront, dtype=float) + h, w = intensity.shape + cy, cx = h // 2, w // 2 + y_ref, x_ref = np.indices((h, w), dtype=float) + + for wavelength, weight in zip(wavelengths, weights): + phase = 2 * np.pi * wavefront / wavelength + phase_safe = np.zeros_like(phase) + phase_safe[mask] = phase[mask] + pupil = aperture * np.exp(1j * phase_safe) + field = fft2(pupil) + intensity_lambda = fftshift(np.abs(field) ** 2) + + scale = lambda_ref / wavelength + y_lambda = (y_ref - cy) * scale + cy + x_lambda = (x_ref - cx) * scale + cx + + # Interpolate intensity_lambda at the scaled coordinates + intensity_rescaled = map_coordinates( + intensity_lambda, + np.array( + [y_lambda, x_lambda] + ), # map_coordinates expects a single array or tuple + order=3, + mode="constant", + cval=0.0, + ) + + # Accumulate and protect against minor cubic interpolation under-shoots + intensity += weight * np.clip(intensity_rescaled, 0.0, None) # / (scale**2) + + sigma_pix = max(self.focal_plane_smoothing_sigma_pix.tel[tel_id], 0.0) + if sigma_pix > 0: + intensity = gaussian_filter(intensity, sigma=sigma_pix, mode="nearest") + + # Convert from discrete probability per FFT pixel to + # probability density per angular area. + total = intensity.sum() + if not np.isfinite(total) or total <= 0: + raise RuntimeError( + f"Invalid PSF normalization: total intensity is {total!r}" + ) + intensity /= total + + psf_extent = self.psf_extent.tel[tel_id].to_value(u.deg) + pixel_scale = psf_extent / (self.pupil_size.tel[tel_id] - 1) + pixel_area = pixel_scale**2 + + total_volume = intensity.sum() * pixel_area + if not np.isfinite(total_volume) or total_volume <= 0: + raise RuntimeError( + f"Invalid PSF normalization: total integrated volume is {total_volume!r}" + ) + + intensity /= total_volume + + return intensity + + @u.quantity_input( + lon=u.deg, + lat=u.deg, + lon0=u.deg, + lat0=u.deg, + ) + def pdf(self, tel_id, lon, lat, lon0, lat0): + dx = np.asarray((lon - lon0).to_value(u.deg)) + dy = np.asarray((lat - lat0).to_value(u.deg)) + input_shape = dx.shape + + lon0_deg = lon0.to_value(u.deg) + lat0_deg = lat0.to_value(u.deg) + intensity = self._build_psf(tel_id, lon0_deg, lat0_deg) + + full_field = self.psf_extent.tel[tel_id].to_value(u.deg) + half_field = 0.5 * full_field + if half_field <= 0: + raise ValueError("psf_extent must be > 0") + + n = self.pupil_size.tel[tel_id] + xpix = (dx + half_field) / (2 * half_field) * (n - 1) + ypix = (dy + half_field) / (2 * half_field) * (n - 1) + + # map_coordinates requires the "points" dimension to have rank >= 1; + # scalar lon/lat collapse to 0-d, so flatten for the call and restore + # the original shape (including scalar) afterward. + coords = np.array([np.atleast_1d(ypix).ravel(), np.atleast_1d(xpix).ravel()]) + + psf = map_coordinates( + intensity, + coords, + order=1, + mode="constant", + cval=0.0, + prefilter=False, + ) + + psf = np.asarray(psf, dtype=float).reshape(input_shape) + psf = np.clip(psf, 0.0, None) + + return psf.item() if psf.shape == () else psf diff --git a/src/ctapipe/instrument/tests/test_psf_model.py b/src/ctapipe/instrument/tests/test_psf_model.py index daa34467e25..07aa3d4d571 100644 --- a/src/ctapipe/instrument/tests/test_psf_model.py +++ b/src/ctapipe/instrument/tests/test_psf_model.py @@ -9,6 +9,8 @@ from ctapipe.compat import trapz_func from ctapipe.instrument.optics import PSFModel +pytest.importorskip("zernike") + @pytest.fixture(scope="session") def coma_psf(example_subarray): @@ -33,29 +35,62 @@ def coma_psf(example_subarray): return psf -def test_asymptotic_behavior(coma_psf): +@pytest.fixture(scope="session") +def zernike_psf(example_subarray): + return PSFModel.from_name("ZernikePSFModel", subarray=example_subarray) + + +# Source position and evaluation grids used by all tests below. These are +# identical for every PSF model; only the model instance under test changes. +SOURCE_LON0 = 2.13 * u.deg +SOURCE_LAT0 = -0.37 * u.deg + +ASYMPTOTIC_LON = 20.0 * u.deg +ASYMPTOTIC_LAT = 0.0 * u.deg +ASYMPTOTIC_LON0 = 2.0 * u.deg +ASYMPTOTIC_LAT0 = 0.0 * u.deg + +NORM_LON = np.linspace(-5.0, 7.0, 601) * u.deg +NORM_LAT = np.linspace(-4.0, 4.0, 401) * u.deg + +CENTER_LON = np.linspace(-1.0, 1.0, 201) * u.deg +CENTER_LAT = np.linspace(-1.0, 1.0, 201) * u.deg + +SOURCE_LON = 2.0 * u.deg +SOURCE_LAT = 0.0 * u.deg + +NORM_REL = 0.05 +NORM_ABS = 0.02 + + +@pytest.fixture(params=["coma_psf", "zernike_psf"]) +def psf_model(request): + return request.getfixturevalue(request.param) + + +def test_asymptotic_behavior(psf_model): assert np.isclose( - coma_psf.pdf( + psf_model.pdf( tel_id=1, - lon=20.0 * u.deg, - lat=0.0 * u.deg, - lon0=2.0 * u.deg, - lat0=0.0 * u.deg, + lon=ASYMPTOTIC_LON, + lat=ASYMPTOTIC_LAT, + lon0=ASYMPTOTIC_LON0, + lat0=ASYMPTOTIC_LAT0, ), 0.0, atol=1e-7, ) -def test_normalization(coma_psf): - lon0 = 2.13 * u.deg - lat0 = -0.37 * u.deg +def test_normalization(psf_model): + lon0 = SOURCE_LON0 + lat0 = SOURCE_LAT0 - lon = np.linspace(-5.0, 7.0, 601) * u.deg - lat = np.linspace(-4.0, 4.0, 401) * u.deg + lon = NORM_LON + lat = NORM_LAT lon_grid, lat_grid = np.meshgrid(lon, lat, indexing="xy") - pdf = coma_psf.pdf( + pdf = psf_model.pdf( tel_id=1, lon=lon_grid, lat=lat_grid, @@ -63,22 +98,25 @@ def test_normalization(coma_psf): lat0=lat0, ) + assert np.isfinite(pdf).all() + assert np.all(pdf >= 0.0) + integral = trapz_func( trapz_func(pdf, lon.to_value(u.deg), axis=1), lat.to_value(u.deg) ) - assert integral == pytest.approx(1.0, rel=0.05, abs=0.02) + assert integral == pytest.approx(1.0, rel=NORM_REL, abs=NORM_ABS) -def test_normalization_at_camera_center(coma_psf): +def test_normalization_at_camera_center(psf_model): lon0 = 0.0 * u.deg lat0 = 0.0 * u.deg - lon = np.linspace(-1.0, 1.0, 201) * u.deg - lat = np.linspace(-1.0, 1.0, 201) * u.deg + lon = CENTER_LON + lat = CENTER_LAT lon_grid, lat_grid = np.meshgrid(lon, lat, indexing="xy") - pdf = coma_psf.pdf( + pdf = psf_model.pdf( tel_id=1, lon=lon_grid, lat=lat_grid, @@ -90,24 +128,24 @@ def test_normalization_at_camera_center(coma_psf): trapz_func(pdf, lon.to_value(u.deg), axis=1), lat.to_value(u.deg) ) - assert integral == pytest.approx(1.0, rel=0.05, abs=0.02) + assert integral == pytest.approx(1.0, rel=NORM_REL, abs=NORM_ABS) -def test_finite_at_source_position(coma_psf): - value = coma_psf.pdf( +def test_finite_at_source_position(psf_model): + value = psf_model.pdf( tel_id=1, - lon=2.0 * u.deg, - lat=0.0 * u.deg, - lon0=2.0 * u.deg, - lat0=0.0 * u.deg, + lon=SOURCE_LON, + lat=SOURCE_LAT, + lon0=ASYMPTOTIC_LON0, + lat0=ASYMPTOTIC_LAT0, ) assert np.isfinite(value) assert value > 0.0 -def test_finite_at_camera_center(coma_psf): - value = coma_psf.pdf( +def test_finite_at_camera_center(psf_model): + value = psf_model.pdf( tel_id=1, lon=0.0 * u.deg, lat=0.0 * u.deg,