add spectral index implementation

4 files changed, 87 insertions, 39 deletions

diff --git a/utils/enums.py b/utils/enums.py
index a5875a8..4637eeb 100644
--- a/utils/enums.py
+++ b/utils/enums.py
@@ -16,6 +16,11 @@ class DataType(Enum):
     ASMIMOV = 3
 
 
+class EnergyDependance(Enum):
+    MONO     = 1
+    SPECTRAL = 2
+
+
 class FitPriorsCateg(Enum):
     DEFAULT = 1
     XS      = 2
@@ -42,6 +47,5 @@ class MCMCSeedType(Enum):
 
 
 class SteeringCateg(Enum):
-    P2_0                  = 1
-    P2_1                  = 2
-
+    P2_0 = 1
+    P2_1 = 2
diff --git a/utils/fr.py b/utils/fr.py
index ddcb5d2..4242e97 100644
--- a/utils/fr.py
+++ b/utils/fr.py
@@ -344,7 +344,6 @@ def u_to_fr(source_fr, matrix):
         array([ 0.33740075,  0.33176584,  0.33083341])
 
     """
-    # TODO(shivesh): energy dependence
     composition = np.einsum(
         'ai, bi, a -> b', abs(matrix)**2, abs(matrix)**2, source_fr
     )
diff --git a/utils/likelihood.py b/utils/likelihood.py
index 5c9af51..9bba79a 100644
--- a/utils/likelihood.py
+++ b/utils/likelihood.py
@@ -59,26 +59,64 @@ def triangle_llh(theta, args, asimov_paramset, mcmc_paramset, fitter):
     for param in mcmc_paramset.from_tag(ParamTag.NUISANCE):
         hypo_paramset[param.name].value = param.value
 
+    if args.energy_dependance is EnergyDependance.SPECTRAL:
+        bin_centers = np.sqrt(args.binning[:-1]*args.binning[1:])
+        bin_width = np.abs(np.diff(args.binning))
+
     if args.fix_source_ratio:
-        fr1, fr2, fr3 = args.source_ratio
+        if args.energy_dependance is EnergyDependance.MONO:
+            source_flux = args.source_ratio
+        elif args.energy_dependance is EnergyDependance.SPECTRAL:
+	    source_flux = np.array(
+		[fr * np.power(bin_centers, SPECTRAL_INDEX)
+		 for fr in args.source_ratio]
+	    ).T
     else:
-        fr1, fr2, fr3 = fr_utils.angles_to_fr(
-            mcmc_paramset.from_tag(ParamTag.SRCANGLES, values=True)
-        )
+        if args.energy_dependance is EnergyDependance.MONO:
+            source_flux = fr_utils.angles_to_fr(
+                mcmc_paramset.from_tag(ParamTag.SRCANGLES, values=True)
+            )
+        elif args.energy_dependance is EnergyDependance.SPECTRAL:
+            source_flux = np.array(
+                [fr * np.power(bin_centers, args.spectral_index)
+                 for fr in angles_to_fr(theta[-2:])]
+            ).T
+
     bsm_angles = mcmc_paramset.from_tag(
         [ParamTag.SCALE, ParamTag.MMANGLES], values=True
     )
 
-    u = fr_utils.params_to_BSMu(
-        theta      = bsm_angles,
-        dim        = args.dimension,
-        energy     = args.energy,
-        no_bsm     = args.no_bsm,
-        fix_mixing = args.fix_mixing,
-        fix_scale  = args.fix_scale,
-        scale      = args.scale
-    )
-    fr = fr_utils.u_to_fr((fr1, fr2, fr3), u)
+    if args.energy_dependance is EnergyDependance.MONO:
+        u = fr_utils.params_to_BSMu(
+            theta      = bsm_angles,
+            dim        = args.dimension,
+            energy     = args.energy,
+            no_bsm     = args.no_bsm,
+            fix_mixing = args.fix_mixing,
+            fix_scale  = args.fix_scale,
+            scale      = args.scale
+        )
+        fr = fr_utils.u_to_fr(source_flux, u)
+    elif args.energy_dependance is EnergyDependance.SPECTRAL:
+        mf_perbin = []
+        for i_sf, sf_perbin in enumerate(source_flux):
+            u = fr_utils.params_to_BSMu(
+                theta      = bsm_angles,
+                dim        = args.dimension,
+                energy     = args.energy,
+                no_bsm     = args.no_bsm,
+                fix_mixing = args.fix_mixing,
+                fix_scale  = args.fix_scale,
+                scale      = args.scale
+            )
+            fr = fr_utils.u_to_fr(sf_perbin, u)
+            mf_perbin.append(fr)
+        measured_flux = np.array(mf_perbin).T
+        intergrated_measured_flux = np.sum(measured_flux * bin_width, axis=1)
+        averaged_measured_flux = (1./(args.binning[-1] - args.binning[0])) * \
+            intergrated_measured_flux
+        fr = averaged_measured_flux / np.sum(averaged_measured_flux)
+
     for idx, param in enumerate(hypo_paramset.from_tag(ParamTag.BESTFIT)):
         param.value = fr[idx]
 
diff --git a/utils/plot.py b/utils/plot.py
index a24c69c..a659b78 100644
--- a/utils/plot.py
+++ b/utils/plot.py
@@ -21,7 +21,7 @@ from getdist import plots
 from getdist import mcsamples
 
 from utils import misc as misc_utils
-from utils.enums import ParamTag
+from utils.enums import Likelihood, ParamTag
 from utils.fr import angles_to_u, angles_to_fr
 
 rc('text', usetex=False)
@@ -68,40 +68,47 @@ def plot_Tchain(Tchain, axes_labels, ranges):
 
 def gen_figtext(args):
     """Generate the figure text."""
+    t = ''
     mr1, mr2, mr3 = args.measured_ratio
     if args.fix_source_ratio:
         sr1, sr2, sr3 = args.source_ratio
         if args.fix_scale:
-            return 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \
+            t += 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \
                     'observed flavour ratio = [{3:.2f}, {4:.2f}, ' \
-                    '{5:.2f}]\nSigma = {6:.3f}\nDimension = {7}\nEnergy = ' \
-                    '{8} GeV\nScale = {9}'.format(
-                        sr1, sr2, sr3, mr1, mr2, mr3, args.sigma_ratio,
-                        args.dimension, int(args.energy), args.scale
+                    '{5:.2f}]\nDimension = {7}\nScale = {9}'.format(
+                        sr1, sr2, sr3, mr1, mr2, mr3, args.dimension,
+                        int(args.energy), args.scale
                     )
         else:
-            return 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \
+            t += 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \
                     'observed flavour ratio = [{3:.2f}, {4:.2f}, ' \
-                    '{5:.2f}]\nSigma = {6:.3f}\nDimension = {7}\nEnergy = {8} ' \
-                    'GeV'.format(
-                        sr1, sr2, sr3, mr1, mr2, mr3, args.sigma_ratio,
-                        args.dimension, int(args.energy)
+                    '{5:.2f}]\nDimension = {7}'.format(
+                        sr1, sr2, sr3, mr1, mr2, mr3, args.dimension,
+                        int(args.energy)
                     )
     else:
         if args.fix_scale:
-            return 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \
-                    '{2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} ' \
-                    'GeV\nScale = {6}'.format(
-                        mr1, mr2, mr3, args.sigma_ratio, args.dimension,
-                        int(args.energy), args.scale
+            t += 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \
+                    '{2:.2f}]\nDimension = {4}\nScale = {6}'.format(
+                        mr1, mr2, mr3, args.dimension, int(args.energy),
+                        args.scale
                     )
 	else:
-            return 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \
-                    '{2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} ' \
-                    'GeV'.format(
-                        mr1, mr2, mr3, args.sigma_ratio, args.dimension,
-                        int(args.energy)
+            t += 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \
+                    '{2:.2f}]\nDimension = {4}'.format(
+                        mr1, mr2, mr3, args.dimension, int(args.energy)
                     )
+    if args.likelihood is Likelihood.GAUSSIAN:
+        t += '\nSigma = {0:.3f}'.format(args.sigma_ratio)
+    if args.energy_dependance is EnergyDependance.SPECTRAL:
+        t += '\nSpectral Index = {0}\nBinning = [{7}, {8}] GeV - {9} bins'.format(
+            int(args.spectral_index), int(10**args.binning[0]),
+            int(10**args.binning[1]), int(args.binning[2])
+        )
+    elif args.energy_dependance is EnergyDependance.MONO:
+        t += '\nEnergy = {0} GeV'.format(args.energy)
+    return t
+
 
 def chainer_plot(infile, outfile, outformat, args, mcmc_paramset):
     """Make the triangle plot."""