Created
December 8, 2019 18:06
-
-
Save astrofrog/11f7f82b879b7d07bb8f7f555ef8d0d0 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment