Source code for astrophot.models.airy

import torch
import numpy as np

from ..utils.decorators import ignore_numpy_warnings, combine_docstrings
from .psf_model_object import PSFModel
from .mixins import RadialMixin
from ..param import forward
from ..backend_obj import backend, ArrayLike

__all__ = ("AiryPSF",)




@combine_docstrings
class AiryPSF(RadialMixin, PSFModel):
    """The Airy disk is an analytic description of the diffraction pattern
    for a circular aperture.

    The diffraction pattern is described exactly by the configuration
    of the lens system under the assumption that all elements are
    perfect. This expression goes as:

    .. math::

       I(\\theta) = I_0\\left[\\frac{2J_1(x)}{x}\\right]^2

    .. math::

       x = ka\\sin(\\theta) = \\frac{2\\pi a r}{\\lambda R}

    where :math:`I(\\theta)` is the intensity as a function of the
    angular position within the diffraction system along its main
    axis, :math:`I_0` is the central intensity of the airy disk,
    :math:`J_1` is the Bessel function of the first kind of order one,
    :math:`k = \\frac{2\\pi}{\\lambda}` is the wavenumber of the
    light, :math:`a` is the aperture radius, :math:`r` is the radial
    position from the center of the pattern, :math:`R` is the distance
    from the circular aperture to the observation plane.

    In the ``Airy_PSF`` class we combine the parameters
    :math:`a,R,\\lambda` into a single ratio to be optimized (or fixed
    by the optical configuration).

    :param I0: The central intensity of the airy disk in flux/arcsec^2.
    :param aRL: The ratio of the aperture radius to the
        product of the wavelength and the distance from the aperture to the
        observation plane, :math:`\\frac{a}{R \\lambda}`.

    """

    _model_type = "airy"
    _parameter_specs = {
        "I0": {
            "units": "flux/arcsec^2",
            "value": 1.0,
            "shape": (),
            "dynamic": False,
            "description": "The central intensity of the airy disk in flux/arcsec^2.",
        },
        "aRL": {
            "units": "a/(R lambda)",
            "shape": (),
            "dynamic": True,
            "description": "The ratio of the aperture radius to the product of the wavelength and the distance from the aperture to the observation plane.",
        },
    }
    usable = True



    @torch.no_grad()
    @ignore_numpy_warnings
    def initialize(self):
        super().initialize()

        if self.I0.initialized and self.aRL.initialized:
            return
        icenter = self.target.targpixel_to_mypixel(*self.center.value)

        if not self.I0.initialized:
            mid_chunk = self.target._data[
                int(icenter[0]) - 2 : int(icenter[0]) + 2,
                int(icenter[1]) - 2 : int(icenter[1]) + 2,
            ]
            self.I0.value = backend.mean(mid_chunk) / self.target.upsample**2
        if not self.aRL.initialized:
            self.aRL.value = (5.0 / 8.0) * 2




    @forward
    def radial_model(self, R: ArrayLike, I0: ArrayLike, aRL: ArrayLike) -> ArrayLike:
        x = 2 * np.pi * aRL * R
        return I0 * (2 * backend.bessel_j1(x) / x) ** 2