GammaLib

#! /usr/bin/env python

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

# Fit event rates

# --------------------------------------------------------------------------

# Copyright (C) 2023 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 gammalib

import glob

import math

import matplotlib.pyplot as plt

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

# Rate model class #

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

class rate_model(gammalib.GPythonOptimizerFunction):

    # Constructor

    def __init__(self, arrays, ieng, norm):

        # Call base class constructor

        gammalib.GPythonOptimizerFunction.__init__(self)

        # Set eval method

        self._set_eval(self.eval)

        # Set data

        self._counts = arrays['counts']

        self._ieng   = ieng

        self._norm   = norm

        # Extract dimensions

        self._noads   = self._counts.shape()[0]

        self._nphibar = self._counts.shape()[1]

        # Multiply veto and constant rate by EHA cut correction. Rates

        # are zero in case that the vetorate is zero.

        self._vetorate  = gammalib.GNdarray(self._noads, self._nphibar)

        self._constrate = gammalib.GNdarray(self._noads, self._nphibar)

        self._diffrate  = gammalib.GNdarray(self._noads, self._nphibar)

        for ioad in range(self._noads):

            vetorate = arrays['vetorate'][ioad]

            if vetorate > 0.0:

                constrate = arrays['constrate'][ioad]

                for iphibar in range(self._nphibar):

                    self._vetorate[ioad, iphibar]  = vetorate  * arrays['ehacutcorr'][ioad,iphibar]

                    self._constrate[ioad, iphibar] = constrate * arrays['ehacutcorr'][ioad,iphibar]

                    self._diffrate[ioad, iphibar]  = self._vetorate[ioad, iphibar] - self._constrate[ioad, iphibar]

        # Compute rate sums

        self._vetorate_sum  = gammalib.GNdarray(self._nphibar)

        self._constrate_sum = gammalib.GNdarray(self._nphibar)

        self._diffrate_sum  = gammalib.GNdarray(self._nphibar)

        for iphibar in range(self._nphibar):

            for ioad in range(self._noads):

                self._vetorate_sum[iphibar]  += self._vetorate[ioad, iphibar]

                self._constrate_sum[iphibar] += self._constrate[ioad, iphibar]

            self._diffrate_sum[iphibar] = self._vetorate_sum[iphibar] - self._constrate_sum[iphibar]

        # Return

        return

    # Methods

    def eval(self):

        """

        Evaluate function

        """

        # Recover parameters

        fprompt = self._pars()[0].value()

        fconst  = 1.0 - fprompt

        # Allocate phibar normalistion arrays

        nevents = gammalib.GNdarray(self._nphibar)

        nmodel  = gammalib.GNdarray(self._nphibar)

        norm    = gammalib.GNdarray(self._nphibar)

        # Compute model

        model = fprompt * self._vetorate + fconst * self._constrate

        # Phibar normalise model

        model_norm = model.copy()

        for iphibar in range(self._nphibar):

            for ioad in range(self._noads):

                if model_norm[ioad, iphibar] > 0.0:

                    nevents[iphibar] += self._counts[ioad, iphibar, self._ieng]

                    nmodel[iphibar]  += model_norm[ioad, iphibar]

            if nmodel[iphibar] > 0.0:

                norm[iphibar] = self._norm * nevents[iphibar] / nmodel[iphibar]

                for ioad in range(self._noads):

                    model_norm[ioad, iphibar] *= norm[iphibar]

        # Compute LogL

        logL = 0.0

        for ioad in range(self._noads):

            for iphibar in range(self._nphibar):

                if model_norm[ioad, iphibar] > 0.0:

                    logL -= self._counts[ioad, iphibar, self._ieng] * \

                            math.log(model_norm[ioad, iphibar]) - \

                            model_norm[ioad, iphibar]

        # Evaluate gradient and curvature

        grad = 0.0

        for iphibar in range(self._nphibar):

            # Precompute normalisation gradient

            nmodel2 = nmodel[iphibar] * nmodel[iphibar]

            g_norm  = -self._norm * nevents[iphibar] / nmodel2 * self._diffrate_sum[iphibar]

            # Loop over superpackets

            for ioad in range(self._noads):

                # Continue only if model is positive

                if model_norm[ioad, iphibar] > 0.0:

                    # Pre computation

                    fb = self._counts[ioad, iphibar, self._ieng] / model_norm[ioad, iphibar]

                    fc = (1.0 - fb)

                    fa = fb / model_norm[ioad, iphibar]

                    # Compute parameter gradient

                    g = norm[iphibar] * self._diffrate[ioad, iphibar] +  g_norm * model[ioad, iphibar]

                    # Add to factor gradient

                    grad += g

                    # Update gradient

                    self.gradient()[0] += fc * g

                    # Update curvature

                    self.curvature()[0,0] += fa * g * g

        # Set value

        self._set_value(logL)

        # Return

        return

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

