ctools

import matplotlib.pyplot as plt

import matplotlib.gridspec as gridspec

from gammalib import *

import asciidata

import math

import os

import sys

import pyfits

import numpy as np

def ascii_dict(fname,names=[],**kwargs):

    d = dict()

    fdata = asciidata.open(fname,kwargs)

    if len(names)==0:

        for col in fdata:

            d[col.colname] = col.tonumpy()

    else:       

        for name in names:

            if fdata.find(name) >= 0:

                d[name] = fdata[name].tonumpy()

            else:

                d[name] = np.array([])

                print "WARNING: Column\"",name,"\" not found in",fname

    return d

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,

BLUE="#337EF5"

GREY='grey'

GREEN = '#21851D'

RED = '#BD0B0B'

YELLOW = '#D4CF44'#D1A206'

ORANGE = '#D19534'#EEBA76'

def xfontsize(sbplt,fontsize=18):  

    for tick in sbplt.xaxis.get_major_ticks():

        tick.label1.set_fontsize(fontsize)

def yfontsize(sbplt,fontsize=18):  

    for tick in sbplt.yaxis.get_major_ticks():

        tick.label1.set_fontsize(fontsize)

def xyfontsize(sbplt,fontsize=18):

    xfontsize(sbplt,fontsize)

    yfontsize(sbplt,fontsize)

def xticklines(sbplt,majorsize=7.0,majorwidth=1.5,minorsize=5.0,minorwidth=1.0):

    for line in sbplt.xaxis.get_majorticklines():

        line.set_markersize(majorsize)

        line.set_markeredgewidth(majorwidth)

    for line in sbplt.xaxis.get_minorticklines():

        line.set_markersize(minorsize)

        line.set_markeredgewidth(minorwidth)

def yticklines(sbplt,majorsize=7.0,majorwidth=1.5,minorsize=5.0,minorwidth=1.0):

    for line in sbplt.yaxis.get_majorticklines():

        line.set_markersize(majorsize)

        line.set_markeredgewidth(majorwidth)

    for line in sbplt.yaxis.get_minorticklines():

        line.set_markersize(minorsize)

        line.set_markeredgewidth(minorwidth)

def xyticklines(sbplt,majorsize=7.0,majorwidth=1.5,minorsize=5.0,minorwidth=1.0):

    xticklines(sbplt,majorsize,majorwidth,minorsize,minorwidth)

    yticklines(sbplt,majorsize,majorwidth,minorsize,minorwidth)

