Zernikes¤

def zernike_name(j: int) -> str:
    """
    Gets the name of the jth Zernike polynomial.

    Parameters
    ----------
    j : int
        The Zernike (noll) index.

    Returns
    -------
    name : str
        The name of the Zernike polynomial.
    """
    return zernike_names[int(j)] if j >= 1 and j <= 36 else f"Zernike {int(j)}"

def noll_indices(j: int) -> tuple[int]:
    """
    Calculate the radial and azimuthal orders of the Zernike polynomial.

    Parameters
    ----------
    j : int
        The Zernike (noll) index.

    Returns
    -------
    n, m : tuple[int]
        The radial and azimuthal orders of the Zernike polynomial.
    """
    n = (np.ceil(-1 / 2 + np.sqrt(1 + 8 * j) / 2) - 1).astype(int)
    smallest_j_in_row = n * (n + 1) / 2 + 1
    number_of_shifts = (j - smallest_j_in_row + ~(n & 1) + 2) // 2
    sign_of_shift = -(j & 1) + ~(j & 1) + 2
    base_case = n & 1
    m = (sign_of_shift * (base_case + number_of_shifts * 2)).astype(int)
    return int(n), int(m)

def zernike_factors(j: int) -> tuple[Array]:
    """
    Calculates the normalisation coefficients and powers of the Zernike polynomial.

    Parameters
    ----------
    j : int
        The Zernike (noll) index.

    Returns
    -------
    c, k : tuple[Array]
        The normalisation coefficients and powers of the Zernike polynomial.
    """
    if j < 1:
        raise ValueError("The Zernike index must be greater than 0.")
    n, m = noll_indices(j)

    # Calculate k
    k = np.arange(((n - m) // 2) + 1, dtype=float)

    # Calculate c
    sign = lax.pow(-1.0, k)
    _fact_1 = vmap(dlu.factorial)(np.abs(n - k))
    _fact_2 = vmap(dlu.factorial)(k)
    _fact_3 = vmap(dlu.factorial)(((n + m) // 2) - k)
    _fact_4 = vmap(dlu.factorial)(((n - m) // 2) - k)
    c = sign * _fact_1 / _fact_2 / _fact_3 / _fact_4
    return c, k

def zernike(j: int, coordinates: Array, diameter: float = 2) -> Array:
    """
    Calculates the Zernike polynomial. Note that this function is not-jittable as is
    has dynamic array shapes. To use this function in a jittable way, use the
    zernike_fast function, with the pre-calculated c and k parameters.

    Parameters
    ----------
    j : int
        The Zernike (noll) index.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.
    diameter : float = 2
        The diameter of the aperture to calculate the Zernike polynomial upon.

    Returns
    -------
    zernike : Array
        The Zernike polynomial.
    """
    coordinates = scale_coords(coordinates, diameter / 2)
    polar_coordinates = dlu.cart2polar(coordinates)
    rho = polar_coordinates[0]
    theta = polar_coordinates[1]
    aperture = rho <= 1.0
    n, m = noll_indices(j)
    c, k = zernike_factors(j)
    return aperture * eval_radial(rho, n, c, k) * eval_azimuthal(theta, n, m)

def zernike_fast(n: int, m: int, c: Array, k: Array, coordinates: Array) -> Array:
    """
    Calculates the Zernike polynomial using the pre-calculated c and k parameters, such
    that this function is jittable.

    Parameters
    ----------
    n : int
        The radial order of the Zernike polynomial.
    m : int
        The azimuthal order of the Zernike polynomial.
    c : Array
        The normalisation coefficients of the Zernike polynomial.
    k : Array
        The powers of the Zernike polynomial.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.

    Returns
    -------
    zernike : Array
        The Zernike polynomial.
    """
    polar_coordinates = dlu.cart2polar(coordinates)
    rho = polar_coordinates[0]
    theta = polar_coordinates[1]
    aperture = rho <= 1.0
    return aperture * eval_radial(rho, n, c, k) * eval_azimuthal(theta, n, m)

def zernike_basis(js: list[int], coordinates: Array, diameter: float = 2) -> Array:
    """
    Calculates the Zernike polynomial basis. Note that this function is not-jittable.

    Parameters
    ----------
    js : list[int]
        The Zernike (noll) indices.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.
    diameter : float = 2
        The diameter of the aperture to calculate the Zernike polynomial upon.

    Returns
    -------
    zernike_basis : Array
        The Zernike polynomial basis.
    """
    return np.array([zernike(j, coordinates, diameter) for j in js])

def polike(nsides: int, j: int, coordinates: Array, diameter: float = 2) -> Array:
    """
    Calculates the Zernike polynomial on an n-sided aperture. Note that this function
    is not-jittable as is has dynamic array shapes. To use this function in a jittable
    way, use the polike_fast function, with the pre-calculated c and k parameters.

    Parameters
    ----------
    nsides : int
        The number of sides of the aperture.
    j : int
        The Zernike (noll) index.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.
    diameter : float = 2
        The diameter of the aperture to calculate the Zernike polynomial upon.

    Returns
    -------
    polike : Array
        The Zernike polynomial on an n-sided aperture.
    """
    if nsides < 3:
        raise ValueError(f"nsides must be >= 3, not {nsides}.")
    coordinates = scale_coords(coordinates, diameter / 2)
    theta = dlu.cart2polar(coordinates)[1]
    alpha = np.pi / nsides
    phi = theta + alpha
    wedge = np.floor((phi + alpha) / (2.0 * alpha))
    u_alpha = phi - wedge * (2 * alpha)
    r_alpha = np.cos(alpha) / np.cos(u_alpha)
    return 1 / r_alpha * zernike(j, coordinates / r_alpha)

def polike_fast(
    nsides: int, n: int, m: int, c: Array, k: Array, coordinates: Array
) -> Array:
    """
    Calculates the Zernike polynomial on an n-sided aperture using the pre-calculated
    c and k parameters, such that this function is jittable.

    Parameters
    ----------
    nsides : int
        The number of sides of the aperture.
    n : int
        The radial order of the Zernike polynomial.
    m : int
        The azimuthal order of the Zernike polynomial.
    c : Array
        The normalisation coefficients of the Zernike polynomial.
    k : Array
        The powers of the Zernike polynomial.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.

    Returns
    -------
    polike : Array
        The Zernike polynomial on an n-sided aperture.
    """
    if nsides < 3:
        raise ValueError(f"nsides must be >= 3, not {nsides}.")
    alpha = np.pi / nsides
    phi = dlu.cart2polar(coordinates)[1] + alpha
    wedge = np.floor((phi + alpha) / (2.0 * alpha))
    u_alpha = phi - wedge * (2 * alpha)
    r_alpha = np.cos(alpha) / np.cos(u_alpha)
    return 1 / r_alpha * zernike_fast(n, m, c, k, coordinates / r_alpha)

def polike_basis(nsides: int, js: list[int], coordinates: Array, diameter: float = 2):
    """
    Calculates the Zernike polynomial basis on an n-sided aperture. Note that this
    function is not-jittable.

    Parameters
    ----------
    nsides : int
        The number of sides of the aperture.
    js : list[int]
        The Zernike (noll) indices.
    coordinates : Array
        The Cartesian coordinates to calculate the Zernike polynomial upon.
    diameter : float = 2
        The diameter of the aperture to calculate the Zernike polynomial upon.

    Returns
    -------
    polike_basis : Array
        The Zernike polynomial basis on an n-sided aperture.
    """
    return np.array([polike(nsides, j, coordinates, diameter) for j in js])