# Load working arrays #

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

def load_working_arrays(filename):

    """

    Load working arrays

    Parameters

    ----------

    filename : str

        Working arrays file name

    Returns

    -------

    arrays : dict

        Working arrays

    """

    # Open FITS file

    fits = gammalib.GFits(filename)

    # Get images

    image_counts     = fits.image('COUNTS')

    image_ehacutcorr = fits.image('EHACUTCORR')

    image_vetorate   = fits.image('VETORATE')

    # Extract dimensions

    noads   = image_counts.naxes(0)

    nphibar = image_counts.naxes(1)

    neng    = image_counts.naxes(2)

    # Allocate Ndarrays

    counts     = gammalib.GNdarray(noads, nphibar, neng)

    ehacutcorr = gammalib.GNdarray(noads, nphibar)

    vetorate   = gammalib.GNdarray(noads)

    constrate  = gammalib.GNdarray(noads)

    # Extract data from images

    for ioad in range(noads):

        for iphibar in range(nphibar):

            for ieng in range(neng):

                counts[ioad,iphibar,ieng] = image_counts[ioad,iphibar,ieng]

            ehacutcorr[ioad,iphibar] = image_ehacutcorr[ioad,iphibar]

        vetorate[ioad]  = image_vetorate[ioad]

        constrate[ioad] = 1.0

    # Normalise rates to unity (excluding bins where vetorate is zero)

    nvetorate  = 0.0

    nconstrate = 0.0

    for ioad in range(noads):

        if vetorate[ioad] > 0.0:

            nvetorate  += vetorate[ioad]

            nconstrate += constrate[ioad]

        else:

            constrate[ioad] = 0.0

    if nvetorate > 0.0:

        for ioad in range(noads):

            vetorate[ioad] /= nvetorate

    if nconstrate > 0.0:

        for ioad in range(noads):

            constrate[ioad] /= nconstrate

    # Put arrays in dictionary

    arrays = {'counts'    : counts,

              'ehacutcorr': ehacutcorr,

              'vetorate'  : vetorate,

              'constrate' : constrate}

    # Return arrays

    return arrays

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

# Compute logL for energy bin #

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

def compute_logL(arrays, ieng, norm, fprompt):

    """

    Compute logL for energy bin

    Parameters

    ----------

    arrays : dict

        Working arrays

    ieng : int

        Energy bin

    norm : float

        Normalisation

    fprompt : float

        Prompt fraction [0-1]

    Returns

    -------

    logL : float

        Log-likelihood value

    """

    # Extract dimensions

    noads   = arrays['counts'].shape()[0]

    nphibar = arrays['counts'].shape()[1]

    # Compute model normalised to number of events for each Phibar

    model = gammalib.GNdarray(noads, nphibar)

    for ioad in range(noads):

        rate = fprompt       * arrays['vetorate'][ioad] + \

               (1.0-fprompt) * arrays['constrate'][ioad]

        for iphibar in range(nphibar):

            model[ioad,iphibar] = rate * arrays['ehacutcorr'][ioad,iphibar]

    for iphibar in range(nphibar):

        nevents = 0.0

        nmodel  = 0.0

        for ioad in range(noads):

            if model[ioad, iphibar] > 0.0:

                nevents += arrays['counts'][ioad, iphibar, ieng]

                nmodel  += model[ioad, iphibar]

        if nmodel > 0.0:

            for ioad in range(noads):

                model[ioad, iphibar] *= (norm * (nevents / nmodel))

    # Compute LogL

    logL = 0.0

    for ioad in range(noads):

        for iphibar in range(nphibar):

            if model[ioad, iphibar] > 0.0:

                logL += arrays['counts'][ioad, iphibar, ieng] * math.log(model[ioad, iphibar]) - \

                        model[ioad, iphibar]

    print(ieng, norm, fprompt, logL)

    # Return logL

    return logL

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

# Show results #

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

