Nathan Kelley-Hoskins / 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,pow,sqrt
import os

try:
        import matplotlib.pyplot as plt
        import matplotlib.gridspec as gridspec
        has_matplotlib = True
except:
        has_matplotlib = False

try:
        import aplpy
        has_aplpy = True
except:
        has_aplpy = False

class options(object):
    def __init__(self):
        self.fmt         = 'o'
        self.color       = BLUE
        self.linecolor   = 'black'
        self.markercolor = self.color
        self.tslim       = 4.0
        self.elinewidth  = 2.0
        self.linewidth   = 2.0
        self.linestyle   = '--'
        self.capsize     = 5.0
        self.markersize  = 5.0
        self.limlength   = 0.4 
        self.label       = "_nolegend_"

import analysisutils 

class SpectralBase(object):
    def __init__(self,sed=True,sed_factor=1.):
        self.m_sed_factor = sed_factor
        self.m_sed        = sed
        
    def _convertSED(self,tab,energy):
        """convert the a table storing dN/dE in sed with right units"""

        if len(tab) != len(energy):
            self.error("input table and energy table does not match")
            return
        if self.m_sed:
            for i in xrange(len(tab)):
                tab[i] *= energy[i]**2*self.m_sed_factor
        return tab

BLUE="#337EF5"
GREY='grey'
GREEN = '#21851D'
RED = '#BD0B0B'
YELLOW = '#D4CF44'
ORANGE = '#D19534'