class ulimgraph(base):

    def __init__(self,filename,sed=True):

        super(ulimgraph,self).__init__()

        self.scalex = 1.0

        self.scaley = 1.0

        self.fmt='o'

        self.color = BLUE

        self.tslim = 4.0

        self.elinewidth = 2.0

        self.capsize=5.0

        self.markersize=5.0

        self.limlength=0.4

        self.sed = sed  

        self.label ="_nolegend_"

        self.xerrors=True  

        self.nlims = -1

        self._load_dict(filename)  

        self.filename =filename

    def _load_dict(self,filename):

        try:

            self.dict = ascii_dict(filename)

        except IOError:

            self.warning("File "+filename+" not existing - cant draw graph")

            self.dict ={}

            self.dict["flux"] = np.array([])

            self.dict["eu_flux"] = np.array([])

            self.dict["ed_flux"] = np.array([])

            self.dict["ulim_flux"] = np.array([])

            self.dict["ener"] = np.array([])

        if not "ener" in self.dict.keys():

            self.warning("Keyword \"ener\" not found in file "+filename)

            if "energy" in self.dict.keys():

                self.info("Found \"energy\" instead")

                self.dict["ener"] = self.dict["energy"]

                if "ed_energy" in self.dict.keys() and "eu_energy" in self.dict.keys():

                    self.dict["ed_ener"] = self.dict["ed_energy"].copy()

                    self.dict["eu_ener"] = self.dict["eu_energy"].copy()

                else:

                    self.dict["ed_energy"] = np.zeros(len(self.dict["ener"]))

                    self.dict["eu_energy"] = np.zeros(len(self.dict["ener"]))

            elif "emean" in self.dict.keys():

                self.info("Found \"emean\" instead")

                self.dict["ener"] = self.dict["emean"]

                if "ed_emean" in self.dict.keys() and "eu_emean" in self.dict.keys():

                    self.dict["ed_ener"] = self.dict["ed_emean"].copy()

                    self.dict["eu_ener"] = self.dict["eu_emean"].copy()

                else:

                    self.dict["ed_ener"] = np.zeros(len(self.dict["ener"]))

                    self.dict["eu_ener"] = np.zeros(len(self.dict["ener"]))

            else:

                self.error("Required parameter \"ener\" not found in file "+filename)

        if not "flux" in self.dict.keys():

            self.error("Required parameter \"flux\" not found in file "+filename)

        if "e_flux" in self.dict.keys():

            self.warning("only symmetrical errors found")

            self.dict["ed_flux"] = self.dict["e_flux"].copy()

            self.dict["eu_flux"] = self.dict["e_flux"].copy()

        if "ulim" in self.dict.keys():

            self.dict["ulim_flux"] = self.dict["ulim"].copy()

        comment = ""

        try:

            f = asciidata.open(filename)

            comment = f["flux"].colcomment

        except IOError:

            pass

        if self.sed:

            if "/MeV" in comment:

                self.dict["flux"]*=pow(self.dict["ener"],2.0)

                self.dict["eu_flux"]*=pow(self.dict["ener"],2.0)

                self.dict["ed_flux"]*=pow(self.dict["ener"],2.0)

                self.dict["ulim_flux"]*=pow(self.dict["ener"],2.0)

        else:

            if "/MeV" not in comment:

                self.dict["flux"]/=pow(self.dict["ener"],2.0)

                self.dict["eu_flux"]/=pow(self.dict["ener"],2.0)

                self.dict["ed_flux"]/=pow(self.dict["ener"],2.0)       

                self.dict["ulim_flux"]/=pow(self.dict["ener"],2.0)       

    def _plt_points(self):

        ed_ener = []

        eu_ener = []

        ener = []

        ed_flux = []

        eu_flux = []

        flux = []

        for i in range(len(self.dict["flux"])):

            if self.dict["TS"][i] >= self.tslim and self.dict["ed_flux"][i]<self.dict["flux"][i]:

                ed_ener.append(self.dict["ed_ener"][i])

                eu_ener.append(self.dict["eu_ener"][i])

                ener.append(self.dict["ener"][i])

                ed_flux.append(self.dict["ed_flux"][i])

                eu_flux.append(self.dict["eu_flux"][i])

                flux.append(self.dict["flux"][i])

                if self.dict["ed_flux"][i]>self.dict["flux"][i]:

                    self.warning("in "+os.path.basename(self.filename).replace(".sp","")+": ")

                    self.warning("\tflux error: "+str(self.dict["ed_flux"][i])+" is larger than flux value "+str(self.dict["flux"][i]))

                    self.warning("\tTS value is however TS="+str(self.dict["TS"][i])+", weird!")

        if self.xerrors:

            plt.errorbar(ener, flux, xerr=[ed_ener,eu_ener],yerr=[ed_flux,eu_flux],fmt=self.fmt,elinewidth=self.elinewidth,label=self.label,mfc = self.color,ecolor=self.color,mec = self.color,ms =self.markersize)               

        else:

            plt.errorbar(ener, flux,yerr=[ed_flux,eu_flux],fmt=self.fmt,elinewidth=self.elinewidth,label=self.label,mfc = self.color,ecolor=self.color,mec = self.color,ms =self.markersize)               

    def _plt_limits(self):

        ed_ener = []

        eu_ener = []

        ener = []

        ed_flux = []

        eu_flux = []

        flux = []

        drawn_lims = 0

        for i in range(len(self.dict["flux"])):

            if self.dict["TS"][i] < self.tslim or self.dict["ed_flux"][i]>self.dict["flux"][i]:

                if drawn_lims >= self.nlims and self.nlims!=-1:

                    break

                ed_ener.append(self.dict["ed_ener"][i])

                eu_ener.append(self.dict["eu_ener"][i])

                ener.append(self.dict["ener"][i])

                ed_flux.append(self.limlength*self.dict["ulim_flux"][i])

                eu_flux.append(0.0)

                flux.append(self.dict["ulim_flux"][i])

                drawn_lims+=1

        if self.xerrors:

            plt.errorbar(ener, flux, xerr=[ed_ener,eu_ener],yerr=[ed_flux,eu_flux],fmt=self.fmt,markersize=0.0,elinewidth=self.elinewidth,lolims=True,capsize=self.capsize,label='_nolegend_',mfc = self.color,ecolor=self.color,ms =self.markersize)      

        else:

            plt.errorbar(ener, flux, yerr=[ed_flux,eu_flux],fmt=self.fmt,markersize=0.0,elinewidth=self.elinewidth,lolims=True,capsize=self.capsize,label='_nolegend_',mfc = self.color,ecolor=self.color,ms =self.markersize)      

    def draw(self):

        for key in self.dict.keys():

            if "ener" in key:

                self.dict[key]*=self.scalex;

            elif "flux" in key:    

                self.dict[key]*=self.scaley;

        self._plt_points()

        self._plt_limits()

