Fri 8 Mar 13:36:58 CST 2019

1 files changed, 180 insertions, 0 deletions

diff --git a/fig2.py b/fig2.py
new file mode 100755
index 0000000..90aca90
--- /dev/null
+++ b/fig2.py
@@ -0,0 +1,180 @@
+#! /usr/bin/env python
+# author : S. Mandalia
+#          s.p.mandalia@qmul.ac.uk
+#
+# date   : February 24, 2019
+
+"""
+HESE BSM Flavour Figure 2
+"""
+
+from __future__ import absolute_import, division
+
+import argparse
+from functools import partial
+
+import numpy as np
+
+from utils import fr as fr_utils
+from utils import misc as misc_utils
+from utils import plot as plot_utils
+from utils.enums import str_enum
+from utils.enums import Likelihood, ParamTag, PriorsCateg
+from utils.param import Param, ParamSet
+
+from matplotlib import pyplot as plt
+
+from pymultinest import Analyzer
+
+
+def define_nuisance():
+    """Define the nuisance parameters."""
+    nuisance = []
+    tag = ParamTag.NUISANCE
+    lg_prior = PriorsCateg.LIMITEDGAUSS
+    nuisance.extend([
+        Param(name='convNorm',        value=1.,  seed=[0.5, 2. ], ranges=[0.1, 10.], std=0.4, prior=lg_prior, tag=tag),
+        Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 20.], std=2.4, prior=lg_prior, tag=tag),
+        Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0. , 10.], std=0.1, tag=tag),
+        Param(name='astroNorm',       value=6.9, seed=[0.,  5. ], ranges=[0. , 20.], std=1.5, tag=tag),
+        Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[-5., 5. ], std=0.1, tag=tag)
+    ])
+    return ParamSet(nuisance)
+
+
+def get_paramsets(args, nuisance_paramset):
+    paramset = []
+    if args.likelihood in [Likelihood.GOLEMFIT, Likelihood.GF_FREQ]:
+        gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)]
+        paramset.extend(gf_nuisance)
+    tag = ParamTag.BESTFIT
+    paramset.extend([
+        Param(name='astroFlavorAngle1', value=0, ranges=[0., 1.], std=0.2, tag=tag),
+        Param(name='astroFlavorAngle2', value=0, ranges=[-1., 1.], std=0.2, tag=tag),
+    ])
+    paramset = ParamSet(paramset)
+    return paramset
+
+
+def process_args(args):
+    """Process the input args."""
+    if args.likelihood is not Likelihood.GOLEMFIT \
+       and args.likelihood is not Likelihood.GF_FREQ:
+        raise AssertionError(
+            'Likelihood method {0} not supported for this '
+            'script!\nChoose either GOLEMFIT or GF_FREQ'.format(
+                str_enum(args.likelihood)
+            )
+        )
+
+def parse_args(args=None):
+    """Parse command line arguments"""
+    parser = argparse.ArgumentParser(
+        description="HESE BSM Flavour Figure 2",
+        formatter_class=misc_utils.SortingHelpFormatter,
+    )
+    parser.add_argument(
+        '--likelihood', default='golemfit',
+        type=partial(misc_utils.enum_parse, c=Likelihood),
+        choices=Likelihood, help='likelihood contour'
+    )
+    parser.add_argument(
+        '--contour-dir', type=str,
+        help='Path to directory containing MultiNest runs'
+    )
+    parser.add_argument(
+        '--outfile', type=str, default='./untitled',
+        help='Output path'
+    )
+    if args is None: return parser.parse_args()
+    else: return parser.parse_args(args.split())
+
+
+def main():
+    args = parse_args()
+    process_args(args)
+    misc_utils.print_args(args)
+
+    paramset = get_paramsets(args, define_nuisance())
+    n_params = len(paramset)
+    print n_params
+
+    # Data
+    data_path = '{0}/{1}/real'.format(
+        args.contour_dir, str_enum(args.likelihood).lower()
+    )
+    prefix = '{0}/_{1}_REAL_mn_'.format(data_path, str_enum(args.likelihood))
+    analyser = Analyzer(
+        outputfiles_basename=prefix, n_params=n_params
+    )
+    print analyser
+
+    pranges = paramset.ranges
+
+    bf = analyser.get_best_fit()['parameters']
+    for i in xrange(len(bf)):
+        bf[i] = (pranges[i][1]-pranges[i][0])*bf[i] + pranges[i][0]
+    print 'bestfit = ', bf
+    print 'bestfit log_likelihood', analyser.get_best_fit()['log_likelihood']
+
+    print
+    print '{0:50} = {1}'.format('global evidence', analyser.get_stats()['global evidence'])
+    print
+
+    chains = analyser.get_data()[:,2:]
+    for x in chains:
+        for i in xrange(len(x)):
+            x[i] = (pranges[i][1]-pranges[i][0])*x[i] + pranges[i][0]
+
+    llh = -0.5 * analyser.get_data()[:,1]
+
+    flavour_angles = chains[:,-2:]
+    flavour_ratios = np.array(
+        map(fr_utils.angles_to_fr, flavour_angles)
+    )
+
+    nbins = 25
+
+    fig = plt.figure(figsize=(8, 8))
+    ax = fig.add_subplot(111)
+    tax = plot_utils.get_tax(ax, scale=nbins)
+
+    plot_utils.flavour_contour(
+        frs = flavour_ratios,
+        ax = ax,
+        nbins = nbins,
+        coverage = 99,
+        linewidth = 2,
+        color = 'green'
+    )
+
+    plot_utils.flavour_contour(
+        frs = flavour_ratios,
+        ax = ax,
+        nbins = nbins,
+        coverage = 90,
+        linewidth = 2,
+        color = 'blue'
+    )
+
+    plot_utils.flavour_contour(
+        frs = flavour_ratios,
+        ax = ax,
+        nbins = nbins,
+        coverage = 68,
+        linewidth = 2,
+        color = 'red'
+    )
+
+    ax.legend()
+
+    fig.savefig('test.png', bbox_inches='tight', dpi=150)
+
+    print "DONE!"
+
+
+main.__doc__ = __doc__
+
+
+if __name__ == '__main__':
+    main()