Prepare a set of file for use with BRATS. Requires MIRIAD.

brats_prep.py

#! /usr/bin/env python

from __future__ import print_function, division

import os
import shutil

from subprocess import Popen
from argparse import ArgumentParser
from astropy.io import fits


def main():
    """
    """

    ps = ArgumentParser()

    ps.add_argument("-i", "--images", nargs="*")
    ps.add_argument("-I", "--imsize", "--size", dest="imsize", default=1000, type=int)
    ps.add_argument("-c", "--coords", type=float, nargs=2)
    ps.add_argument("-s", "--source", "--name", dest="source", type=str,
                    default="SOURCE")

    args = ps.parse_args()

    images = []
    frequencies = []
    bmaj = 0.
    bmin = 0.
    bpa = 0.

    # cdelt1 = 1.e9
    # cdelt2 = 1.e9
    cdelt1 = 0.
    cdelt2 = 0.

    crval1 = args.coords[0]
    crval2 = args.coords[1]
    crpix1 = args.imsize // 2
    crpix2 = args.imsize // 2

    for image in args.images:

        f = fits.open(image, mode="update")

        # Find pixel information:
        if "CD1_1" in f[0].header.keys():
            cd1 = f[0].header["CD1_1"]
            cd2 = f[0].header["CD2_2"]
        elif "CDELT1" in f[0].header.keys():
            cd1 = f[0].header["CDELT1"]
            cd2 = f[0].header["CDELT2"]
        else:
            raise ValueError("No pixel information found for {}".format(image))

        # Use the smallest pixel as our pixel scale:
        if abs(cd1) > abs(cdelt1):
            cdelt1 = cd1
            cdelt2 = cd2

        # Find beam information and determine best beam to use:
        if "BMAJ" in f[0].header.keys():
            if f[0].header["BMAJ"] > bmaj:
                bmaj = f[0].header["BMAJ"]
                bpa = f[0].header["BPA"]
            if f[0].header["BMIN"] > bmin:
                bmin = f[0].header["BMIN"]
        else:
            raise ValueError("No beam information found for {}".format(image))

        # Find frequency information:
        found_freq = False

        if "CRVAL3" in f[0].header.keys() and "CTYPE3" in f[0].header.keys():
            if "FREQ" in f[0].header["CTYPE3"]:
                frequencies.append(f[0].header["CRVAL3"])
                found_freq = True
            

        if not found_freq and "CRVAL4" in f[0].header.keys() and "CTYPE4" in f[0].header.keys():
            if "FREQ" in f[0].header["CTYPE4"]:
                frequencies.append(f[0].header["CRVAL4"])
                found_freq = True


        if not found_freq and "FREQ" in f[0].header.keys():
            frequencies.append(f[0].header["FREQ"])
            found_freq = True

        if not found_freq:
            raise ValueError("No frequency information found for {}".format(image))

        if "BUNIT" not in f[0].header.keys():
            f[0].header["BUNIT"] = "JY/BEAM"
            f[0].header["BTYPE"] = "INTENSITY"

        f.flush()
        f.close()
        
    for image in args.images:

        print("Working on {}".format(image))

        fitsin = "fits in={_in} out={out} op=xyin".format(_in=image,
                                                          out=image.replace(".fits", ".mir"))

        Popen(fitsin, shell=True).wait()

        convol = "convol map={_in} out={out} fwhm={bmaj},{bmin} pa={bpa} " \
                 "options=final".format(_in=image.replace(".fits", ".mir"),
                                        out=image.replace(".fits", ".cvl"),
                                        bmaj=bmaj*3600.,
                                        bmin=bmin*3600.,
                                        bpa=bpa)

        with open("{}_convol.log".format(image.replace(".fits", "")), "w+") as log:
            Popen(convol, shell=True, stdout=log, stderr=log).wait()

        with open("{}_convol.log".format(image.replace(".fits", "")), "r") as log:
            lines = log.readlines()
            for line in lines:
                if "Warning" in line:
                    print("{} did not convolve correctly...".format(image))
                    shutil.rmtree(image.replace(".fits", ".cvl"))
                    break
                elif "Scaling the output by 1.000E+00" in line:
                    print("{} did not convolve correctly...".format(image))
                    shutil.rmtree(image.replace(".fits", ".cvl"))
                    break

        if not os.path.exists(image.replace(".fits", ".cvl")):
            # Image didn't convolve - probably already at the same resolution:
            shutil.copytree(image.replace(".fits", ".mir"), image.replace(".fits", ".cvl"))

        regrid = "regrid in={_in} out={out} axes=1,2 " \
                 "desc={crval1},{crpix1},{cdelt1},{naxis1}," \
                 "{crval2},{crpix2},{cdelt2},{naxis2} " \
                 "project=SIN ".format(_in=image.replace(".fits", ".cvl"),
                                       out=image.replace(".fits", ".rgd"),
                                       crval1=crval1,
                                       crpix1=crpix1,
                                       cdelt1=cdelt1,
                                       naxis1=args.imsize,
                                       crval2=crval2,
                                       crpix2=crpix2,
                                       cdelt2=cdelt2,
                                       naxis2=args.imsize)

        Popen(regrid, shell=True).wait()

        fitsout1 = "fits in={_in} out={out} op=xyout".format(_in=image.replace(".fits", ".rgd"),
                                                             out=image.replace(".fits", "_rgd.fits"))
        fitsout2 = "fits in={_in} out={out} op=xyout".format(_in=image.replace(".fits", ".cvl"),
                                                             out=image.replace(".fits", "_cvl.fits"))

        Popen(fitsout1, shell=True).wait()
        Popen(fitsout2, shell=True).wait()


    for i, image in enumerate(args.images):

        # Ensure we have all the right FITS header items:
        with fits.open(image.replace(".fits", "_rgd.fits"), mode="update") as f:

            f[0].header["BMAJ"] = bmaj
            f[0].header["BMIN"] = bmin
            f[0].header["BPA"] = bpa

            f[0].header["BUNIT"] = "JY/BEAM"
            f[0].header["BTYPE"] = "Intensity"

            f[0].header["OBJECT"] = args.source

            f[0].header["CROTA1"] = 0.
            f[0].header["CROTA2"] = 0.

            f[0].header["FREQ"] = frequencies[i]

            f.flush()


if __name__ == "__main__":
    main()