Convolve a FITS data cube with a Gaussian beam

convolve.py

#!/usr/bin/env python
# convolve - convolve a FITS cube with a gaussian beam

import sys
import os
import getopt
import gzip
from numpy import *
from astropy.convolution import convolve_fft, Gaussian2DKernel
from astropy.io import fits as pyfits

VERSION = "0.2"

global filein, hpbw

def usage():
    """Display usage"""

    print("""Usage: convolve [fitsfile] [beamsize]
                    convolve [batchfile]
Compute the convolution of a fits image or cube with a Gaussian beam.

Arguments:
   fitsfile           fits image or cube to convolve
   beamsize           HPBW of the gaussian beam in arcsec

   batchfile          file containing the fits file to convolve
                      and the beamsizes (one per line).
Options:
   -h,--help          display this help
   -V,--version       display plroute version information

   -s,--spectra=x,y   if the input file is a cube, do not create a
                      convolved cube but only extract a spectra
                      from it at a given offset, expressed in
                      pixels from the cube center
   -d,--delete        delete original file
""")

def version():
    """Display version number"""

    print("This is convolve, version %s" % VERSION)
    print("""Copyright (c) 2005-2014 Sebastien Maret

This is free software. You may redistribute copies of it under the terms
of the GNU General Public License. There is NO WARRANTY, to the extent
permitted by law.""")

def readfits (filein):
    """Read a fits file and return a tuple of arrays
    containing the data, RA, DEC and channels."""

    fits = pyfits.open (filein)

    # Get the RA, DEC and channel values from the header
    # The axis 1 is the RA, the axis 2 is the DEC and the
    # axis 3 is the velocity channels

    header = fits[0].header    # Primary HDU header
    naxis1 = header['naxis1']  # Number of values in each axis
    naxis2 = header['naxis2']  # naxis1 and naxis2 should be equal
    naxis3 = header['naxis3']

    cdelt1 = header['cdelt1']  # Difference between two pixel
    cdelt2 = header['cdelt2']  # cdelt1 and cdelt2 should be opposite
    cdelt3 = header['cdelt3']  # cdelt1 negative

    crpix1 = header['crpix1']  # Reference pixel
    crpix2 = header['crpix2']  # All of them should be 0.
    crpix3 = header['crpix3']

    crval1 = header['crval1']  # Value at the reference pixel
    crval2 = header['crval2']
    crval3 = header['crval3']

    # Compute the values of RA and DEC

    ra = arange(naxis1, dtype = double)
    dec = arange(naxis2, dtype = double)
    veloc = arange(naxis3, dtype = double)
    for i in range(naxis1):
        ra[i] = (i + 1 - crpix1) * cdelt1 + crval1
    for i in range(naxis2):
        dec[i] = (i + 1 - crpix2) * cdelt2 + crval2
    for i in range(naxis3):
        veloc[i] = (i + 1 - crpix3) * cdelt3 + crval3

    # The values of RA and DEC are in degrees in the fits file
    # Velocities are in m/s. We convert them in arcsec and km/s

    ra = ra * 3600
    dec = dec * 3600
    veloc = veloc / 1e3

    # Fix keyword for velocity
    header['ctype3'] = 'VELO-LSR'

    # Now read the brightness temperatures
    tb = fits[0].data

    return (ra, dec, veloc, tb, header)

def writefits (fileout, tb, header):
    """Write the convolved datacube in a fits file."""

    hdu = pyfits.PrimaryHDU (tb, header = header)
    hdu.writeto (fileout)

