diff options
Diffstat (limited to 'scripts/mc_x.py')
| -rwxr-xr-x | scripts/mc_x.py | 203 |
1 files changed, 203 insertions, 0 deletions
diff --git a/scripts/mc_x.py b/scripts/mc_x.py new file mode 100755 index 0000000..8c5e386 --- /dev/null +++ b/scripts/mc_x.py @@ -0,0 +1,203 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +emcee sample SM points for (x, 1-x, 0) +""" + +from __future__ import absolute_import, division + +import argparse +from copy import deepcopy +from functools import partial + +import numpy as np + +from utils import fr as fr_utils +from utils import llh as llh_utils +from utils import mcmc as mcmc_utils +from utils import misc as misc_utils +from utils import plot as plot_utils +from utils.enums import MCMCSeedType, ParamTag, PriorsCateg +from utils.param import Param, ParamSet + + +def define_nuisance(): + """Define the nuisance parameters.""" + tag = ParamTag.SM_ANGLES + nuisance = [] + g_prior = PriorsCateg.GAUSSIAN + lg_prior = PriorsCateg.LIMITEDGAUSS + e = 1e-9 + nuisance.extend([ + Param(name='s_12_2', value=0.307, seed=[0.26, 0.35], ranges=[0., 1.], std=0.013, tex=r's_{12}^2', prior=lg_prior, tag=tag), + Param(name='c_13_4', value=(1-(0.02206))**2, seed=[0.950, 0.961], ranges=[0., 1.], std=0.00147, tex=r'c_{13}^4', prior=lg_prior, tag=tag), + Param(name='s_23_2', value=0.538, seed=[0.31, 0.75], ranges=[0., 1.], std=0.069, tex=r's_{23}^2', prior=lg_prior, tag=tag), + Param(name='dcp', value=4.08404, seed=[0+e, 2*np.pi-e], ranges=[0., 2*np.pi], std=2.0, tex=r'\delta_{CP}', tag=tag), + ]) + tag = ParamTag.SRCANGLES + nuisance.extend([ + Param(name='astroX', value=0.5, seed=[0., 1.], ranges=[0., 1.], std=0.1, tex=r'x', 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 = [] + hypo_paramset = [] + + hypo_paramset.extend( + [x for x in nuisance_paramset.from_tag(( + ParamTag.SM_ANGLES, ParamTag.SRCANGLES + ))] + ) + + for parm in hypo_paramset: + parm.value = args.__getattribute__(parm.name) + + hypo_paramset = ParamSet(hypo_paramset) + + tag = ParamTag.BESTFIT + asimov_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), + ]) + asimov_paramset = ParamSet(asimov_paramset) + + return asimov_paramset, hypo_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.""" + pass + + +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( + '--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 store chains' + ) + mcmc_utils.mcmc_argparse(parser) + nuisance_argparse(parser) + misc_utils.remove_option(parser, 'injected_ratio') + misc_utils.remove_option(parser, 'plot_angles') + misc_utils.remove_option(parser, 'plot_elements') + if args is None: return parser.parse_args() + else: return parser.parse_args(args.split()) + + +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] + + return 1. # Flat 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()) + + prefix = '' + outfile = args.datadir + '/mc_x' + prefix + print '== {0:<25} = {1}'.format('outfile', outfile) + + print 'asimov_paramset', asimov_paramset + print 'hypo_paramset', hypo_paramset + + if args.run_mcmc: + 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.threads + ) + + nsamples = len(samples) + srcs = [fr_utils.normalise_fr((x, 1-x, 0)) for x in samples.T[-1]] + mmxs = map(fr_utils.angles_to_u, samples.T[:-1].T) + frs = np.array( + [fr_utils.u_to_fr(srcs[i], mmxs[i]) for i in xrange(nsamples)], + dtype=np.float64 + ) + mcmc_utils.save_chains(frs, outfile) + + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() + |