ctools

# ==========================================================================

# This script provides a number of functions that are useful for handling

# CTA observations.

#

# Copyright (C) 2011-2014 David Sanchez Michael Mayer Rolf Buehler Juergen Knoedlseder

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program.  If not, see <http://www.gnu.org/licenses/>.

#

# ==========================================================================

import ctools as ct

import gammalib as gl

from math import log10,floor

import sys

class base(object):

    def __init__(self):

        self.classname = self.__class__.__name__

        self.errorcolor = "\033[31m"#red

        self.infocolor = "\033[34m"#blue

        self.warningcolor = "\033[33m"#yellow

        self.successcolor = "\033[32m"#green

        self.endcolor = "\033[0m"#reset

        self.prependfunction = True

        self.basemembers = ["classname","errorcolor","infocolor","warningcolor","successcolor","endcolor","prependfunction"]

    def error(self,message, functionname = ""):

        printstring = ""

        if functionname == "":

            printstring = "\n"+self.errorcolor+"*** Error ["+self.classname+"]: "+message+" ***\n"+self.endcolor

        else:

            printstring = "\n"+self.errorcolor+"*** Error ["+self.classname+"::"+functionname+"]: "+message+" ***\n"+self.endcolor

        sys.exit(printstring)

    def info(self,message,newline=True):

        printstring = self.infocolor+"["+self.classname+"]: "+message+self.endcolor     

        if newline:

            print printstring

        else:

            print self.infocolor+message+self.endcolor,

            sys.stdout.flush()  

    def warning(self,message,functionname = ""):

        printstring = ""

        if functionname == "":

            printstring = self.warningcolor+"["+self.classname+"] Warning: "+message+self.endcolor

        else:

            printstring = self.warningcolor+"["+self.classname+"::"+functionname+"] Warning: "+message+self.endcolor

        print printstring

    def success(self,message):

        printstring = self.successcolor+"["+self.classname+"]: "+message+self.endcolor

        print printstring

    def progress(self,message="."):

        string = self.infocolor+message+self.endcolor

        print string

class ResultsSorage(dict,base):

    """class to store the results in the good format"""

    def __init__(self,*arg,**kw):

        super(ResultsSorage, self).__init__(*arg, **kw)

        base.__init__(self)

        self["dnde"]      = []

        self["eu_dnde"]   = []

        self["ed_dnde"]   = []

        self["ulim_dnde"] = []

        self["ener"]      = []

        self["eu_ener"]   = []

        self["ed_ener"]   = []

        self["TS"]        = []

        self["Iflux"]     = []

        self["e_Iflux"]   = []

    def Print(self):

        self.info("Results are :")

        for k in self.iterkeys():

            print "\t",k,"\t",self[k]

