Flip OSIRIS Datacube axes.

flipaxes.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""OSIRIS Datacubes have axes which are (WAVE, RA, DEC) in FITS files. Most programs
for reading and analyzing data cubes expect to have axes which are (RA, DEC, ...).
This script moves the (RA, DEC) axes to the front of the data cube.

The data is transformed using :func:`numpy.rollaxis` and the WCS solution is transformed
using :meth:`~astropy.wcs.WCS.reorient_celestial_first`. All other FITS header keywords
remain unchanged.

Note that FITS uses fortran-order for axes. Therefore, in Python (which uses c-order),
we flip the cube around so that the `WAVE` axis is first, followed by the celestial
axes.

"""

import textwrap
import argparse
import os
import sys

try:
    from astropy.io import fits
    from astropy.wcs import WCS
except ImportError as e:
    print("{}: {}".format(e.__class__.__name__, str(e)))
    print("\n".join(textwrap.wrap("This script requires astropy>=1.0. Try installing astropy via pip:"))+"\npip install astropy")
    sys.exit(1)

import numpy as np

_cli_epilog = textwrap.dedent("\n".join(textwrap.wrap(__doc__.split("\n\n")[0])))

def postfix(filename, postfix, sep="_"):
    """Add a postfix to a filename."""
    base, ext = os.path.splitext(filename)
    return sep.join([base, postfix]) + ext

def flip_celestial_first(HDU):
    """Re-arrange an HDU so that the celestial data is first, retunring a new HDU."""
    data = HDU.data.copy()
    wcs = WCS(HDU.header)
    
    if set(['WAVE','RA','DEC']) - set(wcs.axis_type_names):
        raise ValueError("More axes than expected! Not sure what to do: {}".format(wcs.axis_type_names))
    
    wave_axis = data.ndim - wcs.axis_type_names.index('WAVE') - 1
    rwcs = wcs.reorient_celestial_first()
    rdata = np.rollaxis(data, wave_axis)
    rHDU = HDU.__class__(rdata, header=HDU.header)
    rHDU.header.update(rwcs.to_header())
    return rHDU

def main():
    """Script entry point for fliping data cubes around."""
    parser = argparse.ArgumentParser(
        description="\n".join(textwrap.wrap("Flip data cube axes around so that the celestial dimensions are first.")),
        epilog=_cli_epilog,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('filename', help="name of the file to change axes.")
    parser.add_argument('-o', dest='postfix', help="append a string to the name of the output file.", default='flip', metavar='flip')
    parser.add_argument('-ext', dest='ext', help="FITS extension to flip.", default=0)
    parser.add_argument('--clobber', action='store_true', help="allow script to overwrite the input file.")
    opt = parser.parse_args()
    filename = opt.filename
    output = postfix(filename, opt.postfix)
    
    with fits.open(filename) as HDUs:
        HDUs[opt.ext] = flip_celestial_first(HDUs[opt.ext])
        HDUs[opt.ext].writeto(output, clobber=((output != filename) or opt.clobber))
    
    print("Flipped data cube saved to {}".format(output))
    return 0


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