class residuals(base):

    def __init__(self,specfilename,modelfilename):

        super(residuals,self).__init__()     

        self.scalex = 1.0

        self.modelfilename = modelfilename

        self.source = os.path.basename(modelfilename).replace(".optmodel","")

        self.xerrors = True

        self.fmt = 'o'

        self.markercolor = BLUE

        self.linecolor = 'black'

        self.tslim = 4.0

        self.elinewidth=2.0

        self.linewidth=2.0

        self.linestyle = '--'

        self._load_dict(specfilename)

    def _load_dict(self,filename):       

        gr = ulimgraph(filename,False)

        self.dict = gr.dict.copy()

    def _plt_residuals(self):

        ener = []

        ed_ener = []

        eu_ener = []

        flux = []

        ed_flux = []

        eu_flux = []

        time = GTime()

        for i in range(len(self.dict["flux"])):

            if self.dict["TS"][i] >= self.tslim and self.dict["ed_flux"][i]<self.dict["flux"][i]:

                th_val = self.spectral.eval(self.energies[i],time)

                ed_ener.append(self.dict["ed_ener"][i])

                eu_ener.append(self.dict["eu_ener"][i])

                ener.append(self.dict["ener"][i])              

                flux.append((self.dict["flux"][i] - th_val) / th_val)

                ed_flux.append(self.dict["ed_flux"][i]/th_val)

                eu_flux.append(self.dict["ed_flux"][i]/th_val)

        if self.xerrors:

            plt.errorbar(ener, flux, xerr=[ed_ener,eu_ener],yerr=[ed_flux,eu_flux],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 

        else:

            plt.errorbar(ener, flux,yerr=[ed_flux,eu_flux],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 

    def draw(self):

        models = GModels(self.modelfilename)

        try:       

            self.spectral = models[self.source].spectral().clone() 

        except:

            self.error("Source \""+self.source+"\" not found in model container "+self.modelfilename)       

        self.energies = []

        for e in self.dict["ener"]:

            self.energies.append(GEnergy(e,"MeV"))

        for key in self.dict.keys():

            if "ener" in key:

                self.dict[key]*=self.scalex;

        self._plt_residuals()

        plt.axhline(0.0,color=self.linecolor,lw=self.linewidth,ls=self.linestyle)

class plotter(base):

    def __init__(self,roiname):

        super(plotter,self).__init__()

        self.xlim = [0,0]

        self.ylim = [0,0]

        self.ylim_res = [0,0]

        self.resultpath = os.environ["PWD"]

        self.roiname = roiname 

        self.tslim = 4.0

        self.scalex = 1e-6

        self.scaley = 1.60217e-6

    def run(self):

        prefix = self.resultpath+"/"+self.roiname

        spfile = prefix+".sp"

        optmodel = prefix+".xml"

        models = GModels(optmodel)

        residualplot = residuals(spfile,optmodel)

        residualplot.scalex = self.scalex

        residualplot.source=self.roiname

        residualplot.tslim = self.tslim

        gr1 = ulimgraph(spfile)

        gr1.tslim = self.tslim

        gr1.scalex = self.scalex

        gr1.scaley = self.scaley

        plt.figure(num=self.roiname,figsize=(12,7),edgecolor="w")

        gs = gridspec.GridSpec(2, 1,height_ratios=[3,1])

        sb1 = plt.subplot(gs[0])

        plt.loglog()

        gr1.draw()

        plt.ylabel("E$^{2}$ dN/dE [erg cm$^{2}$ s$^{-1}$]",fontsize=20)

        xyfontsize(sb1, fontsize=18)

        xyticklines(sb1,7.0,1.5,5.0,1.0)

        if not self.xlim[0] == 0: 

            plt.xlim(xlim)

        else:

            self.xlim = plt.xlim()

        if not self.ylim[0] == 0:

            plt.ylim(ylim)   

        plt.setp(plt.gca(), 'xticklabels', [])

        plt.tight_layout(pad=0.3,w_pad=0.3, h_pad=0.5) 

        sb2 = plt.subplot(gs[1])    

        plt.xscale('log')

        residualplot.draw()

        plt.xlim(self.xlim) 

        if not self.ylim_res[0] == 0:

            plt.ylim(self.ylim_res)

        plt.locator_params(axis='y', nbins=5)

        plt.xlabel("Energy [TeV]",fontsize=20)

        plt.ylabel("residual",fontsize=20)

        xyfontsize(sb2, fontsize=18)

        xyticklines(sb2,7.0,1.5,5.0,1.0)

        plt.subplots_adjust(left=0.11, right=0.97, top=0.9, bottom=0.11)

        plt.show()

if __name__ == "__main__":

    pl = plotter("CrabNebula")

    pl.run()