class Analyser(base):

    def __init__(self):

        super(Analyser,self).__init__()

        self.info("Creating "+self.classname+" object")

        self.m_obs = None

        self.like  = None

        self.m_evtfile  = "events_selected.fits"        

        self.m_raw_evt  = "events.fits"

        self.m_cntfile  = "cntmap.fits"

        self.m_ebinalg  = "LOG"

        self.m_eunit    = "TeV"

        self.m_ra       = 83.63

        self.m_dec      = 22.01

        self.m_roi      = 2.

        self.m_tmin     = 0.

        self.m_tmax     = 0.

        self.m_emin     = 0.1#emin

        self.m_emax     = 10.#emax

        self.m_enumbins = 10#nbins

        self.m_usepnt   = True # Use pointing for map centre

        self.m_nxpix    = 200#npix

        self.m_nypix    = 200#npix

        self.m_binsz    = 0.05#binsz

        self.m_coordsys = "GAL"#coord

        self.m_proj     = "TAN"#proj

        self.m_binned   = False

        self.m_stat     = "POISSON"#stat

        self.m_caldb    = "dummy"

        self.m_irf      = "cta_dummy_irf"

        self.m_edisp    = False

        self.m_xml      = "$PWD/crab.xml"

        self.m_xml_out  = "$PWD/crab_results.xml"

    def set_obs(self,obs):

        self.m_obs = obs #set the GObservation container

    def set_energy_boundary(self,emin,emax):

        self.m_emin = emin

        self.m_emax = emax

        ebounds = gl.GEbounds(gl.GEnergy(emin, self.m_eunit), \

                                gl.GEnergy(emax, self.m_eunit))

        if self.m_obs:

            for obs in self.m_obs:

                obs.events().ebounds(ebounds)

    def ctselect(self):

        # ctselect application and set parameters

        self.info("Running ctselect to cut on events")

        if self.m_obs:

            filter = ct.ctselect(self.m_obs)

        else:

            filter = ct.ctselect()

        filter["infile"] = self.m_raw_evt               

        filter["outfile"] = self.m_evtfile  

        filter["usepnt"].boolean(self.m_usepnt) # Use pointing for map centre

        filter["ra"]  = self.m_ra

        filter["dec"]  = self.m_dec

        filter["rad"]  = self.m_roi

        filter["tmin"] = self.m_tmin

        filter["tmax"] = self.m_tmax

        filter["emin"].real(self.m_emin)

        filter["emax"].real(self.m_emax)

        filter["expr"] = ""

        filter.logFileOpen()

        filter.run()

        if not(self.m_obs):

            filter.save()

        if self.m_obs:

            # Make a deep copy of the observation that will be returned

            # (the ctbin object will go out of scope one the function is

            # left)

            self.m_obs = filter.obs().copy()

    def ctbin(self,log=False,debug=False):

        # ctbin application and set parameters

        self.info("Running ctbin to create count map")

        if self.m_obs:

            bin = ct.ctbin(self.m_obs)

        else:

            bin = ct.ctbin()

            bin["evfile"] = self.m_evtfile

        bin["outfile"] = self.m_cntfile

        bin["ebinalg"].string(self.m_ebinalg)

        bin["emin"].real(self.m_emin)

        bin["emax"].real(self.m_emax)

        Nbdecade = log10(self.m_emax)-log10(self.m_emin)#Compute the number of decade

        bin["enumbins"].integer(int(self.m_enumbins*Nbdecade))

        bin["usepnt"].boolean(self.m_usepnt) # Use pointing for map centre

        bin["nxpix"].integer(self.m_nxpix)

        bin["nypix"].integer(self.m_nypix)

        bin["binsz"].real(self.m_binsz)

        bin["coordsys"].string(self.m_coordsys)

        bin["proj"].string(self.m_proj)

        # Optionally open the log file

        if log:

            bin.logFileOpen()

        # Optionally switch-on debugging model

        if debug:

            bin["debug"].boolean(True)

        # Run ctbin application. This will loop over all observations in

        # the container and bin the events in counts maps

        bin.excute()

        if self.m_obs:

            # Make a deep copy of the observation that will be returned

            # (the ctbin object will go out of scope one the function is

            # left)

            self.m_obs = bin.obs().copy()

    def create_fit(self,log=False,debug=False):

        # create ctlike instance with given parameters

        self.info("Fitting Data using ctlike")

        if self.m_obs:

            self.like = ct.ctlike(self.m_obs)

        else:

            self.like = ct.ctlike()

            if self.m_binned:

                self.like["infile"] = self.m_cntfile

            else:

                self.like["infile"] = self.m_evtfile

            self.like["stat"].string(self.m_stat)

            self.like["caldb"].string(self.m_caldb)

            self.like["irf"].string(self.m_irf)

            self.like["srcmdl"] = self.m_xml 

        self.like["edisp"].boolean(self.m_edisp)

        self.like["outmdl"] = self.m_xml_out

        # Optionally open the log file

        if log:

            self.like.logFileOpen()

        # Optionally switch-on debugging model

        if debug:

            self.like["debug"].boolean(True)

    def fit(self,log=False,debug=False):

        if not(self.like):

            self.warning("ctlike object not created, creating now")

            self.create_fit(log,debug)

        # Run ctlike application.

        self.like.run()

        # Save the results in XML

        self.like.save()

        self.success("Fit performed")

    def PrintResults(self,srcname = ""):

        self.info("Results of the Fit")

        for m in self.like.obs().models():

            if srcname == m.name() or srcname=="":

                print "Model : "+m.name()

                print m

    def GetSrcResuls(self,Res,srcname,E0=1,factor = 1e6):

        results = self.GetSrcParams(srcname)

        Res["ener"].append(E0*factor)

        Res["ed_ener"].append((E0-self.m_emin)*factor)

        Res["eu_ener"].append((self.m_emax-E0)*factor)

        Res["Iflux"].append(results["flux"])

        Res["e_Iflux"].append(results["eflux"])

        m = self.GetModel(srcname)

        Res["dnde"].append(m.spectral().eval(gl.GEnergy(E0,self.m_eunit),gl.GTime(0)))

        m.spectral().eval_gradients(gl.GEnergy(E0,self.m_eunit),gl.GTime(0))

        # Assume that par 0 is the norm

        Res["eu_dnde"].append(m.spectral()[0].factor_gradient())

        Res["ed_dnde"].append(m.spectral()[0].factor_gradient())

        Res["TS"].append(100)

        # Res["TS"].append(self.like.obs().models()[srcname].ts())

        Res["ulim_dnde"].append(Res["dnde"][-1])

        return Res

    def GetSrcParams(self,srcname):

        results = {}

        found = False

        m = self.GetModel(srcname)

        results["flux"] = m.spectral().flux(gl.GEnergy(self.m_emin,self.m_eunit),gl.GEnergy(self.m_emax,self.m_eunit))

        results["eflux"] = m.spectral().eflux(gl.GEnergy(self.m_emin,self.m_eunit),gl.GEnergy(self.m_emax,self.m_eunit))

        for par in m.spectral():

            if par.is_free():

                results[par.name()] = par.value()

                results["e"+par.name()] = par.error()

        return results

    def GetModel(self,srcname):

        found = False

        for m in self.like.obs().models():

            if m.name() == srcname:

                return m

        self.error("source "+srcname+" not in the source list")

