Flip Keck/OSIRIS IFU Cube axes from (Lambda, RA, Declination) to (RA, Declination, Lambda)

flip_osiris_axes.py

#!/usr/bin/env python

"""
flip_osiris_axes.py: Convert a Keck/OSIRIS IFU datacube from axes of
(Lambda, RA, Dec) to (RA, Dec, Lambda) so that it can be read by, e.g.,
QFitsView.

OSIRIS, for a frankly mystifying reason, (related to FITS FORTRAN axes
orders, I think), defaults to axes of (Lambda, RA, Dec.). This script
converts the cube to (Ra, Dec, Lambda), which requires both a reorder of the
data axes and a reorientation of the WCS axes. This is accomplished
with numpy.rollaxis and astropy.wcs.WCS.reorient_celestial_first().

Uses some code from @alexrudy.
"""

__author__ = "Grant R. Tremblay, Harvard-Smithsonian Center for Astrophysics"
__licence__ = "MIT"

import os
import sys
import argparse
import warnings

try:
    from astropy.io import fits
    from astropy.wcs import WCS
    from astropy.utils.exceptions import AstropyWarning
except ImportError:
    print("Looks like Astropy isn't installed on your Python stack.")
    print("See installation instructions here: www.astropy.org")
    sys.exit(1)
try:
    import numpy as np
except ImportError:
    print("Looks like Numpy isn't installed on your Python stack.")
    print("See installation instructions here: www.numpy.org")
    sys.exit(1)

def parse_args():
    """Parse and return command line arguments"""

    parser = argparse.ArgumentParser()

    parser.add_argument('filename',
                        help="The input (Lamda, RA, Dec) FITS cube")
    parser.add_argument('-ext', dest="ext", help="Extension to flip",
                        default=0)

    parser.add_argument('-o', required=True, dest='outfile', help="Output file postfix")

    parser.add_argument('--clobber', action='store_true',
                        help="Allow overwrite?")

    args = parser.parse_args()

    return args

def flipcube(HDU):
    """Flip the cube from (Lambda, RA, Dec) to (RA, Dec, Lambda)"""

    # Ensure you're working on a copy of the data
    data = HDU.data.copy()
    print("Input data copied.")

    wcs = WCS(HDU.header)
    print("WCS info read.")

    wave_axis = data.ndim - wcs.axis_type_names.index('WAVE') - 1
    print("Wavelength axis currently assigned to: {}".format(wave_axis))

    # Reorient the WCS such that the celestial axes are first, built-in method
    flippedWCS = wcs.reorient_celestial_first()
    print("WCS reoriented to (RA, Dec, Lambda)")

    flippeddata = np.rollaxis(data, wave_axis)
    flippedHDU = HDU.__class__(flippeddata, header=HDU.header)

    flippedHDU.header.update(flippedWCS.to_header())
    return flippedHDU


def main():
    args = parse_args()

    filename = args.filename
    ext = args.ext
    outfile = args.outfile
    clobber = args.clobber

    with fits.open(filename) as HDUs:
        warnings.simplefilter('ignore', category=AstropyWarning)
        HDUs[args.ext] = flipcube(HDUs[args.ext])
        HDUs[args.ext].writeto(outfile, clobber=((outfile != filename) or clobber))

    print("Flipped cube saved to {}".format(outfile))
    return 0

if __name__ == '__main__':
    sys.exit(main())