casa_coordsys_ape14.py

import numpy as np

from astropy.wcs.wcsapi import BaseLowLevelWCS
from astropy.coordinates import SkyCoord, EarthLocation
from astropy import units as u

# The following dictionary maps axis names to UCD1+ terms from
# http://www.ivoa.net/documents/latest/UCDlist.html

NAMES_TO_UCD = {
    'Right Ascension': 'pos.eq.ra',
    'Declination': 'pos.eq.dec',
    'Longitude': 'pos.galactic.lon',
    'Latitude': 'pos.galactic.lat',
    'Frequency': 'em.freq',  # Frequency
}


class CASAWCS(BaseLowLevelWCS):
    """
    APE-14 compliant WCS interface to CASA coordsys objects.

    Parameters
    ----------
    cs : `casatools.coordsys.coordsys`
        The CASA coordsys object
    """

    def __init__(self, coordsys):
        self._coordsys = coordsys

    @property
    def pixel_n_dim(self):
        return self._coordsys.naxes()

    @property
    def world_n_dim(self):
        return self._coordsys.naxes()

    @property
    def world_axis_physical_types(self):
        return [NAMES_TO_UCD.get(name) for name in self._coordsys.names()]

    @property
    def world_axis_names(self):
        return self._coordsys.names()

    @property
    def world_axis_units(self):
        return self._coordsys.units()

    def pixel_to_world_values(self, *pixel_arrays):
        return list(self._coordsys.toworldmany(np.array(pixel_arrays))['numeric'])

    def array_index_to_world_values(self, *index_arrays):
        return list(self._coordsys.toworldmany(np.array(index_arrays)[::-1])['numeric'])

    def world_to_pixel_values(self, *world_arrays):
        return list(self._coordsys.topixelmany(np.array(world_arrays))['numeric'])

    def world_to_array_index_values(self, *world_arrays):
        pixel_arrays = list(self._coordsys.topixelmany(np.array(world_arrays))['numeric'])
        array_indices = tuple(np.asarray(np.floor(pixel + 0.5), dtype=np.int_) for pixel in pixel_arrays)
        return array_indices

    @property
    def world_axis_object_components(self):
        return self._get_components_and_classes()[0]

    @property
    def world_axis_object_classes(self):
        return self._get_components_and_classes()[1]

    def _get_components_and_classes(self):

        # The aim of this function is to return whatever is needed for
        # world_axis_object_components and world_axis_object_classes. It's easier
        # to figure it out in one go and then return the values and let the
        # properties return part of it.

        components = [None] * self.world_n_dim
        classes = {}

        # Let's start off by checking whether the WCS has a pair of celestial
        # components

        coordinate_type = self._coordsys.coordinatetype()
        names = self._coordsys.names()
        units = self._coordsys.units()

        if 'Direction' in coordinate_type:

            refcode = self._coordsys.referencecode('Direction')[0]

            if refcode == 'J2000':
                frame = 'fk5'
                lng = names.index('Right Ascension')
                lat = names.index('Declination')
            else:
                raise NotImplementedError(f"Did not recognize refcode={refcode}")

            kwargs = {}
            kwargs['frame'] = frame
            kwargs['unit'] = (units[lng], units[lat])

            classes['celestial'] = (SkyCoord, (), kwargs)

            components[lng] = ('celestial', 0, 'spherical.lon.' + units[lng])
            components[lat] = ('celestial', 1, 'spherical.lat.' + units[lng])

        # Fallback: for any remaining components that haven't been identified, just
        # return Quantity as the class to use

        for i in range(self.world_n_dim):
            if components[i] is None:
                name = names[i].lower()
                if name == '':
                    name = 'world'
                while name in classes:
                    name += "_"
                classes[name] = (u.Quantity, (), {'unit': units[i]})
                components[i] = (name, 0, 'value')

        return components, classes