ctools

#!/usr/bin/env python

import gammalib

import sys

import os

import glob

import json

def print_help():

    print ""

    print "tsmerge.py expects two keyword argumens: use via command line as follows"

    print "tsmerge.py files=<list of files> outfile=<desired outfile name>"

    print "-------------------------------------------------------------------------"

    print "*** \"outfile\" should be a path to a fits file"

    print "*** \"files\" should be either be:"

    print "     - an ascii-file listing the files your want to merge"

    print "  OR - a wildcard string describing the files for merging"

    print "-------------------------------------------------------------------------"

    print ""

    print "Example: python tsmerge,py files=tsmap*.fits outfile=tsmap.fits"

    print ""

    sys.exit(-1)

# Maps class holding the different sky maps 

# created by cttsmap

class maps:

    # Constructor

    def __init__(self,fitsfile):

        self.filename = fitsfile

        if not os.path.isfile(fitsfile):

            self.error("Maps instance cannot be constructed: File \""+self.filename+"\" does not exist")

        # open FITS file

        fits = gammalib.GFits(fitsfile)

        # get TS map

        self.tsmap = gammalib.GSkymap()

        self.tsmap.read(fits[0])

        # Get status map

        if not fits.contains("STATUS MAP"):

            self.error("No extension \"STATUS MAP\" found in \""+self.filename+"\"")

        self.statusmap = gammalib.GSkymap()

        self.statusmap.read(fits["STATUS MAP"])

        # Get other maps 

        self.maps = []

        self.mapnames = []

        for hdu in fits:

            if hdu.extname() != "IMAGE" and hdu.extname() != "STATUS MAP":

                skymap = gammalib.GSkymap()

                skymap.read(hdu)

                self.maps.append(skymap)

                self.mapnames.append(hdu.extname())

    # len operator

    def __len__(self):

        return len(self.maps)

    # operator []

    def __getitem__(self,index):

        return self.maps[index]

    # Check the status map for completeness

    def check_status(self):

        check = True

        for pix in self.statusmap:

            if pix < 0.5:

                check = False

                break

        return check

    # Save maps to outfile

    def save(self,outfile):

        fits = gammalib.GFits()

        self.tsmap.write(fits)

        for i in range(len(self.maps)):

            self.maps[i].write(fits)

        for i in range(len(self.mapnames)):   

            fits[i+1].extname(self.mapnames[i])

        if not self.check_status():

            self.statusmap.write(fits)

            fits[fits.size()-1].extname("STATUS MAP")

        fits.saveto(outfile,True)

    # Add map instances to each other

    # in this way the TS maps get merged  

    def add(self,addmaps):

        if not len(addmaps) == len(self.maps):

            self.error("Maps cannot be added: Maps from \""+self.filename+"\" do not match from \""+addmaps.filename+"\"")

        # Loop over bins    

        for i in range(self.tsmap.npix()):

            if addmaps.statusmap[i] > 0.5:

                self.tsmap[i] = addmaps.tsmap[i]

                self.statusmap[i] = addmaps.statusmap[i]

                for j in range(len(self.maps)):

                    self.maps[j][i] = addmaps[j][i]

    # Unify error messages

    def error(self,message):  

        print "*** ERROR ***: "+message

        sys.exit(-1)

def pars_from_argv(argv):

    if argv is None:

        argv = sys.argv

    args = []

    kwargs = {}

    for s in argv[1:]:

        if s.count('=') == 1:

            key, value = s.split('=', 1)

        else:

            key, value = None, s

        try:

            value = json.loads(value) 

        except ValueError:

            pass

        if key:

            kwargs[key] = value

        else:

            args.append(value)

    return args, kwargs

def get_kwargs(argv):

    kwargs = {}

    for s in argv:

        if s.count('=') == 1:

            key, value = s.split('=', 1) 

            kwargs[key] = value

        else:

            print_help()         

    return kwargs   

if __name__ == "__main__":

    # Initialise list of files

    kwargs = get_kwargs(sys.argv[1:])

    if not kwargs.has_key("outfile") or not kwargs.has_key("files"):

        print_help()

    outfile = kwargs["outfile"]

    files = kwargs["files"]

    allfiles = []

    if os.path.isfile(files):

        allfiles = open(files).readlines()   

    else:

        allfiles = glob.glob(files)

    # Initialise files to be merged

    mergefiles = []

    # Test files for the entry status map

    # and add them to the file list

    for fitsfile in allfiles:    

        fits = gammalib.GFits(fitsfile)

        if fits.contains("STATUS MAP"):

            mergefiles.append(fitsfile)

        fits.close()

    # create a maps instance from the first file

    finalmaps = maps(mergefiles[0])

    # add other files to the maps instance

    for fitsfile in mergefiles[1:]:

        finalmaps.add(maps(fitsfile))

    # save the filled maps to file

    finalmaps.save(outfile=outfile)