GammaLib

#! /usr/bin/env python

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

# Test H.E.S.S. energy dispersion matrix and GCTAEdisp2D class

#

# Copyright (C) 2018 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 math

import gammalib

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

# Test H.E.S.S. energy dispersion #

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

def test_edisp(infile):

    """

    """

    # Open FITS file

    fits = gammalib.GFits(infile)

    # Get energy dispersion extension

    edisp = fits['EDISP']

    # Get true energies

    etrue_lo = edisp['ENERG_LO']

    etrue_hi = edisp['ENERG_HI']

    etrues   = []

    for i in range(etrue_lo.number()):

        etrue = math.sqrt(etrue_lo[0,i] * etrue_hi[0,i])

        etrues.append(etrue)

    # Get migration values

    migra_lo = edisp['MIGRA_LO']

    migra_hi = edisp['MIGRA_HI']

    migras   = []

    for i in range(migra_lo.number()):

        migra = 0.5*(migra_lo[0,i] + migra_hi[0,i])

        migras.append(migra)

    # Get energy dispersion matrix

    matrix = edisp['MATRIX']

    # Get matrix dimensions

    netrue  = matrix.dim()[0]

    nmigra  = matrix.dim()[1]

    noffset = matrix.dim()[2]

    npix    = netrue * nmigra

    # Loop over offset angles

    #for ioffset in range(noffset):

    for ioffset in range(1):

        # Integrate over migra

        sum_migra = []

        for ietrue in range(netrue):

            sum = 0.0

            for imigra in range(nmigra):

                i       = ioffset * npix + imigra * netrue + ietrue

                dmigra  = migra_hi[0,imigra] - migra_lo[0,imigra]

                sum    += matrix[0,i] * dmigra

            sum_migra.append(sum)

        # Integrate over ereco

        #

        # Edisp(Ereco) = Edisp(migra) / Etrue

        #

        sum_ereco = []

        for ietrue in range(netrue):

            sum = 0.0

            for imigra in range(nmigra):

                i          = ioffset * npix + imigra * netrue + ietrue

                ereco_min  = migra_lo[0,imigra] * etrues[ietrue]

                ereco_max  = migra_hi[0,imigra] * etrues[ietrue]

                dereco     = ereco_max - ereco_min

                sum       += matrix[0,i] / etrues[ietrue] * dereco

            sum_ereco.append(sum)

        # Integrate over log(ereco)

        #

        # Edisp(Ereco) = Edisp(migra) * (Ereco / Etrue)

        #

        sum_logereco = []

        for ietrue in range(netrue):

            sum = 0.0

            for imigra in range(nmigra):

                i             = ioffset * npix + imigra * netrue + ietrue

                ereco_min     = migra_lo[0,imigra] * etrues[ietrue]

                ereco_max     = migra_hi[0,imigra] * etrues[ietrue]

                ereco         = math.sqrt(ereco_min * ereco_max)

                logereco_min  = math.log(ereco_min)

                logereco_max  = math.log(ereco_max)

                dlog10ereco   = logereco_max - logereco_min

                sum          += matrix[0,i] * ereco / etrues[ietrue] * dlog10ereco

            sum_logereco.append(sum)

        # Integrate over log10(ereco)

        #

        # Edisp(Ereco) = Edisp(migra) * (ln(10) * Ereco) / Etrue)

        #

        sum_log10ereco = []

        for ietrue in range(netrue):

            sum = 0.0

            for imigra in range(nmigra):

                i               = ioffset * npix + imigra * netrue + ietrue

                ereco_min       = migra_lo[0,imigra] * etrues[ietrue]

                ereco_max       = migra_hi[0,imigra] * etrues[ietrue]

                ereco           = math.sqrt(ereco_min * ereco_max)

                log10ereco_min  = math.log10(ereco_min)

                log10ereco_max  = math.log10(ereco_max)

                dlog10ereco     = log10ereco_max - log10ereco_min

                sum            += matrix[0,i] * math.log(10.0) * ereco / etrues[ietrue] * dlog10ereco

            sum_log10ereco.append(sum)

        # Dump results for all energy bins

        for ietrue in range(netrue):

            result = '%2d %3d %8.5f %8.5f %8.5f %8.5f' % \

                     (ioffset, ietrue, sum_migra[ietrue], sum_ereco[ietrue],

                      sum_logereco[ietrue], sum_log10ereco[ietrue])

            print(result)

    # Return

    return

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

# Test GCTAEdisp2D class #

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

def test_GCTAEdisp2D(infile):

    """

    """

    # Load energy dispersion

    edisp2D = gammalib.GCTAEdisp2D(infile)

    # Fetch energy dispersion data

    edisp2D.fetch()

    # Get matrix dimensions

    netrue  = edisp2D.table().axis_bins(0)

    nmigra  = edisp2D.table().axis_bins(1)

    noffset = edisp2D.table().axis_bins(2)

    npix    = netrue * nmigra

    # Loop over offset angles

    for ioffset in range(noffset):

        # Get theta

        theta = edisp2D.table().axis_nodes(2)[ioffset]

        # Loop over true energies

        for ietrue in range(netrue):

            # Get log10 Etrue

            logEsrc = edisp2D.table().axis_nodes(0)[ietrue]

            etrue   = gammalib.GEnergy()

            etrue.log10TeV(logEsrc)

            # Get energy boundaries

            ereco_bounds = edisp2D.ebounds_obs(logEsrc, theta)

            #

            ereco_min = ereco_bounds.emin()

            ereco_max = ereco_bounds.emax()

            #

            prob = edisp2D.prob_erecobin(ereco_min, ereco_max, etrue, theta)

            result = 'Prob: %2d %3d %8.5f' % (ioffset, ietrue, prob)

            print(result)

    # Return

    return

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

# Main routine entry point #

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

if __name__ == '__main__':

    # Set filename

    filename = '/project-data/hess/hess_dl3_dr1/data/hess_dl3_dr1_obs_id_020326.fits'

    # Test file

    test_edisp(filename)

    # Test GCTAEdisp2D

    test_GCTAEdisp2D(filename)