# 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)
if n < 0 or m < 0 or abs(m) > n:
raise ValueError("m and n values are incorrect.")
for k in range((n - m) / 2 + 1):
radial_component += pre_fac(k, n, m) * rho ** (n - 2.0 * k)

@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)

@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

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
```