Source code for africanus.rime.zernike

import numpy as np


from africanus.util.numba import jit


@jit(nogil=True, nopython=True, cache=True)
def fac(x):
    if x < 0:
        raise ValueError("Factorial input is negative.")
    if x == 0:
        return 1
    factorial = 1
    for i in range(1, x + 1):
        factorial *= i
    return factorial


@jit(nogil=True, nopython=True, cache=True)
def pre_fac(k, n, m):
    numerator = (-1.0) ** k * fac(n - k)
    denominator = fac(k) * fac((n + m) / 2.0 - k) * fac((n - m) / 2.0 - k)
    return numerator / denominator


@jit(nogil=True, nopython=True, cache=True)
def zernike_rad(m, n, rho):
    if n < 0 or m < 0 or abs(m) > n:
        raise ValueError("m and n values are incorrect.")
    radial_component = 0
    for k in range((n - m) / 2 + 1):
        radial_component += pre_fac(k, n, m) * rho ** (n - 2.0 * k)
    return radial_component


@jit(nogil=True, nopython=True, cache=True)
def zernike(j, rho, phi):
    if rho > 1:
        return 0.0
    j += 1
    n = 0
    j1 = j - 1
    while j1 > n:
        n += 1
        j1 -= n
    m = (-1) ** j * ((n % 2) + 2 * int((j1 + ((n + 1) % 2)) / 2.0))
    if m > 0:
        return zernike_rad(m, n, rho) * np.cos(m * phi)
    if m < 0:
        return zernike_rad(-m, n, rho) * np.sin(-m * phi)
    return zernike_rad(0, n, rho)


@jit(nogil=True, nopython=True, cache=True)
def _convert_coords(l_coords, m_coords):
    rho, phi = ((l_coords**2 + m_coords**2) ** 0.5), np.arctan2(l_coords, m_coords)
    return rho, phi


@jit(nogil=True, nopython=True, cache=True)
def nb_zernike_dde(
    coords,
    coeffs,
    noll_index,
    out,
    parallactic_angles,
    frequency_scaling,
    antenna_scaling,
    pointing_errors,
):
    sources, times, ants, chans, corrs = out.shape
    npoly = coeffs.shape[-1]

    for s in range(sources):
        for t in range(times):
            for a in range(ants):
                sin_pa = np.sin(parallactic_angles[t, a])
                cos_pa = np.cos(parallactic_angles[t, a])

                for c in range(chans):
                    l, m, freq = coords[:, s, t, a, c]

                    l_coords = l * frequency_scaling[c]
                    m_coords = m * frequency_scaling[c]

                    l_coords += pointing_errors[t, a, c, 0]
                    m_coords += pointing_errors[t, a, c, 1]

                    vl = l_coords * cos_pa - l_coords * sin_pa
                    vm = m_coords * sin_pa + m * cos_pa

                    vl *= antenna_scaling[a, c, 0]
                    vm *= antenna_scaling[a, c, 1]

                    rho, phi = _convert_coords(vl, vm)

                    for co in range(corrs):
                        zernike_sum = 0

                        for p in range(npoly):
                            zc = coeffs[a, c, co, p]
                            zn = noll_index[a, c, co, p]
                            zernike_sum += zc * zernike(zn, rho, phi)

                        out[s, t, a, c, co] = zernike_sum

    return out


[docs] def zernike_dde( coords, coeffs, noll_index, parallactic_angles, frequency_scaling, antenna_scaling, pointing_errors, ): """Wrapper for :func:`nb_zernike_dde`""" _, sources, times, ants, chans = coords.shape # ant, chan, corr_1, ..., corr_n, poly corr_shape = coeffs.shape[2:-1] npoly = coeffs.shape[-1] # Flatten correlation dimensions for numba function fcorrs = np.prod(corr_shape) ddes = np.empty((sources, times, ants, chans, fcorrs), coeffs.dtype) coeffs = coeffs.reshape((ants, chans, fcorrs, npoly)) noll_index = noll_index.reshape((ants, chans, fcorrs, npoly)) result = nb_zernike_dde( coords, coeffs, noll_index, ddes, parallactic_angles, frequency_scaling, antenna_scaling, pointing_errors, ) # Reshape to full correlation size return result.reshape((sources, times, ants, chans) + corr_shape)
_ZERNICKE_DOCSTRING = """ Computes Direction Dependent Effects by evaluating `Zernicke Polynomials <zernike_wiki_>`_ defined by coefficients ``coeffs`` and noll indexes ``noll_index`` at the specified coordinates ``coords``. Decomposition of a voxel beam cube into Zernicke polynomial coefficients can be achieved through the use of the eidos_ package. .. _zernike_wiki: https://en.wikipedia.org/wiki/Zernike_polynomials .. _eidos: https://github.com/kmbasad/eidos/ Parameters ---------- coords : :class:`numpy.ndarray` Float coordinates at which to evaluate the zernike polynomials. Has shape :code:`(3, source, time, ant, chan)`. The three components in the first dimension represent l, m and frequency coordinates, respectively. coeffs : :class:`numpy.ndarray` complex Zernicke polynomial coefficients. Has shape :code:`(ant, chan, corr_1, ..., corr_n, poly)` where ``poly`` is the number of polynomial coefficients and ``corr_1, ..., corr_n`` are a variable number of correlation dimensions. noll_index : :class:`numpy.ndarray` Noll index associated with each polynomial coefficient. Has shape :code:`(ant, chan, corr_1, ..., corr_n, poly)`. correlation dimensions. parallactic_angles : :class:`numpy.ndarray` Parallactic angle rotation. Has shape :code:`(time, ant)`. frequency_scaling : :class:`numpy.ndarray` The scaling of frequency of the beam. Has shape :code:`(chan,)`. antenna_scaling : :class:`numpy.ndarray` The antenna scaling. Has shape :code:`(ant, chan, 2)`. pointing_errors : :class:`numpy.ndarray` The pointing error. Has shape :code:`(time, ant, chan, 2)`. Returns ------- dde : :class:`numpy.ndarray` complex values with shape :code:`(source, time, ant, chan, corr_1, ..., corr_n)` """ zernike_dde.__doc__ = _ZERNICKE_DOCSTRING