create script which produces the flavour ratio ternary contour

2 files changed, 155 insertions, 81 deletions

diff --git a/contour.py b/contour.py
index db3e21a..5615fe2 100755
--- a/contour.py
+++ b/contour.py
@@ -10,20 +10,23 @@ HESE flavour ratio contour
 
 from __future__ import absolute_import, division
 
-import os
 import argparse
 from functools import partial
 
 import numpy as np
 
+from utils import fr as fr_utils
 from utils import gf as gf_utils
 from utils import llh as llh_utils
 from utils import misc as misc_utils
+from utils import mn as mn_utils
 from utils import plot as plot_utils
 from utils.enums import str_enum
 from utils.enums import DataType, Likelihood, ParamTag
 from utils.param import Param, ParamSet, get_paramsets
 
+from pymultinest import Analyzer, run
+
 
 def define_nuisance():
     """Define the nuisance parameters."""
@@ -34,7 +37,7 @@ def define_nuisance():
         Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 50.], std=0.05, tag=tag),
         Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0. , 50.], std=0.1,  tag=tag),
         Param(name='astroNorm',       value=6.9, seed=[0.1, 10.], ranges=[0. , 50.], std=0.1,  tag=tag),
-        Param(name='astroDeltaGamma', value=2.5, seed=[1. , 3. ], ranges=[-5., 5. ], std=0.1,  tag=tag)
+        Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[-5., 5. ], std=0.1,  tag=tag)
     ])
     return ParamSet(nuisance)
 
@@ -49,6 +52,7 @@ def nuisance_argparse(parser):
 
 def process_args(args):
     """Process the input args."""
+    args.plot_angles = args.plot_chains
     if args.likelihood is not Likelihood.GOLEMFIT \
        and args.likelihood is not Likelihood.GF_FREQ:
         raise AssertionError(
@@ -74,6 +78,14 @@ def parse_args(args=None):
         help='Do the scan from scratch'
     )
     parser.add_argument(
+        '--plot-chains', type=misc_utils.parse_bool, default='False',
+        help='Plot the (joint) posteriors'
+    )
+    parser.add_argument(
+        '--plot-triangle', type=misc_utils.parse_bool, default='False',
+        help='Project the posterior contour on the flavour triangle'
+    )
+    parser.add_argument(
         '--seed', type=misc_utils.seed_parse, default='25',
         help='Set the random seed value'
     )
@@ -89,8 +101,10 @@ def parse_args(args=None):
         gf_utils.gf_argparse(parser)
     except: pass
     llh_utils.likelihood_argparse(parser)
+    mn_utils.mn_argparse(parser)
     nuisance_argparse(parser)
     misc_utils.remove_option(parser, 'sigma_ratio')
+    misc_utils.remove_option(parser, 'mn_output')
     if args is None: return parser.parse_args()
     else: return parser.parse_args(args.split())
 
@@ -98,11 +112,74 @@ def parse_args(args=None):
 def gen_identifier(args):
     f = '_{0}_{1}'.format(*map(str_enum, (args.likelihood, args.data)))
     if args.data is not DataType.REAL:
-        ir1, ir2, ir3 = solve_ratio(args.injected_ratio)
-        f += '_INJ_{1:03d}_{2:03d}_{3:03d}'.format(ir1, ir2, ir3)
+        ir1, ir2, ir3 = misc_utils.solve_ratio(args.injected_ratio)
+        f += '_INJ_{0:03d}_{1:03d}_{2:03d}'.format(ir1, ir2, ir3)
     return f
 
 
+def gen_figtext(args, asimov_paramset):
+    f = ''
+    if args.data is DataType.REAL:
+        f += 'IceCube Preliminary\n'
+    else:
+        ir1, ir2, ir3 = misc_utils.solve_ratio(args.injected_ratio)
+        f += 'Injected ratio = [{0}, {1}, {2}]\n'.format(ir1, ir2, ir3)
+        for param in asimov_paramset:
+            f += 'Injected {0} = {1:.3f}\n'.format(
+                param.name, param.nominal_value
+            )
+    return f
+
+
+def triangle_llh(theta, args, hypo_paramset, fitter):
+    """Log likelihood function for a given theta."""
+    if len(theta) != len(hypo_paramset):
+        raise AssertionError(
+            'Dimensions of scan is not the same as the input '
+            'params\ntheta={0}\nparamset]{1}'.format(theta, hypo_paramset)
+        )
+    for idx, param in enumerate(hypo_paramset):
+        param.value = theta[idx]
+
+    if args.likelihood is Likelihood.GOLEMFIT:
+        llh = gf_utils.get_llh(fitter, hypo_paramset)
+    elif args.likelihood is Likelihood.GF_FREQ:
+        llh = gf_utils.get_llh_freq(fitter, hypo_paramset)
+
+    return llh
+
+
+def ln_prob(theta, args, hypo_paramset, fitter):
+    lp = llh_utils.lnprior(theta, paramset=hypo_paramset)
+    if not np.isfinite(lp):
+        return -np.inf
+    return lp + triangle_llh(
+        theta,
+        args           = args,
+        hypo_paramset  = hypo_paramset,
+        fitter         = fitter
+    )
+
+
+def lnProb(cube, ndim, n_params, hypo_paramset, args, fitter):
+    if ndim != len(hypo_paramset):
+        raise AssertionError(
+            'Length of MultiNest scan paramset is not the same as the input '
+            'params\ncube={0}\nmn_paramset]{1}'.format(cube, hypo_paramset)
+        )
+    pranges = hypo_paramset.ranges
+    for i in xrange(ndim):
+        hypo_paramset[i].value = (pranges[i][1]-pranges[i][0])*cube[i] + pranges[i][0]
+    theta = hypo_paramset.values
+    llh = ln_prob(
+        theta          = theta,
+        args           = args,
+        hypo_paramset  = hypo_paramset,
+        fitter         = fitter
+    )
+    return llh
+
+
 def main():
     args = parse_args()
     process_args(args)
