Skip to content

Instantly share code, notes, and snippets.

@alexrudy
Last active August 29, 2015 14:17
Show Gist options
  • Save alexrudy/64b15a1e7e965b08e938 to your computer and use it in GitHub Desktop.
Save alexrudy/64b15a1e7e965b08e938 to your computer and use it in GitHub Desktop.
Flip OSIRIS Datacube axes.
#!/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())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment