GammaLib

#! /Applications/Anaconda/bin/python

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

# This script tests the ON-OFF analysis functionality of the gammalib

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

from ctools import *

from gammalib import *

from math import *

import os

import glob

import sys

import numpy as np

import matplotlib.pyplot as plt

import matplotlib.cm as cm

import matplotlib.patches as ptch

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

# Setup single observation #

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

def single_obs(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \

            emin=0.1, emax=100.0, rad=5.0, \

            irf="South_50h", caldb="prod2", id="000000", instrument="CTA"):

    """

    Returns a single CTA observation.

    Parameters:

     pntdir   - Pointing direction [GSkyDir]

    Keywords:

     tstart     - Start time [seconds] (default: 0.0)

     duration   - Duration of observation [seconds] (default: 1800.0)

     deadc      - Deadtime correction factor (default: 0.95)

     emin       - Minimum event energy [TeV] (default: 0.1)

     emax       - Maximum event energy [TeV] (default: 100.0)

     rad        - ROI radius used for analysis [deg] (default: 5.0)

     irf        - Instrument response function (default: cta_dummy_irf)

     caldb      - Calibration database path (default: "dummy")

     id         - Run identifier (default: "000000")

     instrument - Intrument (default: "CTA")

    """

    # Allocate CTA observation

    obs_cta = GCTAObservation(instrument)

    # Set calibration database

    db = GCaldb()

    if (os.path.isdir(caldb)):

        db.rootdir(caldb)

    else:

        db.open("cta", caldb)

    # Set pointing direction

    pnt = GCTAPointing()

    pnt.dir(pntdir)

    obs_cta.pointing(pnt)

    # Set ROI

    roi     = GCTARoi()

    instdir = GCTAInstDir()

    instdir.dir(pntdir)

    roi.centre(instdir)

    roi.radius(rad)

    # Set GTI

    gti = GGti()

    gti.append(GTime(tstart), GTime(tstart+duration))

    # Set energy boundaries

    ebounds = GEbounds(GEnergy(emin, "TeV"), \

                                GEnergy(emax, "TeV"))

    # Allocate event list

    events = GCTAEventList()

    events.roi(roi)

    events.gti(gti)

    events.ebounds(ebounds)

    obs_cta.events(events)

    # Set instrument response

    obs_cta.response(irf, db)

    # Set ontime, livetime, and deadtime correction factor

    obs_cta.ontime(duration)

    obs_cta.livetime(duration*deadc)

    obs_cta.deadc(deadc)

    obs_cta.id(id)

    # Return CTA observation

    return obs_cta

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

# Add CTA background model to observations #

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

def add_background_model(models,name="Background",instru="CTA",id="",bgd_file="",bgd_sigma=0.0):

    """

    Add CTA background model to a model container.

    Parameters:

     models - a model container

     bgd_file - file function for background spectral dependence

     bgd_sigma - sigma parameter for the gaussian offset angle dependence of the background rate

    """

    # Gaussian radial dependence model

    if bgd_sigma > 0.0:

        # Spatial part

        bgd_radial = GCTAModelRadialGauss(bgd_sigma)

        # Spectral part

        if len(bgd_file) > 0:

            bgd_spectrum = GModelSpectralFunc(bgd_file,1.0)

        else:

            pivot_nrj=GEnergy(1.0e6, "MeV")

            bgd_spectrum = GModelSpectralPlaw(1.0e-6, -2.0, pivot_nrj)

            bgd_spectrum["Prefactor"].value(61.8e-6)

            bgd_spectrum["Prefactor"].scale(1.0e-6)

            bgd_spectrum["PivotEnergy"].value(1.0)

            bgd_spectrum["PivotEnergy"].scale(1.0e6)

            bgd_spectrum["Index"].value(-1.85)

            bgd_spectrum["Index"].scale(1.0)

            # Full model

            bgd_model = GCTAModelRadialAcceptance(bgd_radial, bgd_spectrum)

    # IRF-defined model

    else:

            # Spectral model

            pivot_nrj=GEnergy(1.0e6, "MeV")

            bgd_spectrum = GModelSpectralPlaw(1.0e-6, -2.0, pivot_nrj)

            bgd_spectrum["Prefactor"].value(1.0)

            bgd_spectrum["Prefactor"].scale(1.0)

            bgd_spectrum["PivotEnergy"].value(1.0)

            bgd_spectrum["PivotEnergy"].scale(1.0e6)

            bgd_spectrum["Index"].value(0.0)

            bgd_spectrum["Index"].scale(1.0)

            # Full model

            bgd_model = GCTAModelIrfBackground(bgd_spectrum)

    # Set other background properties

    if (len(name) > 0):

        bgd_model.name(name)

    else:

        bgd_model.name("Background")

    if (len(instru) > 0):

        bgd_model.instruments(instru)

    else:

        bgd_model.instruments("CTA")

    if (len(id) > 0):

        bgd_model.ids(id)

    # By default, set model to be fitted

    if (len(bgd_file) > 0):

        bgd_model['Normalization'].free()

    else:

        bgd_model['Prefactor'].free()

        bgd_model['Index'].free()

    # Add background model to container

    models.append(bgd_model)

    # Return model container

    return models

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

# Set any general type source #

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

def add_source_model(models, name, ra, dec, coord='CEL', \

                     skymap='', specfile='', \

                     flux=5.7e-16, index=-2.5, pivot=3.0e5, \

                     type="point", sigma=1.0, radius=1.0, width=0.1, \

                     nnode=0, emin=1.0, emax=10.0, \

                     pivot_scale=1e6, flux_scale=1.0e-16):

    """

    Adds a single point-source with power-law spectrum to a model container.

    The default is crab-like.

    Parameters:

     models - a model container

     name   - Unique source name

     ra     - RA of source location [deg]

     dec    - Declination of source location [deg]

    Keywords:

     flux   - Source flux density in 1.0e-16 ph/cm2/s/MeV

     index  - Source spectral index

     pivot  - Pivot energy in TeV

     type   - Source spatial model

     sigma/radius/width  - Spatial model parameters

     nnode  - Number of nodes (if > 0, spectral model is a node function)

    """

    # Set source spectral model

    if (nnode > 0) and (emin != emax):

        spectrum = GModelSpectralNodes()

        dloge=log10(emax/emin)/nnode

        logemin=math.log10(emin)

        spectrum.reserve(nnode)

        for i in range(nnode):

                    enode = GEnergy()

                    enode.TeV(10.0**(logemin+(i+0.5)*dloge))

                    inode=flux*(enode.MeV()/pivot)**(index)

                    spectrum.append(enode, inode)

                    spectrum[i*2].scale(pivot_scale)

                    spectrum[i*2].fix()

                    spectrum[i*2+1].scale(flux_scale)

                    spectrum[i*2+1].free()    

    # Set source spectral model

    elif (len(specfile) > 0):

        spectrum = GModelSpectralFunc(specfile)

        spectrum["Normalization"].value(1.0)

        spectrum["Normalization"].free()

    else:

        pivot_nrj=GEnergy()

        pivot_nrj.TeV(pivot*pivot_scale/1e6)

        spectrum = GModelSpectralPlaw(flux, index,pivot_nrj)

        spectrum["Prefactor"].scale(flux_scale)

        spectrum["Prefactor"].value(flux)

        spectrum["Prefactor"].free()

        spectrum["PivotEnergy"].scale(pivot_scale)

        spectrum["PivotEnergy"].value(pivot)

        spectrum["Index"].value(index)

        spectrum["Index"].scale(1.0)

        spectrum["Index"].free()        

    # Set source spatial model

    location = GSkyDir()

    if (coord == 'CEL'):

        location.radec_deg(ra, dec)

    else:

        location.lb_deg(ra, dec)    

    if (len(skymap) > 0):

        spatial = GModelSpatialDiffuseMap(skymap,1.0)

        source = GModelSky(spatial, spectrum)

        spatial[0].free()

    elif type == "point":

        spatial = GModelSpatialPointSource(location)

        spatial[0].free()

        spatial[1].free()

        source  = GModelSky(spatial, spectrum)

    elif type == "gauss":

        radial = GModelSpatialRadialGauss(location, sigma)

        radial[0].free()

        radial[1].free()

        radial[2].free()

        source = GModelSky(radial, spectrum)

    elif type == "disk":

        radial = GModelSpatialRadialDisk(location, radius)

        radial[0].free()

        radial[1].free()

        radial[2].free()

        source = GModelSky(radial, spectrum)

    elif type == "shell":

        radial = GModelSpatialRadialShell(location, radius, width)

        radial[0].free()

        radial[1].free()

        radial[2].free()

        radial[3].free()

        source = GModelSky(radial, spectrum)

    else:

        print "ERROR: Unknown source type '"+type+"'."

        return None

    # Set source name

    source.name(name)

    # Append source model to model container

    models.append(source)

    # Return source

    return models

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

# Set the fit parameters for a given model #

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

def set_fit_param(model, parname=[], parfit=[], parval=[]):

    """

    Set the fit parameters for a given model

    Arguments

    - model: a model container

    - parname: array of parameters names

    - parfit: array of fit/fix flags (True=fit,False=fix)

    - parval: array of initial/fix parameters values

    Notes

    - Arrays should have same numbers of elements

    - A way to not set a value but just the fit/fix flag is to pass/leave parval empty

    - If among several parameters some values must be set and others not, separate these and call the function two times

    """

    # Count/check number of parameters to be set

    if (len(parname) != len(parfit)):

        print "Incorrect input. Fit parameters not set or modified."

        return 0

    else:

        num_toset=len(parname)

    # Loop over parameters

    for i in range(model.size()):

        name=model[i].name()

        if (name in parname):

            # Get index of matching name

            i_par=parname.index(name)

            # Set the parameter to be fix/free

            if (parfit[i_par]):

                model[name].free()

            else:

                model[name].fix()

            # Optionally set the value

            if (len(parval) > 0): 

                model[name].value(parval[i_par])

    # Exit point

    return model

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

# Simulate observations #

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

def sim(obs, log=False, debug=False, chatter=2, edisp=False, seed=0, nbins=0,

        binsz=0.05, npix=200, proj="TAN", coord="GAL", outfile=""):

    """

    Simulate events for all observations in the container.

    Parameters:

     obs   - Observation container

    Keywords:

     log   - Create log file(s)

     debug - Create console dump?

     edisp - Apply energy dispersion?

     seed  - Seed value for simulations (default: 0)

     nbins - Number of energy bins (default: 0=unbinned)

     binsz - Pixel size for binned simulation (deg/pixel)

     npix  - Number of pixels in X and Y for binned simulation

    """

    # Allocate ctobssim application and set parameters

    sim = ctobssim(obs)

    sim["seed"] = seed

    sim["edisp"] = edisp

    sim["outevents"] = outfile

    # Optionally open the log file

    if log:

        sim.logFileOpen()

    # Optionally switch-on debugging model

    if debug:

        sim["debug"] = True

    # Set chatter level

    sim["chatter"] = chatter

    # Run ctobssim application. This will loop over all observations in the

    # container and simulation the events for each observation. Note that

    # events are not added together, they still apply to each observation

    # separately.

    sim.run()

    # Binned option?

    if nbins > 0:

        # Determine common energy boundaries for observations

        emin = None

        emax = None

        for run in sim.obs():

            run_emin = run.events().ebounds().emin().TeV()

            run_emax = run.events().ebounds().emax().TeV()

            if emin == None:

                emin = run_emin

            elif run_emin > emin:

                emin = run_emin

            if emax == None:

                emax = run_emax

            elif run_emax > emax:

                emax = run_emax

        # Allocate ctbin application and set parameters

        bin = ctbin(sim.obs())

        bin["ebinalg"] = "LOG"

        bin["emin"] = emin

        bin["emax"] = emax

        bin["enumbins"] = nbins

        bin["usepnt"] = True # Use pointing for map centre

        bin["nxpix"] = npix

        bin["nypix"] = npix

        bin["binsz"] = binsz

        bin["coordsys"] = coord

        bin["proj"] = proj

        # Optionally open the log file

        if log:

            bin.logFileOpen()

        # Optionally switch-on debugging model

        if debug:

            bin["debug"] = True

        # Set chatter level

        bin["chatter"] = chatter

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

        # the container and bin the events in counts maps

        bin.run()

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

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

        # left)

        obs = bin.obs().copy()

    else:

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

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

        # left)

        obs = sim.obs().copy()

    # Optionally save file

    if (len(outfile) > 0):

        sim.save()

    # Delete the simulation

    del sim

    # Return observation container

    return obs

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

# Make a simple model #

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

def make_simple_model(models,name,ra,dec,flux,idx):

    # Add background model (one for each pointing)

    models = add_background_model(models,bgd_sigma=0.0,instru='CTA',id="0001",name="Background left")

    set_fit_param(models["Background left"],parname=['Sigma','Normalization','Prefactor','Index'],parfit=[False,True,True,False])

    models = add_background_model(models,bgd_sigma=0.0,instru='CTA',id="0002",name="Background right")

    set_fit_param(models["Background right"],parname=['Sigma','Normalization','Prefactor','Index'],parfit=[False,True,True,False])

    # Add source model

    models = add_source_model(models, name, ra, dec, flux=flux, index=idx, pivot=1.0e6, type="point", pivot_scale=1e6, flux_scale=1.0e-16)

    set_fit_param(models["Test source"],parname=['Prefactor','Index','RA','DEC'],parfit=[True,True,False,False])

    # Exit

    return models

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

# Make a simple model #

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

def reset_simple_model(models,ra,dec,flux,idx):

    # Add background model (one for each pointing)

    set_fit_param(models["Background left"],parname=['Sigma','Normalization','Prefactor','Index'],parfit=[False,True,True,False],parval=[0.0,1.0,1.0,0.0])

    set_fit_param(models["Background right"],parname=['Sigma','Normalization','Prefactor','Index'],parfit=[False,True,True,False],parval=[0.0,1.0,1.0,0.0])

    set_fit_param(models["Test source"],parname=['Prefactor','Index','RA','DEC'],parfit=[True,True,False,False],parval=[flux,idx,ra,dec])

    # Exit

    return models

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

# Make a histogram #

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

def make_histogram(arr,n):

    # Make histogram

    dbin=(max(arr)-min(arr))/n

    hist=np.histogram(np.array(arr), bins=n)

    yhist=hist[0]

    xhist=np.linspace(min(arr)+0.5*dbin, max(arr)-0.5*dbin, num=n, endpoint=True)

    # Exit

    return xhist,yhist

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

# Make a plot #

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

def plot_histogram(x,n):

    # Plot

    plt.figure()

    plt.hist(x,n)

    plt.xscale('linear')

    plt.yscale('linear')

    plt.xlabel('Parameter',labelpad=10)

    plt.ylabel('Number',labelpad=10)

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

# Set plotting environment #

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

def prepare_plots():

    # Set

    params = {'legend.fontsize': 14,

              'legend.labelspacing': 0.2,

              'figure.figsize': (8,8),

              'figure.facecolor': 'white',

              'lines.linewidth': 2,

              'axes.titlesize': 'large',

              'axes.labelsize': 'large'}

    # Assign

    plt.rcParams.update(params)

    return    

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

# Make a plot #

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

def plot_results(x,y):

    # Plot

    plt.figure()

    plt.plot(x,y)

    plt.xscale('linear')

    plt.yscale('linear')

    plt.xlabel('Parameter',labelpad=10)

    plt.ylabel('Number',labelpad=10)

    xmin=0.95*min(x)

    xmax=1.05*max(x)

    ymin=0.95*min(y)

    ymax=1.05*max(y)

    plt.xlim([xmin,xmax])

    plt.ylim([ymin,ymax])

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

# Print out fit results #

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

def print_fit_results(obslist):

    # Print

    print "Background left prefactor: %.3e +/- %.3e " % (obslist.models()["Background left"]['Prefactor'].value(),obslist.models()["Background left"]['Prefactor'].error())

    print "Background left index: %.3e +/- %.3e " % (obslist.models()["Background left"]['Index'].value(),obslist.models()["Background left"]['Index'].error())

    print "Background right prefactor: %.3e +/- %.3e " % (obslist.models()["Background right"]['Prefactor'].value(),obslist.models()["Background right"]['Prefactor'].error())

    print "Background right index: %.3e +/- %.3e " % (obslist.models()["Background right"]['Index'].value(),obslist.models()["Background right"]['Index'].error())

    print "Source prefactor: %.3e +/- %.3e " % (obslist.models()["Test source"]['Prefactor'].value(),obslist.models()["Test source"]['Prefactor'].error())

    print "Source index: %.3e +/- %.3e " % (obslist.models()["Test source"]['Index'].value(),obslist.models()["Test source"]['Index'].error())

    print "Fit ended with value %.3e and status %d after %d iterations." % (opt.value(),opt.status(),opt.iter())

    # Exit

    return

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

# Routine entry point #

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

if __name__ == '__main__':

    """

    Aims:

    This script tests the ON-OFF analysis functionality of the gammalib

    Arguments:

    - None

    Notes:

    - Hard-coded source, observations, and regions (more flexible script later...)

    - Approximate method to set OFF regions (works only on equator)

    - Have to set up path to data base (db and irf)

    - Source flux set as flux density at 1TeV

    """

    # Job parameters

    # CTA

    db="prod2"

    irf="South_50h"

    bgd=""

    rsp_sigma=4.0

    bgd_sigma=4.0

    duration=1800.0

    rad=5.0

    # Energy

    emin=0.1

    emax=100.0

    nbins=9

    nnodes=5

    # Source

    name="Test source"

    ra=0.0

    dec=0.0

    flux=1.0e-16

    idx=-2.3

    dflux=1.5  # Offset factor by which true value is rescaled before fit

    didx=0.3   # Offset factor by which true value is rescaled before fit

    # ON-OFF

    onshift=1.0

    onsize=0.3

    noff=3

    # Others

    npull=100

    nhist=10

    chatter=4

    seed=5

    # Create observation container

    obslist = GObservations()

    # Add two offset observations

    leftdir = GSkyDir()

    leftdir.radec_deg(ra+onshift, dec)

    obs=single_obs(leftdir, tstart=0.0, duration=duration, deadc=0.95, \

                   emin=emin, emax=emax, rad=rad, \

                   irf=irf, caldb=db, id="000000", instrument="CTA")

    obs.name("Left")

    obs.id("0001")

    obslist.append(obs)

    rightdir = GSkyDir()

    rightdir.radec_deg(ra-onshift, dec)

    obs=single_obs(rightdir, tstart=0.0, duration=duration, deadc=0.95, \

                   emin=emin, emax=emax, rad=rad, \

                   irf=irf, caldb=db, id="000000", instrument="CTA")

    obs.name("Right")

    obs.id("0002")

    obslist.append(obs)    

    # Create a model container

    models = GModels()   

    # Make a simple model for two observations

    models = make_simple_model(models,name,ra,dec,flux,idx)

    # Add model to the container

    obslist.models(models)

    # Define ON region

    ondir = GSkyDir()

    ondir.radec_deg(0.0,0.0)

    on = GSkyRegions()

    onreg=GSkyRegionCircle(ondir,onsize)

    on.append(onreg)

    # Define OFF regions

    offleft = GSkyRegions()

    offright = GSkyRegions()

    for i in range(noff):

        phi=(i+1)*360./(noff+1.0)

        # Left pointing  

        offdir = GSkyDir(leftdir)       

        offdir.rotate_deg(phi-90., onshift)

        offreg=GSkyRegionCircle(offdir, onsize)

        offleft.append(offreg)

        # Right pointing

        offdir = GSkyDir(rightdir)       

        offdir.rotate_deg(phi+90.0, onshift)       

        offreg=GSkyRegionCircle(offdir, onsize)

        offright.append(offreg)

    # Create ON-OFF analysis object

    onofflist = GCTAOnOffObservations()

    # Define energy binning

    etrue = GEbounds(nbins, GEnergy(emin, "TeV"), GEnergy(emax, "TeV"))

    ereco = GEbounds(nbins, GEnergy(emin, "TeV"), GEnergy(emax, "TeV"))

    # Append ON-OFF observations

    obs=GCTAOnOffObservation(ereco, on, offleft)

    obs.name("Left")

    obs.id("0001")

    onofflist.append(obs)

    obs=GCTAOnOffObservation(ereco, on, offright)

    obs.name("Right")

    obs.id("0002")

    onofflist.append(obs)

    # Create arrays to store pull results

    onoff_flux=[]

    onoff_idx=[]

    unbinned_flux=[]

    unbinned_idx=[]

    # Repeat the run...

    # (make each time new copies of obslist and onofflist otherwise something accumulates and goes wrong)

    for i in range(npull):

        # Print

        print '\nRun ',i

        # Make event list

        theobslist=GObservations(obslist)

        theobslist = sim(theobslist, log=False, debug=False, chatter=chatter, edisp=False, seed=seed+i, nbins=0)

        # Load data and save (fills ON and OFF spectra, ON and OFF observing time, and alpha parameter)

        theonofflist=GCTAOnOffObservations(onofflist)

        theonofflist[0].fill(theobslist[0])

        theonofflist[1].fill(theobslist[1])

            # Compute responses and save

        theonofflist[0].compute_response(theobslist[0],models,etrue)

        theonofflist[1].compute_response(theobslist[1],models,etrue)    

        # Reset model container and append it to observation list

        models=reset_simple_model(models,ra,dec,dflux*flux,idx-didx)

        theonofflist.models(models)

        # Set optimizer and logger

        log=GLog('fit_onoff_%d.log' % (i),True)

        log.chatter(chatter)

        opt=GOptimizerLM(log)

        # Optimize for ON-OFF analysis

        print '\nON-OFF fitting...'

        theonofflist.optimize(opt)

        theonofflist.errors(opt)

        print_fit_results(theonofflist)

        # Append results

        onoff_flux.append(theonofflist.models()["Test source"]['Prefactor'].value())

        onoff_idx.append(theonofflist.models()["Test source"]['Index'].value())

        # Free memory

        del opt

        del log

        del theonofflist

        # Reset model container and append it to observation list

        models=reset_simple_model(models,ra,dec,dflux*flux,idx-didx)

        theobslist.models(models)

        # Set optimizer and logger

        log=GLog('fit_unbinned_%d.log' % (i),True)

        log.chatter(chatter)

        opt=GOptimizerLM(log)

        # Optimize for Fermi-style analysis

        print '\nUnbinned fitting...'

        theobslist.optimize(opt)

        theobslist.errors(opt)

        print_fit_results(theobslist)

        # Append results

        unbinned_flux.append(theobslist.models()["Test source"]['Prefactor'].value())

        unbinned_idx.append(theobslist.models()["Test source"]['Index'].value())

        # Free memory

        del opt

        del log

        del theobslist

    # Print out statistics

    print '\nTrue values: prefactor=%.3e and index=%.3f \n' % (flux,idx)

    print 'ON-OFF analysis'

    print 'Prefactor: mean= %.3e +/- %.3e' % (np.mean(np.array(onoff_flux)),np.std(np.array(onoff_flux)))

    print 'Index= %.3f +/- %.3f \n' % (np.mean(np.array(onoff_idx)),np.std(np.array(onoff_idx)))

    print 'Unbinned analysis'

    print 'Prefactor: mean= %.3e +/- %.3e' % (np.mean(np.array(unbinned_flux)),np.std(np.array(unbinned_flux)))

    print 'Index= %.3f +/- %.3f \n' % (np.mean(np.array(unbinned_idx)),np.std(np.array(unbinned_idx)))

    # Plot results

    prepare_plots()

    plot_histogram(unbinned_flux,nhist)

    plot_histogram(unbinned_idx,nhist)

    plot_histogram(onoff_flux,nhist)

    plot_histogram(onoff_idx,nhist)

    plt.show()