Merge pull request #127 from lpsinger/from_header

Add HEALPix.from_header

astropy_healpix/high_level.py

@@ -7,10 +7,11 @@
 from astropy.coordinates.representation import UnitSphericalRepresentation
 
 from .core import (nside_to_pixel_area, nside_to_pixel_resolution,
-                   nside_to_npix, healpix_to_lonlat, lonlat_to_healpix,
-                   bilinear_interpolation_weights, interpolate_bilinear_lonlat,
-                   ring_to_nested, nested_to_ring, healpix_cone_search,
-                   boundaries_lonlat, neighbours)
+                   nside_to_npix, npix_to_nside, healpix_to_lonlat,
+                   lonlat_to_healpix, bilinear_interpolation_weights,
+                   interpolate_bilinear_lonlat, ring_to_nested, nested_to_ring,
+                   healpix_cone_search, boundaries_lonlat, neighbours)
+from .utils import parse_input_healpix_data
 
 __all__ = ['HEALPix']
 
@@ -52,6 +53,40 @@ def __init__(self, nside=None, order='ring', frame=None):
         self.order = order
         self.frame = frame
 
+    @classmethod
+    def from_header(cls, input_data, field=0, hdu_in=None, nested=None):
+        """
+        Parameters
+        ----------
+        input_data : str or `~astropy.io.fits.TableHDU` or `~astropy.io.fits.BinTableHDU` or tuple
+            The input data to reproject. This can be:
+
+                * The name of a HEALPIX FITS file
+                * A `~astropy.io.fits.TableHDU` or `~astropy.io.fits.BinTableHDU`
+                  instance
+                * A tuple where the first element is a `~numpy.ndarray` and the
+                  second element is a `~astropy.coordinates.BaseCoordinateFrame`
+                  instance or a string alias for a coordinate frame.
+
+        hdu_in : int or str, optional
+            If ``input_data`` is a FITS file, specifies the HDU to use.
+            (the default HDU for HEALPIX data is 1, unlike with image files where
+            it is generally 0)
+        nested : bool, optional
+            The order of the healpix_data, either nested (True) or ring (False).
+            If a FITS file is passed in, this is determined from the header.
+
+        Returns
+        -------
+        healpix : `~astropy_healpix.HEALPix`
+            A HEALPix pixellization corresponding to the input data.
+        """
+        array_in, frame, nested = parse_input_healpix_data(
+            input_data, field=field, hdu_in=hdu_in, nested=nested)
+        nside = npix_to_nside(len(array_in))
+        order = 'nested' if nested else 'ring'
+        return cls(nside=nside, order=order, frame=frame)
+
     @property
     def pixel_area(self):
         """

astropy_healpix/tests/test_high_level.py

@@ -8,6 +8,7 @@
 
 from astropy import units as u
 from astropy.coordinates import Longitude, Latitude, Galactic, SkyCoord
+from astropy.io.fits import Header
 
 from ..high_level import HEALPix
 
@@ -108,6 +109,25 @@ class TestCelestialHEALPix:
     def setup_class(self):
         self.pix = HEALPix(nside=256, order='nested', frame=Galactic())
 
+    def test_healpix_from_header(self):
+        """Test instantiation from a FITS header.
+
+        Notes
+        -----
+        We don't need to test all possible options, because
+        :meth:`~astropy_healpix.HEALPix.from_header` is just a wrapper around
+        :meth:`~astropy_healpix.utils.parse_input_healpix_data`, which is
+        tested exhaustively in :mod:`~astropy_healpix.tests.test_utils`.
+        """
+
+        pix = HEALPix.from_header(
+            (np.empty(self.pix.npix), 'G'),
+            nested=self.pix.order == 'nested')
+
+        assert pix.nside == self.pix.nside
+        assert type(pix.frame) == type(self.pix.frame)
+        assert pix.order == self.pix.order
+
     def test_healpix_to_skycoord(self):
         coord = self.pix.healpix_to_skycoord([1, 2, 3])
 

astropy_healpix/tests/test_utils.py

@@ -0,0 +1,56 @@
+import numpy as np
+
+import pytest
+from astropy.coordinates import FK5, Galactic
+from astropy.io import fits
+
+from ..utils import parse_coord_system, parse_input_healpix_data
+
+
+def test_parse_coord_system():
+
+    frame = parse_coord_system(Galactic())
+    assert isinstance(frame, Galactic)
+
+    frame = parse_coord_system('fk5')
+    assert isinstance(frame, FK5)
+
+    with pytest.raises(ValueError) as exc:
+        frame = parse_coord_system('e')
+    assert exc.value.args[0] == "Ecliptic coordinate frame not yet supported"
+
+    frame = parse_coord_system('g')
+    assert isinstance(frame, Galactic)
+
+    with pytest.raises(ValueError) as exc:
+        frame = parse_coord_system('spam')
+    assert exc.value.args[0] == "Could not determine frame for system=spam"
+
+
+def test_parse_input_healpix_data(tmpdir):
+
+    data = np.arange(3072)
+
+    col = fits.Column(array=data, name='flux', format="E")
+    hdu = fits.BinTableHDU.from_columns([col])
+    hdu.header['NSIDE'] = 512
+    hdu.header['COORDSYS'] = "G"
+
+    # As HDU
+    array, coordinate_system, nested = parse_input_healpix_data(hdu)
+    np.testing.assert_allclose(array, data)
+
+    # As filename
+    filename = tmpdir.join('test.fits').strpath
+    hdu.writeto(filename)
+    array, coordinate_system, nested = parse_input_healpix_data(filename)
+    np.testing.assert_allclose(array, data)
+
+    # As array
+    array, coordinate_system, nested = parse_input_healpix_data((data, "galactic"))
+    np.testing.assert_allclose(array, data)
+
+    # Invalid
+    with pytest.raises(TypeError) as exc:
+        parse_input_healpix_data(data)
+    assert exc.value.args[0] == "input_data should either be an HDU object or a tuple of (array, frame)"

astropy_healpix/utils.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+from astropy.io import fits
+from astropy.io.fits import TableHDU, BinTableHDU
+from astropy.extern import six
+from astropy.coordinates import BaseCoordinateFrame, frame_transform_graph, Galactic, ICRS
+
+FRAMES = {
+    'g': Galactic(),
+    'c': ICRS()
+}
+
+
+def parse_coord_system(system):
+    if isinstance(system, BaseCoordinateFrame):
+        return system
+    elif isinstance(system, six.string_types):
+        system = system.lower()
+        if system == 'e':
+            raise ValueError("Ecliptic coordinate frame not yet supported")
+        elif system in FRAMES:
+            return FRAMES[system]
+        else:
+            system_new = frame_transform_graph.lookup_name(system)
+            if system_new is None:
+                raise ValueError("Could not determine frame for system={0}".format(system))
+            else:
+                return system_new()
+
+
+def parse_input_healpix_data(input_data, field=0, hdu_in=None, nested=None):
+    """
+    Parse input HEALPIX data to return a Numpy array and coordinate frame object.
+    """
+
+    if isinstance(input_data, (TableHDU, BinTableHDU)):
+        data = input_data.data
+        header = input_data.header
+        coordinate_system_in = parse_coord_system(header['COORDSYS'])
+        array_in = data[data.columns[field].name].ravel()
+        if 'ORDERING' in header:
+            nested = header['ORDERING'].lower() == 'nested'
+    elif isinstance(input_data, six.string_types):
+        hdu = fits.open(input_data)[hdu_in or 1]
+        return parse_input_healpix_data(hdu, field=field)
+    elif isinstance(input_data, tuple) and isinstance(input_data[0], np.ndarray):
+        array_in = input_data[0]
+        coordinate_system_in = parse_coord_system(input_data[1])
+    else:
+        raise TypeError("input_data should either be an HDU object or a tuple of (array, frame)")
+
+    return array_in, coordinate_system_in, nested