def MakeEbin(analyse,nbins,srcname,binned = False):

    #Store the results

    Res = ResultsSorage()

    #check the source name, exit is the srcname is false

    analyse.GetModel(srcname)

    #compute the conversion factor to have energy in MeV

    Energy_list = {"MeV":1,"GeV":1e3,"TeV":1e6} #TODO

    factor = Energy_list[analyse.m_eunit]

    for i in xrange(nbins):

        obs_copy = analyse.like.obs().copy()

        emin = pow(10,log10(analyse.m_emin) + (log10(analyse.m_emax)-log10(analyse.m_emin))/(nbins)*i)

        emax = pow(10,log10(analyse.m_emin) + (log10(analyse.m_emax)-log10(analyse.m_emin))/(nbins)*(i+1))

        E0 = pow(10, (log10(emin) + log10(emax)) / 2)

        print "Making energy bin between %2.1e"%(emin)+analyse.m_eunit+" and %2.1e"%(emax)+analyse.m_eunit

        # change the PivotEnergy to E0, froze all spectral parameters but the prefactor

        # m = obs_copy.models()[srcname]

        # for par in m.spectral():

            # if par.name() == "PivotEnergy":

                # par.value(E0*factor)

            # if not(par.name() == "Prefactor") and par.is_free(): # assume that all model have a prefactor

                # par.fix()

            # if par.name() == "Prefactor":

                # Set the norm to a better guess

                # new_val = m.spectral().eval(gl.GEnergy(E0,analyse.m_eunit),gl.GTime(0))

                # m.spectral()["Prefactor"].value(new_val)

        ana_bin = Analyser()

        #copy the obs object

        ana_bin.set_obs(obs_copy)

        # Set energy boundaries

        ana_bin.set_energy_boundary(emin,emax)

        ana_bin.ctselect()

        # if binned:

            # ana_bin.ctbin()

        # Create the fit object

        ana_bin.create_fit(log=False,debug=False)

        # Run ctlike application.

        ana_bin.like.run()

        #Retrive the results in the good format

        Res = ana_bin.GetSrcResuls(Res,srcname,E0,factor)

        del ana_bin

    return Res