Strip extra axes from FITS file.

strip_naxis.py

import numpy as np
from astropy.io import fits

def strip_naxis(fitsimage, outname=None, strip_wcs=True):
       
    """Strip additional axes from FITS file.

    Strips axes 3 and/or 4 if present.
    
    Parameters
    ----------
    fitsimage : str or astropy.io.fits.HDUList
    
    """

    def get_hdulist(fitsimage):
    """Get HDUList object from filepath or form HDUList object."""

    if isinstance(fitsimage, str):
        hdulist = fits.open(fitsimage)
        opened = True
    elif isinstance(fitsimage, fits.HDUList):
        hdulist = fitsimage
        opened = False
    else:
        raise ValueError("FITS image must be a filename/path or an " \
                         "astropy.io.fits.HDUList object.")

    return hdulist, opened
    
    
    image, opened = get_hdulist(fitsimage)

    if image[0].header["NAXIS"] == 3:

        image[0].data = image[0].data[0, :, :]


    elif image[0].header["NAXIS"] == 4:

        image[0].data = image[0].data[0, 0, :, :]

    else:

        pass

    if strip_wcs:

        try: image[0].header["FREQ"] = image[0].header["CRVAL3"]
        except KeyError: pass

        for i in ["3", "4"]:
            for j in ["CRVAL"+i, "CTYPE"+i, "CDELT"+i, "CRPIX"+i, "NAXIS"+i, "CUNIT"+i, "CROTA"+i]:
                try: del image[0].header[j]
                except KeyError: pass

    image[0].header["NAXIS"] = 2

    if outname is None:
        outname = fitsimage

    image.writeto(outname, clobber=True)
    if opened:
        image.close()