From 402f8b53dd892b8fd44ae5ad45eac91b5f6b3750 Mon Sep 17 00:00:00 2001 From: Shivesh Mandalia Date: Fri, 28 Feb 2020 18:39:45 +0000 Subject: reogranise into a python package --- scripts/contour.py | 260 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 260 insertions(+) create mode 100755 scripts/contour.py (limited to 'scripts/contour.py') diff --git a/scripts/contour.py b/scripts/contour.py new file mode 100755 index 0000000..db9a933 --- /dev/null +++ b/scripts/contour.py @@ -0,0 +1,260 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : November 26, 2018 + +""" +HESE flavour ratio contour +""" + +from __future__ import absolute_import, division + +import os +import argparse +from copy import deepcopy +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 mcmc as mcmc_utils +from utils import plot as plot_utils +from utils.enums import str_enum +from utils.enums import DataType, Likelihood, MCMCSeedType, ParamTag +from utils.enums import PriorsCateg +from utils.param import Param, ParamSet + + +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='convNorm', value=1., seed=[0.5, 2. ], ranges=[0.1, 10.], std=0.4, tag=tag), + # Param(name='promptNorm', value=0., seed=[0., 6. ], ranges=[0., 20.], std=2.4, 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), + # Param(name='CRDeltaGamma', value=0., seed=[-0.1, 0.1 ], ranges=[-1., 1. ], std=0.1, tag=tag), + # Param(name='NeutrinoAntineutrinoRatio', value=1., seed=[0.8, 1.2 ], ranges=[0., 2. ], std=0.1, tag=tag), + # Param(name='anisotropyScale', value=1., seed=[0.8, 1.2 ], ranges=[0., 2. ], std=0.1, tag=tag), + # Param(name='domEfficiency', value=0.99, seed=[0.8, 1.2 ], ranges=[0.8, 1.2 ], std=0.1, tag=tag), + # Param(name='holeiceForward', value=0., seed=[-0.8, 0.8 ], ranges=[-4.42, 1.58 ], std=0.1, tag=tag), + # Param(name='piKRatio', value=1.0, seed=[0.8, 1.2 ], ranges=[0., 2. ], std=0.1, tag=tag) + ]) + return ParamSet(nuisance) + + +def get_paramsets(args, nuisance_paramset): + """Make the paramsets for generating the Asmimov MC sample and also running + the MCMC. + """ + asimov_paramset = [] + llh_paramset = [] + + gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)] + llh_paramset.extend(gf_nuisance) + + for parm in llh_paramset: + parm.value = args.__getattribute__(parm.name) + + llh_paramset = ParamSet(llh_paramset) + + if args.data is not DataType.REAL: + flavour_angles = fr_utils.fr_to_angles(args.injected_ratio) + else: + flavour_angles = fr_utils.fr_to_angles([1, 1, 1]) + + tag = ParamTag.BESTFIT + asimov_paramset.extend(gf_nuisance) + asimov_paramset.extend([ + Param(name='astroFlavorAngle1', value=flavour_angles[0], ranges=[ 0., 1.], std=0.2, tag=tag), + Param(name='astroFlavorAngle2', value=flavour_angles[1], ranges=[-1., 1.], std=0.2, tag=tag), + ]) + asimov_paramset = ParamSet(asimov_paramset) + + return asimov_paramset, llh_paramset + + +def nuisance_argparse(parser): + nuisance = define_nuisance() + for parm in nuisance: + parser.add_argument( + '--'+parm.name, type=float, default=parm.value, + help=parm.name+' to inject' + ) + +def process_args(args): + """Process the input args.""" + if args.data is not DataType.REAL: + args.injected_ratio = fr_utils.normalise_fr(args.injected_ratio) + + args.likelihood = Likelihood.GOLEMFIT + + args.mcmc_threads = misc_utils.thread_factors(args.threads)[0] + args.threads = misc_utils.thread_factors(args.threads)[1] + + +def parse_args(args=None): + """Parse command line arguments""" + parser = argparse.ArgumentParser( + description="BSM flavour ratio analysis", + formatter_class=misc_utils.SortingHelpFormatter, + ) + parser.add_argument( + '--injected-ratio', type=float, nargs=3, default=[1, 1, 1], + help='Set the central value for the injected flavour ratio at IceCube' + ) + parser.add_argument( + '--seed', type=misc_utils.seed_parse, default='26', + help='Set the random seed value' + ) + parser.add_argument( + '--threads', type=misc_utils.thread_type, default='1', + help='Set the number of threads to use (int or "max")' + ) + parser.add_argument( + '--datadir', type=str, default='./untitled', + help='Path to output results' + ) + gf_utils.gf_argparse(parser) + mcmc_utils.mcmc_argparse(parser) + nuisance_argparse(parser) + if args is None: return parser.parse_args() + else: return parser.parse_args(args.split()) + + +def gen_identifier(args): + f = '_{0}'.format(str_enum(args.data)) + if args.data is not DataType.REAL: + f += '_INJ_{0}'.format(misc_utils.solve_ratio(args.injected_ratio)) + return f + + +def gen_figtext(args): + """Generate the figure text.""" + t = r'$' + if args.data is DataType.REAL: + t += r'\textbf{IceCube\:Preliminary}$' + elif args.data in [DataType.ASIMOV, DataType.REALISATION]: + t += r'{\rm\bf IceCube\:Simulation}' + '$\n$' + t += r'\rm{Injected\:composition}'+r'\:=\:({0})_\oplus'.format( + solve_ratio(args.injected_ratio).replace('_', ':') + ) + '$' + return t + + +def triangle_llh(theta, args, hypo_paramset): + """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(hypo_paramset) + elif args.likelihood is Likelihood.GF_FREQ: + llh = gf_utils.get_llh_freq(hypo_paramset) + + return llh + + +def ln_prob(theta, args, hypo_paramset): + dc_hypo_paramset = deepcopy(hypo_paramset) + lp = llh_utils.lnprior(theta, paramset=dc_hypo_paramset) + if not np.isfinite(lp): + return -np.inf + return lp + triangle_llh( + theta, + args = args, + hypo_paramset = dc_hypo_paramset, + ) + + +def main(): + args = parse_args() + process_args(args) + misc_utils.print_args(args) + + if args.seed is not None: + np.random.seed(args.seed) + + asimov_paramset, hypo_paramset = get_paramsets(args, define_nuisance()) + hypo_paramset.extend(asimov_paramset.from_tag(ParamTag.BESTFIT)) + + prefix = '' + outfile = args.datadir + '/contour' + prefix + gen_identifier(args) + print '== {0:<25} = {1}'.format('outfile', outfile) + + print 'asimov_paramset', asimov_paramset + print 'hypo_paramset', hypo_paramset + + if args.run_mcmc: + gf_utils.setup_fitter(args, asimov_paramset) + + ln_prob_eval = partial( + ln_prob, + hypo_paramset = hypo_paramset, + args = args, + ) + + if args.mcmc_seed_type == MCMCSeedType.UNIFORM: + p0 = mcmc_utils.flat_seed( + hypo_paramset, nwalkers=args.nwalkers + ) + elif args.mcmc_seed_type == MCMCSeedType.GAUSSIAN: + p0 = mcmc_utils.gaussian_seed( + hypo_paramset, nwalkers=args.nwalkers + ) + + samples = mcmc_utils.mcmc( + p0 = p0, + ln_prob = ln_prob_eval, + ndim = len(hypo_paramset), + nwalkers = args.nwalkers, + burnin = args.burnin, + nsteps = args.nsteps, + args = args, + threads = args.mcmc_threads + ) + mcmc_utils.save_chains(samples, outfile) + + labels = [ + r'$N_{\rm conv}$', + r'$N_{\rm prompt}$', + r'$N_{\rm muon}$', + r'$N_{\rm astro}$', + r'$\gamma_{\rm astro}$', + r'$\text{sin}^4\phi_\oplus$', + r'$\text{cos}\left(2\psi_\oplus\right)$', + ] + + of = outfile[:5]+outfile[5:].replace('data', 'plots')+'_posterior' + plot_utils.chainer_plot( + infile = outfile+'.npy', + outfile = of, + outformat = ['pdf'], + args = args, + llh_paramset = hypo_paramset, + fig_text = gen_figtext(args), + labels = labels + ) + + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() -- cgit v1.2.3