Source code for africanus.rime.transform

#!/usr/bin/env python
# -*- coding: utf-8 -*-


import math

import numpy as np

from africanus.util.numba import jit


@jit(nopython=True, nogil=True, cache=True)
def _nb_transform_sources(
    lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, coords
):
    """
    numba implementation of
    :func:`~africanus.rime.transform_sources`
    """
    _, nsrc, ntime, na, nchan = coords.shape

    for t in range(ntime):
        for a in range(na):
            pa_sin = math.sin(parallactic_angles[t, a])
            pa_cos = math.cos(parallactic_angles[t, a])

            for s in range(nsrc):
                l, m = lm[s]

                # Rotate source coordinate by parallactic angle
                l = l * pa_cos - m * pa_sin  # noqa
                m = l * pa_sin + m * pa_cos

                # Add pointing errors
                l += pointing_errors[t, a, 0]  # noqa
                m += pointing_errors[t, a, 1]

                # Scale by antenna scaling factors
                for c in range(nchan):
                    coords[0, s, t, a, c] = l * antenna_scaling[a, c]
                    coords[1, s, t, a, c] = m * antenna_scaling[a, c]
                    coords[2, s, t, a, c] = frequency[c]

    return coords


[docs] def transform_sources( lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, dtype=None ): """ Creates beam sampling coordinates suitable for use in :func:`~africanus.rime.beam_cube_dde` by: 1. Rotating ``lm`` coordinates by the ``parallactic_angles`` 2. Adding ``pointing_errors`` 3. Scaling by ``antenna_scaling`` Parameters ---------- lm : :class:`numpy.ndarray` LM coordinates of shape :code:`(src,2)` in radians offset from the phase centre. parallactic_angles : :class:`numpy.ndarray` parallactic angles of shape :code:`(time, antenna)` in radians. pointing_errors : :class:`numpy.ndarray` LM pointing errors for each antenna at each timestep in radians. Has shape :code:`(time, antenna, 2)` antenna_scaling : :class:`numpy.ndarray` antenna scaling factor for each channel and each antenna. Has shape :code:`(antenna, chan)` frequency : :class:`numpy.ndarray` frequencies for each channel. Has shape :code:`(chan,)` dtype : :class:`numpy.dtype`, optional Numpy dtype of result array. Should be float32 or float64. Defaults to float64 Returns ------- coords : :class:`numpy.ndarray` coordinates of shape :code:`(3, src, time, antenna, chan)` where each coordinate component represents **l**, **m** and **frequency**, respectively. """ ntime, na = parallactic_angles.shape nsrc = lm.shape[0] assert (ntime, na, 2) == pointing_errors.shape nchan = antenna_scaling.shape[1] assert nchan == frequency.shape[0] dtype = np.float64 if dtype is None else dtype coords = np.empty((3, nsrc, ntime, na, nchan), dtype=dtype) return _nb_transform_sources( lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, coords )