def show_results(results):

    """

    Show results

    """

    # Create figure

    fig = plt.figure(figsize=(6,4))

    fig.subplots_adjust(left=0.15, right=0.99, top=0.93, bottom=0.12)

    plt.title('log-likelihood')

    # Determine number of norms

    nnorms = len(results['norms'])

    # Loop over energies

    for index, ieng in enumerate(results['iengs']):

        # Set label and color

        if ieng == 0:

            label = '0.75-1 MeV'

            color = '#1f77b4' # blue

        elif ieng == 1:

            label = '1-3 MeV'

            color = '#ff7f0e' # orange

        elif ieng == 2:

            label = '3-10 MeV'

            color = '#279e68' # green

        elif ieng == 3:

            label = '10-30 MeV'

            color = '#d62728' # red

        else:

            label = ''

            color = '#b5bd61' # ocker

        # Loop over norms

        xs     = []

        ys     = []

        ymaxs  = []

        styles = []

        labels = []

        for inorm, norm in enumerate(results['norms']):

            # Setup vectors

            x = [f    for f    in results['fprompts']]

            y = [logL for logL in results['logLs'][index*nnorms+inorm]]

            # Append vectors

            xs.append(x)

            ys.append(y)

            ymaxs.append(max(y))

            # Append labels and styles

            if norm == 1.0:

                styles.append('o-')

                labels.append(label)

            elif norm < 1.0:

                styles.append('o:')

                labels.append(None)

            else:

                styles.append('o--')

                labels.append(None)

        # Nomalise y vectors

        ymax = max(ymaxs)

        for k, _ in enumerate(ys):

            for i, _ in enumerate(ys[k]):

                ys[k][i] -= ymax

        # Plot vectors

        for i, _ in enumerate(xs):

            plt.plot(xs[i], ys[i], styles[i], label=labels[i], color=color)

    # Set attributes

    plt.xlabel(r'f$_{\rm prompt}$')

    plt.ylabel(r'$\Delta$logL')

    plt.legend(loc='lower left')

    # Show plot

    plt.show()

    # Return

    return

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

# Create log-likelihood plot #

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

def create_logL_plot(arrays):

    """

    Create log-likelihood plot

    """

    # Setup arrays

    iengs    = [0,1,2,3]

    norms    = [0.99,1.0,1.01]

    fprompts = [0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0]

    # Compute logL

    results = {'iengs'   : iengs,

               'norms'   : norms,

               'fprompts': fprompts,

               'logLs'   : []}

    for ieng in iengs:

        for norm in norms:

            logLs = []

            for fprompt in fprompts:

                logL = compute_logL(arrays, ieng, norm, fprompt)

                logLs.append(logL)

            results['logLs'].append(logLs)

    # Show results

    show_results(results)

    # Return

    return

# ======== #

# Fit logL #

# ======== #

def fit_logL(arrays, ieng, norm, fprompt):

    """

    Fit logL

    Parameters

    ----------

    arrays : dict

        Working arrays

    ieng : int

        Energy bin

    norm : float

        Normalisation

    fprompt : float

        Initial prompt fraction

    Returns

    -------

    logL : float

        Log-likelihood value

    """

    # Set fprompt parameter

    pars = gammalib.GOptimizerPars()

    par1 = gammalib.GOptimizerPar('fprompt', fprompt)

    par1.range(0.0,1.0)

    par1.factor_gradient(1.0) # To avoid zero gradient log message

    pars.append(par1)

    # Set fit function

    fct = rate_model(arrays, ieng, norm)

    # Optimize function and compute errors

    log = gammalib.GLog()

    log.cout(True)

    opt = gammalib.GOptimizerLM(log)

    opt.optimize(fct, pars)

    opt.errors(fct, pars)

    # Retrieve logL

    logL = opt.value()

    # Recover fprompt parameter and errors

    fprompt   = pars[0].value()

    e_fprompt = pars[0].error()

    # Print fit results

    print('%.3f (%7.5f +/- %7.5f)' % (logL, fprompt, e_fprompt))

    # Return logL

    return logL

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

# Test log-likelihood fitting #

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

def test_logL_fit(arrays):

    """

    Test log-likelihood fitting

    """

    # Setup arrays

    iengs    = [0,1,2,3]

    norms    = [1.0]

    fprompts = [0.5]

    # Compute logL

    results = {'iengs'   : iengs,

               'norms'   : norms,

               'fprompts': fprompts,

               'logLs'   : []}

    for ieng in iengs:

        for norm in norms:

            logLs = []

            for fprompt in fprompts:

                logL = fit_logL(arrays, ieng, norm, fprompt)

                logLs.append(logL)

            results['logLs'].append(logLs)

    # Show results

    show_results(results)

    # Return

    return

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

# Main routine entry point #

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

if __name__ == '__main__':

    # Dump header

    print('')

    print('*******************')

    print('* Fit event rates *')

    print('*******************')

    # Load working arrays

    arrays = load_working_arrays('debug_gcomdris_compute_drws_vetorate.fits')

    # Create logL plot

    #create_logL_plot(arrays)

    # Fit fprompt parameter

    #fit_logL(arrays, 0, 1.0)

    test_logL_fit(arrays)