def main ():
    """Main program for convolve"""

    extract_spectra = False
    delete = False

    # Parse command line options and arguments
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hvs:d",
                                   ["help", "version", "spectra=","delete"])
    except getopt.GetoptError:
        print("convolve: error: unknown option", file=sys.stderr)
        print("Type 'convolve --help' for more information", file=sys.stderr)
        sys.exit(1)

    for opt, arg in opts:
        if opt in ("-h", "--help") :
            usage()
            sys.exit(0)
        if opt in ("-v", "--version"):
            version()
            sys.exit(0)
        if opt in ("-s", "--spectra"):
            try:
                extract_spectra = True
                offset = arg.split(",")
                offset = list(map(int, offset))
            except:
                print("convolve: error: incorrect argument for option spectra", file=sys.stderr)
                print("Type 'convolve --help' for more information", file=sys.stderr)
                sys.exit(1)
        if opt in ("-d", "--delete"):
            delete = True

    if len(args) == 1:
        batchfile = args[0]
        fitsfile = []
        hpbw = []
    elif len(args) == 2:
        batchfile = None
        fitsfile = [args[0]]
        hpbw = [float(args[1])]
    else:
        print("convolve: error: incorrect number of arguments", file=sys.stderr)
        print("Type 'convolve --help' for more information", file=sys.stderr)
        sys.exit(1)

    # If we are in batch mode, read the batch file
    if batchfile:
        bf = open (batchfile)
        for line in bf:
            if line[0] == "#":
                continue
            f, h = line.split()
            fitsfile.append (f)
            hpbw.append (float (h))

    # Now convolve each file
    for i in range (len (fitsfile)):

        sys.stdout.write ("Reading %s\n" % fitsfile[i])

        (ra, dec, veloc, tb, header) = readfits (fitsfile[i])

        # Get the normalized beam power pattern
        resol = ra[0] - ra[1]
        stdev = hpbw / (2 * sqrt (2 * log(2)))
        beam = Gaussian2DKernel (stddev = stdev / resol)

        # FixMe: We should check that the fits image is big enough (i.e. bigger
        # than the kernel) before doing the convolution, and pad the image
        # if not. Having a beam size and image size that are a power of two
        # would speed up convolution computation (if FFT is used)

        # Convolve the brightness temperature by the beam pattern
        # in each channel

        dim_ra = len (ra)
        dim_dec = len (dec)
        dim_veloc = len (veloc)

        ta = zeros (dim_ra * dim_dec * dim_veloc, dtype = double)
        ta = reshape (ta, (dim_veloc, dim_ra, dim_dec))

        sys.stdout.write ("Convoluting with Gaussian beam... ")
        sys.stdout.flush ()

        # Compute the convolution using the fft
        for k in range(dim_veloc):
            ta[k,:,:] = convolve_fft(tb[k,:,:], beam)

        sys.stdout.write ("done\n")

        # Extract a spectra from the convolved cube and update the
        # header accordingly
        if extract_spectra:
            try:
                sys.stdout.write ("Extracting spectra at offset (%i,%i) from the cube center\n" % \
                                      (offset[0], offset[1]))
                ra_index = (dim_ra - 1) / 2 + int(offset[0])
                dec_index = (dim_dec -1) / 2 + int(offset[1])
                ta_extract = ta[:,dec_index,ra_index]
                ta = ta_extract

                header['ctype1'] = header['ctype3']
                header['cdelt1'] = header['cdelt3']
                header['crval1'] = header['crval3']
                header['crpix1'] = header['crpix3']
                del header['ctype2']
                del header['cdelt2']
                del header['crval2']
                del header['crpix2']
                del header['ctype3']
                del header['cdelt3']
                del header['crval3']
                del header['crpix3']

            except IndexError:
                print("convolve: error: offset is out of cube bounds", file=sys.stderr)
                sys.exit(1)

        # Write the results in the fits file

        # astropy.io.fits does not support writing gziped fits files, so if the
        # original file was gziped, we first create a temporary fits
        # file, and then we gzip it.

        if fitsfile[i][-7:] == 'fits.gz':
            fileout = fitsfile[i][:-8] + '_conv.fits'
            gziped = True
        else:
            fileout = fitsfile[i][:-5] + '_conv.fits'
            gziped = False
        sys.stdout.write ( "Writing results in %s\n" % fileout)
        if os.path.exists (fileout):
            os.unlink (fileout)
        writefits(fileout, ta, header)
        if gziped:
            sys.stdout.write ( "Compressing file... ")
            sys.stdout.flush ()
            gzfileout = fileout + '.gz'
            if os.path.exists (gzfileout):
                os.unlink (gzfileout)
            gzf = gzip.open (gzfileout, 'w')
            f = open (fileout, 'r')
            gzf.write (f.read ())
            f.close ()
            gzf.close ()
            sys.stdout.write ( "done\n")
            os.unlink (fileout)

        # Delete original file, if requested

        if delete:
            os.unlink(fitsfile[i])

if __name__ == "__main__":
    main()