Source code for africanus.model.spi.component_spi

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


import numpy as np

from africanus.util.docs import DocstringTemplate
from africanus.util.numba import jit


@jit(nopython=True, nogil=True, cache=True)
def _fit_spi_components_impl(data, weights, freqs, freq0, out,
                             jac, ncomps, nfreqs, tol, maxiter):
    w = freqs/freq0
    for comp in range(ncomps):
        eps = 1.0
        k = 0
        alphak = out[0, comp]
        i0k = out[2, comp]
        while eps > tol and k < maxiter:
            alphap = alphak
            i0p = i0k
            jac[1, :] = w**alphak
            model = i0k*jac[1, :]
            jac[0, :] = model * np.log(w)
            residual = data[comp] - model
            lik = 0.0
            hess00 = 0.0
            hess01 = 0.0
            hess11 = 0.0
            jr0 = 0.0
            jr1 = 0.0
            for v in range(nfreqs):
                lik += residual[v] * weights[v] * residual[v]
                jr0 += jac[0, v] * weights[v] * residual[v]
                jr1 += jac[1, v] * weights[v] * residual[v]
                hess00 += jac[0, v] * weights[v] * jac[0, v]
                hess01 += jac[0, v] * weights[v] * jac[1, v]
                hess11 += jac[1, v] * weights[v] * jac[1, v]
            det = hess00 * hess11 - hess01**2
            alphak = alphap + (hess11 * jr0 - hess01 * jr1)/det
            i0k = i0p + (-hess01 * jr0 + hess00 * jr1)/det
            eps = np.maximum(np.abs(alphak - alphap), np.abs(i0k - i0p))
            k += 1
        if k == maxiter:
            print("Warning - max iterations exceeded for component ", comp)
        out[0, comp] = alphak
        out[1, comp] = hess11/det
        out[2, comp] = i0k
        out[3, comp] = hess00/det
    return out


[docs]def fit_spi_components(data, weights, freqs, freq0, alphai=None, I0i=None, tol=1e-4, maxiter=100): ncomps, nfreqs = data.shape jac = np.zeros((2, nfreqs), dtype=data.dtype) out = np.zeros((4, ncomps), dtype=data.dtype) if alphai is not None: out[0, :] = alphai else: out[0, :] = -0.7 * np.ones(ncomps, dtype=data.dtype) if I0i is not None: out[2, :] = I0i else: tmp = np.abs(freqs - freq0) ref_freq_idx = np.argwhere(tmp == tmp.min()).squeeze() out[2, :] = data[:, ref_freq_idx] return _fit_spi_components_impl(data, weights, freqs, freq0, out, jac, ncomps, nfreqs, tol, maxiter)
SPI_DOCSTRING = DocstringTemplate( r""" Computes the spectral indices and the intensity at the reference frequency of a spectral index model: .. math:: I(\nu) = I(\nu_0) \left( \frac{\nu}{\nu_0} \right) ^ \alpha Parameters ---------- data : $(array_type) array of shape :code:`(comps, chan)` The noisy data as a function of frequency. weights : $(array_type) array of shape :code:`(chan,)` Inverse of variance on each frequency axis. freqs : $(array_type) frequencies of shape :code:`(chan,)` freq0 : float Reference frequency alphai : $(array_type), optional array of shape :code:`(comps,)` Initial guess for the alphas. Defaults to -0.7. I0i : $(array_type), optional array of shape :code:`(comps,)` Initial guess for the intensities at the reference frequency. Defaults to 1.0. tol : float, optional Solver absolute tolerance (optional). Defaults to 1e-6. maxiter : int, optional Solver maximum iterations (optional). Defaults to 100. dtype : np.dtype, optional Datatype of result. Should be either np.float32 or np.float64. Defaults to np.float64. Returns ------- out : $(array_type) array of shape :code:`(4, comps)` The fitted components arranged as [alphas, alphavars, I0s, I0vars] """) try: fit_spi_components.__doc__ = SPI_DOCSTRING.substitute( array_type=":class:`numpy.ndarray`") except AttributeError: pass