class residuals(analysisutils.base,options,SpectralBase):
    def __init__(self,spectrum,datapoint,sed,sed_factor,eunit="TeV"):
        super(residuals,self).__init__()
        options.__init__(self)
        SpectralBase.__init__(self,sed,sed_factor)

        # load the gammalib spectrum object
        self.spectrum  = spectrum
        self.SetDataPoint(datapoint)
        
        self.xerrors = True
        self.refUnit = eunit     # reference unit for conversion
        
    def SetDataPoint(self,datapoint):
        self.m_point_info = datapoint
        self.Npt = len(self.m_point_info["dnde"]["value"])

    def _plt_residuals(self):
        th_val = []
        index  = []
        ed_y   = []
        eu_y   = []
        y      = []
        
        Energy_list = {"MeV":1,"GeV":1e-3,"TeV":1e-6} #TODO
        for i in range(self.Npt):
            if self.m_point_info["TS"][i] >= self.tslim:
                th_val.append(self.spectrum.eval(gl.GEnergy(self.m_point_info["ener"]["value"][i],self.m_point_info["ener"]["unit"]),gl.GTime(0)))
                index.append(i)
            else:
                th_val.append(0)
        if self.m_sed:
            self._convertSED(th_val,self.m_point_info["ener"])
        for i in index:
                y.append((self.m_point_info["dnde"]["value"][i] - th_val[i]) / th_val[i])
                ed_y.append(self.m_point_info["dnde"]["ed_value"][i]/th_val[i])
                eu_y.append(self.m_point_info["dnde"]["eu_value"][i]/th_val[i])
                
        if self.xerrors:
            plt.errorbar(self.m_point_info["ener"]["value"], y, xerr=[self.m_point_info["ener"]["ed_value"],self.m_point_info["ener"]["ed_value"]],yerr=[ed_y,eu_y],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 
        else:
            plt.errorbar(self.m_point_info["ener"]["value"], y, yerr=[ed_y,eu_y],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 
        
        plt.axhline(0.0,color=self.linecolor,lw=self.linewidth,ls=self.linestyle)


class ulimgraph(analysisutils.base,options,SpectralBase):
    def __init__(self,datapoint,sed,sed_factor,eunit='TeV'):
        super(ulimgraph,self).__init__()
        #~ super(ulimgraph,self).__init__()
        options.__init__(self)
        SpectralBase.__init__(self,sed,sed_factor)
        
        self.xerrors=True  
        self.nlims = -1
        self.refUnit = eunit
        self.SetDataPoint(datapoint)
    
    def SetDataPoint(self,datapoint):
        self.m_point_info = datapoint
        self.Npt = len(self.m_point_info["dnde"]["value"])
       
    def _plt_points(self):
        results = analysisutils.ResultsSorage()
        
        Energy_list = {"MeV":1,"GeV":1e-3,"TeV":1e-6} #TODO
        for i in range(self.Npt):

            self.m_point_info["ener"]["value"][i]*=Energy_list[self.refUnit]
            self.m_point_info["ener"]["ed_value"][i]*=Energy_list[self.refUnit]
            self.m_point_info["ener"]["eu_value"][i]*=Energy_list[self.refUnit]
            self.m_point_info["ener"]["unit"] = self.refUnit

            if self.m_point_info["TS"][i] >= self.tslim:

                if self.m_point_info["dnde"]["ed_value"][i]>self.m_point_info["dnde"]["value"][i]:
                    self.warning("\tflux error: "+str(self.m_point_info["dnde"]["ed_value"][i])+" is larger than flux value "+str(self.m_point_info["dnde"]["ed_value"][i]))
                    self.warning("\tTS value is however TS="+str(self.m_point_info["TS"][i])+", weird!")
            else :
                self.m_point_info["dnde"]["ed_value"][i] = 0
                self.m_point_info["dnde"]["eu_value"][i] = 0
            
        if self.xerrors:
            plt.errorbar(self.m_point_info["ener"]["value"], self.m_point_info["dnde"]["value"], xerr=[self.m_point_info["ener"]["ed_value"], self.m_point_info["ener"]["eu_value"]],yerr=[self.m_point_info["dnde"]["ed_value"], self.m_point_info["dnde"]["eu_value"]],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 
        else:
            plt.errorbar(self.m_point_info["ener"]["value"], self.m_point_info["dnde"]["value"],yerr=[results["dnde"]["ed_value"],results["dnde"]["eu_value"]],fmt=self.fmt,color=self.markercolor,elinewidth=self.elinewidth) 
        
    def _plt_limits(self):
        ed_flux = []
        eu_flux = []
        flux    = []
        ener    = []
        
        drawn_lims = 0
        Energy_list = {"MeV":1,"GeV":1e-3,"TeV":1e-6} #TODO
        for i in range(self.Npt):
            if self.m_point_info["TS"][i] < self.tslim:
                self.info("Plot an upper limit for E = "+self.m_point_info["ener"]["value"][i]+" with a TS = "+self.m_point_info["TS"][i])
                if drawn_lims >= self.nlims and self.nlims!=-1:
                    break

                ed_flux.append(self.limlength*self.m_point_info["ulim_dnde"][i])
                eu_flux.append(0.0)
                flux.append(self.m_point_info["ulim_dnde"][i])
                drawn_lims+=1
                
        # if self.xerrors:
            # plt.errorbar(self.m_point_info["ener"]["value"], self.m_point_info["dnde"]["value"], xerr=[self.m_point_info["ener"]["ed_value"], self.m_point_info["ener"]["eu_value"]],yerr=[self.m_point_info["dnde"]["ed_value"],self.m_point_info["dnde"]["eu_value"]],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)      
            

class SpectrumPlotter(analysisutils.base,SpectralBase):
    """ Class to compute and plot the spectra (i.e best fit model and butterfly).
        if data points are provided, only change applied  is energy unit convertion to match the butterfly one
        (i.e. no dN/dE to SED conversion).
        It is currently assumed that the unit is MeV for the energy of the data point. this will be change 
        once the I/O container format will be defined."""
    def __init__(self,srcname,like,datapoint,emin=0.1,emax=100,energy = "TeV",SED = True ,npt = 50):
        super(SpectrumPlotter,self).__init__()

        self.m_ebound   = [emin,emax]
        self.m_name     = srcname
        self.m_like     = like
        self.m_spectral = like.obs().models()[srcname].spectral()
        self.m_energy   = []
        self.m_flux     = []
        self.m_error    = []
        self.m_but      = []
        self.m_enebut   = []
        self.m_npt      = npt
        self.m_covar    = gl.GMatrix(1,1)
        self.covar()
        #~ this part take care of the style : SED or differential flux
        #~ in the case of SED, the y-unit is erg/cm2/s
        self.m_sed = SED
        self.eunit = energy
        self.m_sed_factor = 1.
        self._validateUnits()
        SpectralBase.__init__(self,self.m_sed,self.m_sed_factor)

        #Assume MeV for the energy #TODO
        self.points = ulimgraph(datapoint,self.m_sed,gl.MeV2erg,energy)
        self.residuals = residuals(self.m_spectral,datapoint,self.m_sed,gl.MeV2erg,energy)


    def _validateUnits(self):
        Energy_list = {"MeV":1,"GeV":1e3,"TeV":1e6}
        if not(Energy_list.has_key(self.eunit)):
            self.warning("Change energy unit to TeV")
            self.eunit = "TeV"
        
        if self.m_sed:
            self.m_sed_factor = gl.MeV2erg*Energy_list[self.eunit]**2

    def covar(self):
        # Get covariance matrix if fit has converged
        try:
            fullcovar = self.m_like.obs().function().curvature().invert()
        except:
            self.warning("Covariance matrix not determined successfully")
            return
        #~ get the matrix element of the source of interest
        #~ 1) get the right element indices
        idx = 0
        par_index_map = []
        for m in self.m_like.obs().models() :
            if m.name()!= self.m_name:
                idx += m.size()
                continue
            idx += m.spatial().size()
            for par in m.spectral():
                if par.is_free():
                    par_index_map.append(idx)
                    idx += 1

        #~ 2) get the elemets and store them in a matrix
        self.m_covar = gl.GMatrix(len(par_index_map),len(par_index_map))
        i = 0
        for xpar in par_index_map:
            j = 0
            for ypar in par_index_map:
                self.m_covar[i,j] = fullcovar[xpar,ypar]
                j += 1
            i += 1
        self.success("Covariance Matrix succesfuly computed")

    def _makespectrum(self):
        for i in xrange(self.m_npt):
            #Compute the energy
            lene = log10(self.m_ebound[0])+log10(self.m_ebound[1])*i/(self.m_npt-1)
            self.m_energy.append(pow(10,lene))
            #Compute the flux
            self.m_flux.append(self.m_spectral.eval(gl.GEnergy(pow(10,lene),self.eunit),gl.GTime(0)))

            #Compute the gradients
            if self.m_covar.size() != 1:
                self.m_spectral.eval_gradients(gl.GEnergy(pow(10,lene), self.eunit),gl.GTime(0))
                #store them into a GVector
                Derivative = gl.GVector(self.m_covar.size()/2)
                j = 0
                for par in self.m_spectral:
                    if par.is_free():
                        Derivative[j] = par.factor_gradient()
                    j += 1
                #~ computed the error
                self.m_error.append(sqrt(Derivative * (self.m_covar*Derivative)))

        #~ convert the flux in sed if asked
        self.m_flux = self._convertSED(self.m_flux,self.m_energy)
        if self.m_covar.size() != 1:
            self.m_error = self._convertSED(self.m_error,self.m_energy)
            self._makebutterfly() # make the butterfly
        self.success("Spectrum computed")

        
    def _makebutterfly(self):
        """make the butterfly by appending element in a table"""
        for i in xrange(self.m_npt):
            self.m_but.append(self.m_flux[i]+self.m_error[i])
            self.m_enebut.append(self.m_energy[i])
        for i in xrange(self.m_npt):
            idx = self.m_npt-i-1
            self.m_but.append(self.m_flux[idx]-self.m_error[idx])
            self.m_enebut.append(self.m_energy[idx])
        self.m_but.append(self.m_but[0])
        self.m_enebut.append(self.m_enebut[0])

    def write(self):
        self.warning("write function not yet implemented")

    def draw(self):
        """plot the results, provide a file named ptfile for the flux points"""
        gs = gridspec.GridSpec(2, 1,height_ratios=[3,1])
        self._makespectrum()
        self.write()

        if not(has_matplotlib):
            self.warning("matplotlib module not found, can draw")
            return

        plt.figure("results",figsize=(12,7),edgecolor="w")

        sb1 = plt.subplot(gs[0])
        plt.loglog()
        plt.xlabel("Energy ["+ self.eunit+"]",fontsize=15)
        
        if self.m_sed:
            plt.ylabel("E$^{2}$ dN/dE [erg cm$^{2}$ s$^{-1}$]",fontsize=15)
        else:
            plt.ylabel("dN/dE [cm$^{-2}$ s$^{-1}$ "+ self.eunit+"$^{-1}$]",fontsize=15)

        plt.plot(self.m_energy,self.m_flux)
        plt.plot(self.m_enebut,self.m_but)
        
        #~ plot the data points on top of the buterfly
        self.points._plt_points()
        self.points._plt_limits()

        #~ plot the residual points
        sb2 = plt.subplot(gs[1])     
        plt.xscale('log')     
        self.residuals._plt_residuals()

        plt.show()


# ============= #
# Plot skymaps  #
# ============= #
class MapPlotter(analysisutils.base):
    def __init__(self,modelmap, countmap):
        super(MapPlotter,self).__init__()
        self.modelmap = modelmap
        self.countmap = countmap
        
        if not(has_aplpy) :
            self.error("Module aplypy needed")
        if not(has_matplotlib) :
            self.error("Module matplotlib needed")
            
    def draw(self):
        """ Plot the model map.    """
        # Load model map
        fig = plt.figure()
        f1 = aplpy.FITSFigure(self.countmap, figure=fig, subplot=[0.1,0.1,0.35,0.8])
        f1.set_tick_labels_font(size='small')
        f1.set_axis_labels_font(size='small')
        f1.show_colorscale()

        f1.tick_labels.set_yformat('dd:mm:ss')
        f1.tick_labels.set_xformat('hh:mm:ss')
        f1.axis_labels.set_xtext('Right Ascension (J2000)')
        f1.axis_labels.set_ytext('Declination (J2000)')
        f1.ticks.set_length(10.5, minor_factor=0.5)

        f2 = aplpy.FITSFigure(self.modelmap, figure=fig, subplot=[0.5,0.1,0.35,0.8])
        f2.set_tick_labels_font(size='small')
        f2.set_axis_labels_font(size='small')
        f2.show_colorscale()

        f2.hide_yaxis_label()
        f2.hide_ytick_labels()

        fig.canvas.draw()
        fig.show()