1 files changed, 175 insertions, 0 deletions
diff --git a/sens_bayes.py b/sens_bayes.py
new file mode 100755
index 0000000..b0030d4
--- /dev/null
+++ b/sens_bayes.py
@@ -0,0 +1,175 @@
+#! /usr/bin/env python
+# author : S. Mandalia
+#          s.p.mandalia@qmul.ac.uk
+#
+# date   : April 10, 2018
+
+"""
+HESE BSM flavour ratio analysis script
+"""
+
+from __future__ import absolute_import, division
+
+import numpy as np
+import matplotlib as mpl
+mpl.use('Agg')
+from matplotlib import pyplot as plt
+from matplotlib import rc
+
+import GolemFitPy as gf
+import pymultinest
+from pymultinest.solve import solve
+from pymultinest.watch import ProgressPrinter
+
+import fr
+from utils import gf as gf_utils
+from utils import likelihood as llh_utils
+from utils import misc as misc_utils
+from utils.enums import Likelihood, ParamTag
+from utils.plot import plot_BSM_angles_limit
+
+rc('text', usetex=False)
+rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18})
+
+
+RUN = False
+
+
+z = 0.
+scenarios = [
+    [np.sin(np.pi/2.)**2, z, z, z],
+    [z, np.cos(np.pi/2.)**4, z, z],
+    [z, z, np.sin(np.pi/2.)**2, z],
+]
+xticks = [r'$\mathcal{O}_{12}$', r'$\mathcal{O}_{13}$', r'$\mathcal{O}_{23}$']
+
+def fit_flags(flag_dict):
+    flags = gf.FitParametersFlag()
+    for key in flag_dict.iterkeys():
+        flags.__setattr__(key, flag_dict[key])
+    return flags
+
+default_flags = {
+    # False means it's not fixed in minimization
+    'astroFlavorAngle1'         : True,
+    'astroFlavorAngle2'         : True,
+    # 'astroENorm'                : True,
+    # 'astroMuNorm'               : True,
+    # 'astroTauNorm'              : True,
+    'convNorm'                  : False,
+    'promptNorm'                : False,
+    'muonNorm'                  : False,
+    'astroNorm'                 : False,
+    'astroParticleBalance'      : True,
+    'astroDeltaGamma'           : False,
+    'cutoffEnergy'              : True,
+    'CRDeltaGamma'              : False,
+    'piKRatio'                  : False,
+    'NeutrinoAntineutrinoRatio' : True,
+    'darkNorm'                  : True,
+    'domEfficiency'             : True,
+    'holeiceForward'            : True,
+    'anisotropyScale'           : True,
+    'astroNormSec'              : True,
+    'astroDeltaGammaSec'        : True
+}
+
+
+def main():
+    args = fr.parse_args()
+    fr.process_args(args)
+    misc_utils.print_args(args)
+
+    bins = 10
+
+    asimov_paramset, mcmc_paramset = fr.get_paramsets(args)
+
+    sc_range = mcmc_paramset.from_tag(ParamTag.SCALE)[0].ranges
+    scan_scales = np.linspace(sc_range[0], sc_range[1], bins)
+    print 'scan_scales', scan_scales
+
+    p = mcmc_paramset.from_tag([ParamTag.SCALE, ParamTag.MMANGLES], invert=True)
+    n_params = len(p)
+    prior_ranges = p.seeds
+
+    outfile = './sens'
+    if RUN:
+        if args.likelihood is Likelihood.GOLEMFIT:
+            fitter = gf_utils.setup_fitter(args, asimov_paramset)
+            fitter.SetFitParametersFlag(fit_flags(default_flags))
+        elif args.likelihood is Likelihood.GAUSSIAN:
+            fitter = None
+
+        def CubePrior(cube, ndim, nparams):
+            # default are uniform priors
+            return ;
+
+        data = np.zeros((len(scenarios), bins, 2))
+        mm_angles = mcmc_paramset.from_tag(ParamTag.MMANGLES)
+        sc_angles = mcmc_paramset.from_tag(ParamTag.SCALE)[0]
+        for idx, scen in enumerate(scenarios):
+            scales = []
+            evidences = []
+            for yidx, an in enumerate(mm_angles):
+                an.value = scen[yidx]
+            for sc in scan_scales:
+                sc_angles.value = sc
+                def lnProb(cube, ndim, nparams):
+                    for i in range(ndim):
+                        prange = prior_ranges[i][1] - prior_ranges[i][0]
+                        p[i].value = prange*cube[i] + prior_ranges[i][0]
+                    for name in p.names:
+                        mcmc_paramset[name].value = p[name].value
+                    theta = mcmc_paramset.values
+                    # print 'theta', theta
+                    # print 'mcmc_paramset', mcmc_paramset
+                    return llh_utils.triangle_llh(
+                        theta=theta,
+                        args=args,
+                        asimov_paramset=asimov_paramset,
+                        mcmc_paramset=mcmc_paramset,
+                        fitter=fitter
+                    )
+                # TODO(shivesh)
+                prefix = 'mnrun_{0:.0E}'.format(np.power(10, sc)) + '_' + misc_utils.gen_outfile_name(args)[2:]
+                print 'begin running evidence calculation for {0}'.format(prefix)
+                result = pymultinest.run(
+                    LogLikelihood=lnProb,
+                    Prior=CubePrior,
+                    n_dims=n_params,
+                    importance_nested_sampling=True,
+                    n_live_points=args.bayes_live_points,
+                    evidence_tolerance=args.bayes_tolerance,
+                    outputfiles_basename=prefix,
+                    resume=False,
+                    verbose=True
+                )
+
+                analyzer = pymultinest.Analyzer(outputfiles_basename=prefix, n_params=n_params)
+                a_lnZ = analyzer.get_stats()['global evidence']
+                print 'Evidence = {0}'.format(a_lnZ)
+                scales.append(sc)
+                evidences.append(a_lnZ)
+
+            for i, d in enumerate(evidences):
+                data[idx][i][0] = scales[i]
+                data[idx][i][1] = d
+
+        np.save(outfile + '.npy', data)
+
+    plot_BSM_angles_limit(
+        infile=outfile+'.npy',
+        outfile=outfile,
+        xticks=xticks,
+        outformat=['png'],
+        args=args,
+        bayesian=True
+    )
+
+
+main.__doc__ = __doc__
+
+
+if __name__ == '__main__':
+    main()
+