Source code for meersolar.crystalball.wsclean

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

from collections import namedtuple
import logging

from africanus.model.wsclean.file_model import load
from astropy.coordinates import SkyCoord
import numpy as np


log = logging.getLogger(__name__)

WSCleanModel = namedtuple("WSCleanModel", ["source_type", "radec",
                                           "flux", "spi", "ref_freq",
                                           "log_poly", "gauss_shape"])


[docs] def import_from_wsclean(wsclean_comp_list, include_regions=[], point_only=False, num=None): """ Imports sources from wsclean, sorted from brightest to faintest. If ``include_regions`` is specified only those sources within are returned. Similarly if ``num`` is specified, ``num`` brightest sources are returned. Parameters ---------- wsclean_comp_list : str wsclean component list file. include_regions : List[:class:`regions.CircleSkyRegion`] List of valid region's. Only sources within these regions will be loaded. Defaults to ``[]`` in which case, all sources are loaded. point_only : bool, optional Only include point sources. Defaults to False num : integer, optional Number of sources to include. If ``None`` all sources are returned. Returns ------- source_type : :class:`numpy.ndarray` Source types of shape :code:`(source,)` radec : :class:`numpy.ndarray` Source radec coordinates of shape :code:`(source, 2)`. stokes : :class:`numpy.ndarray` Source stokes paramters of shape :code:`(source, 4)`. spectral_index : :class:`numpy.ndarray` Spectral index of shape :code:`(source, spi)` ref_freq : :class:`numpy.ndarray` Reference frequency of shape :code:`(source,)` log_si : bool Boolean indicated whether to use the logarithmic spectral index. gauss_shape : :class:`numpy.ndarray` Gaussian shape parameters of shape :code:`(source, 3)`. The three components are MajorAxis, MinorAxis and Orientation, respectively. """ wsclean_comps = {column: np.asarray(values) for column, values in load(wsclean_comp_list)} if np.unique(wsclean_comps['LogarithmicSI']).shape[0] > 1: raise ValueError("Can't handle Mixed log and ordinary polynomial " "coefficients in '%s'" % wsclean_comp_list) # Sort components by descending flux sort_index = np.ascontiguousarray(np.argsort(wsclean_comps['I'])[::-1]) # re-sort all entries using this wsclean_comps = {key: value[sort_index] for key, value in wsclean_comps.items()} log.info('%s contains %d components', wsclean_comp_list, len(wsclean_comps['Type'])) # make mask of sources to include include = np.ones_like(wsclean_comps['Type'], bool) if include_regions: include[:] = False # NB: regions is *supposed* to have a sensible "contains" interface, # but it doesn't work as of Mar 2019. So hacking # a kludge for circular regions for now from regions import CircleSkyRegion if not all([type(reg) is CircleSkyRegion for reg in include_regions]): raise ValueError('Only circular DS( regions supported for now') coord = SkyCoord(wsclean_comps['Ra'], wsclean_comps['Dec'], unit="rad", frame=include_regions[0].center.frame) include = coord.separation( include_regions[0].center) <= include_regions[0].radius for reg in include_regions[1:]: include |= coord.separation(reg.center) <= reg.radius log.info("%d of which fall within the %d inclusive regions", include.sum(), len(include_regions)) # select points if point_only: include &= (wsclean_comps['Type'] == 'POINT') log.info("%d of which are point sources", include.sum()) # apply filters to component list wsclean_comps = {key: value[include] for key, value in wsclean_comps.items()} # select limited number if asked if num is not None: wsclean_comps = {key: value[:num] for key, value in wsclean_comps.items()} log.info("Selecting up %d brightest sources", num) # print if small subset if (num is not None and num < 100) or include_regions: flux_types = zip(wsclean_comps['I'], wsclean_comps['Type']) it = enumerate(sorted(flux_types, reverse=True)) for i, (flux, src_type) in it: log.info('%d: %s %s Jy', i, src_type, flux) log.info('Total flux of %d selected components is %f Jy', len(wsclean_comps['I']), wsclean_comps['I'].sum()) # Create radec array radec = np.concatenate((wsclean_comps['Ra'][:, None], wsclean_comps['Dec'][:, None]), axis=1) # Create gaussian shapes gauss_shape = np.stack((wsclean_comps['MajorAxis'], wsclean_comps['MinorAxis'], wsclean_comps['Orientation']), axis=-1) return WSCleanModel(wsclean_comps['Type'], radec, wsclean_comps['I'], wsclean_comps['SpectralIndex'], wsclean_comps['ReferenceFrequency'], wsclean_comps['LogarithmicSI'], gauss_shape)