@@ -111,13 +188,86 @@ def main():
     if args.seed is not None:
         np.random.seed(args.seed)
 
-    asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance())
+    asimov_paramset, hypo_paramset = get_paramsets(args, define_nuisance())
+    hypo_paramset.extend(asimov_paramset.from_tag(ParamTag.BESTFIT))
     outfile = args.outfile + gen_identifier(args)
     print '== {0:<25} = {1}'.format('outfile', outfile)
 
+    n_params = len(hypo_paramset)
+    prefix = outfile + '_mn_'
+    misc_utils.make_dir(prefix)
+
+    print 'asimov_paramset', asimov_paramset
+    print 'hypo_paramset', hypo_paramset
+
     if args.run_scan:
         fitter = gf_utils.setup_fitter(args, asimov_paramset)
 
+        lnProbEval = partial(
+            lnProb,
+            hypo_paramset  = hypo_paramset,
+            args           = args,
+            fitter         = fitter
+        )
+
+        print 'Running evidence calculation for {0}'.format(prefix)
+        run(
+            LogLikelihood              = lnProbEval,
+            Prior                      = mn_utils.CubePrior,
+            n_dims                     = n_params,
+            n_live_points              = args.mn_live_points,
+            evidence_tolerance         = args.mn_tolerance,
+            outputfiles_basename       = prefix,
+            importance_nested_sampling = True,
+            resume                     = False,
+            verbose                    = True
+        )
+
+    # Analyze
+    analyser = Analyzer(
+        outputfiles_basename=prefix, n_params=n_params
+    )
+    print analyser
+
+    pranges = hypo_paramset.ranges
+    fig_text = gen_figtext(args, asimov_paramset)
+
+    if args.plot_chains:
+        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]
+        plot_utils.chainer_plot(
+            infile       = chains,
+            outfile      = outfile[:5]+outfile[5:].replace('data', 'plots')+'_posterior',
+            outformat    = ['pdf'],
+            args         = args,
+            llh_paramset = hypo_paramset,
+            fig_text     = fig_text
+        )
+
+    if args.plot_triangle:
+        llh = analyser.get_data()[:,1]
+
+        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]
+        flavour_angles = chains[:,-2:]
+        flavour_ratios = np.array(
+            map(fr_utils.angles_to_fr, flavour_angles), dtype=np.float
+        )
+
+        plot_utils.triangle_project(
+            frs          = flavour_ratios,
+            llh          = llh,
+            outfile      = outfile[:5]+outfile[5:].replace('data', 'plots')+'_triangle',
+            outformat    = ['png'],
+            args         = args,
+            llh_paramset = hypo_paramset,
+            fig_text     = fig_text
+        )
+
     print "DONE!"
 
 
diff --git a/plot_llh/calc_fr_MCMC.py b/plot_llh/calc_fr_MCMC.py
deleted file mode 100644
index 1bf016a..0000000
--- a/plot_llh/calc_fr_MCMC.py
+++ /dev/null
@@ -1,76 +0,0 @@
-#! /usr/bin/env python
-from __future__ import absolute_import, division
-
-import sys
-sys.path.extend(['.', '..'])
-
-import numpy as np
-import tqdm
-
-from utils import fr as fr_utils
-from utils import misc as misc_utils
-from utils.enums import MixingScenario
-
-binning = np.logspace(np.log10(6e4), np.log10(1e7), 21)
-dimension = 6
-source = [0, 1, 0]
-scenario = MixingScenario.T13
-
-def get_fr(theta, source, binning, dimension, scenario):
-    sm_mixings = theta[:6]
-    nuisance   = theta[6:11]
-    scale      = np.power(10., theta[-1])
-
-    index = -nuisance[-1]
-
-    bin_centers = np.sqrt(binning[:-1]*binning[1:])
-    bin_width = np.abs(np.diff(binning))
-
-    # TODO(shivesh): test with astroNorm
-    source_flux = np.array(
-        [fr * np.power(bin_centers, index) for fr in source]
-    ).T
-
-    mass_eigenvalues = sm_mixings[-2:]
-    sm_u = fr_utils.angles_to_u(sm_mixings[:-2])
-
-    mf_perbin = []
-    for i_sf, sf_perbin in enumerate(source_flux):
-        u = fr_utils.params_to_BSMu(
-            theta             = [],
-            dim               = dimension,
-            energy            = bin_centers[i_sf],
-            mass_eigenvalues  = mass_eigenvalues,
-            sm_u              = sm_u,
-            no_bsm            = False,
-            fix_mixing        = scenario,
-            fix_mixing_almost = False,
-            fix_scale         = True,
-            scale             = 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./(binning[-1] - binning[0])) * \
-        intergrated_measured_flux
-    fr = averaged_measured_flux / np.sum(averaged_measured_flux)
-
-    return map(float, fr)
-
-if len(sys.argv)< 2:
-    print sys.argv
-    print "Usage: calc_fr_MCMC.py input_filepath."
-    exit(1)
-
-infile = sys.argv[1]
-outfile = infile[:-4] + '_proc.npy'
-
-d = np.load(infile)
-
-p = []
-for x in tqdm.tqdm(d, total=len(d)):
-    p.append(get_fr(x, source, binning, dimension, scenario))
-p = np.array(p)
-
-np.save(outfile, p)