Source code for africanus.rime.wsclean_predict

# -*- coding: utf-8 -*-

import numpy as np

from africanus.constants import two_pi_over_c, c as lightspeed
from africanus.util.docs import DocstringTemplate
from africanus.util.numba import JIT_OPTIONS, overload, njit
from africanus.model.wsclean.spec_model import spectra


@njit(**JIT_OPTIONS)
def wsclean_predict_main(uvw, lm, source_type, gauss_shape, frequency, spectrum, dtype):
    fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0))
    fwhminv = 1.0 / fwhm
    gauss_scale = fwhminv * np.sqrt(2.0) * np.pi / lightspeed

    nrow = uvw.shape[0]
    nchan = frequency.shape[0]
    ncorr = 1

    nsrc = spectrum.shape[0]
    n1 = lm.dtype.type(1)

    scaled_freq = frequency * frequency.dtype.type(gauss_scale)

    vis = np.zeros((nrow, nchan, ncorr), dtype=dtype)

    for s in range(nsrc):
        l = lm[s, 0]  # noqa
        m = lm[s, 1]
        n = np.sqrt(n1 - l * l - m * m) - n1

        if source_type[s] == "POINT":
            for r in range(nrow):
                u = uvw[r, 0]
                v = uvw[r, 1]
                w = uvw[r, 2]

                real_phase = two_pi_over_c * (u * l + v * m + w * n)

                for f in range(nchan):
                    p = real_phase * frequency[f]
                    re = np.cos(p) * spectrum[s, f]
                    im = np.sin(p) * spectrum[s, f]

                    vis[r, f, 0] += re + im * 1j
        elif source_type[s] == "GAUSSIAN":
            emaj, emin, angle = gauss_shape[s]

            # Convert to l-projection, m-projection, ratio
            el = emaj * np.sin(angle)
            em = emaj * np.cos(angle)
            er = emin / (1.0 if emaj == 0.0 else emaj)

            for r in range(nrow):
                u = uvw[r, 0]
                v = uvw[r, 1]
                w = uvw[r, 2]

                # Compute phase term
                real_phase = two_pi_over_c * (u * l + v * m + w * n)

                # Gaussian shape term bits
                u1 = (u * em - v * el) * er
                v1 = u * el + v * em

                for f in range(nchan):
                    p = real_phase * frequency[f]
                    re = np.cos(p) * spectrum[s, f]
                    im = np.sin(p) * spectrum[s, f]

                    # Calculate gaussian shape component and multiply in
                    fu1 = u1 * scaled_freq[f]
                    fv1 = v1 * scaled_freq[f]
                    shape = np.exp(-(fu1 * fu1 + fv1 * fv1))
                    re *= shape
                    im *= shape

                    vis[r, f, 0] += re + im * 1j
        else:
            raise ValueError("source_type must be " "POINT or GAUSSIAN")

    return vis


[docs] @njit(**JIT_OPTIONS) def wsclean_predict( uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency ): return wsclean_predict_impl( uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency )
def wsclean_predict_impl( uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency ): raise NotImplementedError @overload(wsclean_predict_impl, jit_options=JIT_OPTIONS) def nb_wsclean_predict( uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency ): arg_dtypes = tuple( np.dtype(a.dtype.name) for a in (uvw, lm, flux, coeffs, ref_freq, frequency) ) dtype = np.result_type(np.complex64, *arg_dtypes) def impl( uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency ): spectrum = spectra(flux, coeffs, log_poly, ref_freq, frequency) return wsclean_predict_main( uvw, lm, source_type, gauss_shape, frequency, spectrum, dtype ) return impl WSCLEAN_PREDICT_DOCS = DocstringTemplate( """ Predict visibilities from a `WSClean sky model <https://sourceforge.net/p/wsclean/wiki/ComponentList/>`_. Parameters ---------- uvw : $(array_type) UVW coordinates of shape :code:`(row, 3)` lm : $(array_type) Source LM coordinates of shape :code:`(source, 2)`, in radians. Derived from the ``Ra`` and ``Dec`` fields. source_type : $(array_type) Strings defining the source type of shape :code:`(source,)`. Should be either ``"POINT"`` or ``"GAUSSIAN"``. Contains the ``Type`` field. flux : $(array_type) Source flux of shape :code:`(source,)`. Contains the ``I`` field. coeffs : $(array_type) Source Polynomial coefficients of shape :code:`(source, coeffs)`. Contains the ``SpectralIndex`` field. log_poly : $(array_type) Source polynomial type of shape :code:`(source,)`. If True, logarithmic polynomials are used. If False, standard polynomials are used. Contains the ``LogarithmicSI`` field. ref_freq : $(array_type) Source Reference frequency of shape :code:`(source,)`. Contains the ``ReferenceFrequency`` field. gauss_shape : $(array_type) Gaussian shape parameters of shape :code:`(source, 3)` used when the corresponding ``source_type`` is ``"GAUSSIAN"``. The 3 components should contain the ``MajorAxis``, ``MinorAxis`` and ``Orientation`` fields in radians, respectively. frequency : $(array_type) Frequency of shape :code:`(chan,)`. Returns ------- visibilities : $(array_type) Complex visibilities of shape :code:`(row, chan, 1)` """ ) wsclean_predict.__doc__ = WSCLEAN_PREDICT_DOCS.substitute( array_type=":class:`numpy.ndarray`" )