# -*- 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`"
)