1 files changed, 163 insertions, 0 deletions
diff --git a/austin.py b/austin.py
new file mode 100755
index 0000000..cddbed1
--- /dev/null
+++ b/austin.py
@@ -0,0 +1,163 @@
+#! /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
+import json
+
+
+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 = '/home/aschneider/programs/GOLEMSPACE/sources/GolemFit/scripts/diffuse/mcmcs/results/dpl_numu_prior_flavor_20190302-162221-a747f528-8aa6-4488-8c80-059572c099fe.json'
+    data_path = '/home/aschneider/programs/GOLEMSPACE/sources/GolemFit/scripts/diffuse/mcmcs/results/spl_flavor_20190311-170924-5297d736-3c6e-447f-8de7-4a0653a51bb6.json'
+    with open(data_path) as f:
+        d_json = json.load(f)
+
+    names = d_json['func_args']
+    chains = np.array(d_json['chain'])
+    print 'names', names
+    print 'chains.shape', chains.shape
+
+    flavour_angles = chains[:,4:6]
+    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_austin.png', bbox_inches='tight', dpi=150)
+
+    print "DONE!"
+
+
+main.__doc__ = __doc__
+
+
+if __name__ == '__main__':
+    main()