diff options
Diffstat (limited to 'fr.py')
| -rwxr-xr-x | fr.py | 203 |
1 files changed, 112 insertions, 91 deletions
@@ -12,19 +12,18 @@ 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 mcmc as mcmc_utils from utils import misc as misc_utils from utils import plot as plot_utils -from utils.enums import EnergyDependance, Likelihood, MixingScenario from utils.enums import MCMCSeedType, ParamTag, PriorsCateg -from utils.param import Param, ParamSet, get_paramsets +from utils.param import Param, ParamSet def define_nuisance(): @@ -32,45 +31,45 @@ def define_nuisance(): tag = ParamTag.SM_ANGLES nuisance = [] g_prior = PriorsCateg.GAUSSIAN - hg_prior = PriorsCateg.LIMITEDGAUSS + 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=g_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=g_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=g_prior, tag=tag), + Param(name='s_12_2', value=0.307, seed=[0.26, 0.35], ranges=[0., 1.], std=0.013, tex=r's_{12}^2', 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', 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', 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), - Param( - name='m21_2', value=7.40E-23, seed=[7.2E-23, 7.6E-23], ranges=[6.80E-23, 8.02E-23], - std=2.1E-24, tex=r'\Delta m_{21}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag - ), - Param( - name='m3x_2', value=2.494E-21, seed=[2.46E-21, 2.53E-21], ranges=[2.399E-21, 2.593E-21], - std=3.3E-23, tex=r'\Delta m_{3x}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag - ) - # Param(name='s_12_2', value=0.307, ranges=[0., 1.], std=20, tex=r's_{12}^2', prior=g_prior, tag=tag), - # Param(name='c_13_4', value=(1-(0.02206))**2, ranges=[0., 1.], std=20, tex=r'c_{13}^4', prior=g_prior, tag=tag), - # Param(name='s_23_2', value=0.538, ranges=[0., 1.], std=20, tex=r's_{23}^2', prior=g_prior, tag=tag), - # Param(name='dcp', value=4.08404, ranges=[0., 2*np.pi], std=20, tex=r'\delta_{CP}', tag=tag), - # Param( - # name='m21_2', value=7.40E-23, ranges=[6.80E-23, 8.02E-23], - # std=2.1E-24, tex=r'\Delta m_{21}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag - # ), - # Param( - # name='m3x_2', value=2.494E-21, ranges=[2.399E-21, 2.593E-21], - # std=3.3E-23, tex=r'\Delta m_{3x}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag - # ) - ]) - tag = ParamTag.NUISANCE - 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): + """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)] + ) + + for parm in hypo_paramset: + parm.value = args.__getattribute__(parm.name) + + hypo_paramset = ParamSet(hypo_paramset) + + tag = ParamTag.BESTFIT + flavour_angles = fr_utils.fr_to_angles(args.injected_ratio) + + 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, hypo_paramset + + def nuisance_argparse(parser): nuisance = define_nuisance() for parm in nuisance: @@ -82,24 +81,7 @@ def nuisance_argparse(parser): def process_args(args): """Process the input args.""" - # if args.fix_mixing is not MixingScenario.NONE and args.fix_scale: - # raise NotImplementedError('Fixed mixing and scale not implemented') - if args.fix_mixing is not MixingScenario.NONE and args.fix_mixing_almost: - raise NotImplementedError( - '--fix-mixing and --fix-mixing-almost cannot be used together' - ) - - args.measured_ratio = fr_utils.normalise_fr(args.measured_ratio) - if args.fix_source_ratio: - args.source_ratio = fr_utils.normalise_fr(args.source_ratio) - - if args.energy_dependance is EnergyDependance.SPECTRAL: - args.binning = np.logspace( - np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+1 - ) - - if not args.fix_scale: - args.scale, args.scale_region = fr_utils.estimate_scale(args) + args.injected_ratio = fr_utils.normalise_fr(args.injected_ratio) def parse_args(args=None): @@ -109,7 +91,11 @@ def parse_args(args=None): formatter_class=misc_utils.SortingHelpFormatter, ) parser.add_argument( - '--seed', type=misc_utils.seed_parse, default='25', + '--injected-ratio', type=float, nargs=3, default=[1, 2, 0], + help='Set the central value for the injected flavour ratio at source' + ) + parser.add_argument( + '--seed', type=misc_utils.seed_parse, default='26', help='Set the random seed value' ) parser.add_argument( @@ -117,20 +103,57 @@ def parse_args(args=None): help='Set the number of threads to use (int or "max")' ) parser.add_argument( - '--outfile', type=str, default='./untitled', - help='Path to output results' + '--datadir', type=str, default='./untitled', + help='Path to store chains' ) - fr_utils.fr_argparse(parser) - try: - gf_utils.gf_argparse(parser) - except: pass - llh_utils.likelihood_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 = '_INJ_{0}'.format(misc_utils.solve_ratio(args.injected_ratio)) + return f + + +def gen_figtext(args, asimov_paramset): + f = '' + f += 'Injected ratio = {0}'.format( + misc_utils.solve_ratio(args.injected_ratio) + ) + for param in asimov_paramset: + f += '\nInjected {0:20s} = {1:.3f}'.format( + param.name, param.nominal_value + ) + return f + + +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) @@ -139,60 +162,58 @@ def main(): if args.seed is not None: np.random.seed(args.seed) - asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance()) - outfile = misc_utils.gen_outfile_name(args) - print '== {0:<25} = {1}'.format('outfile', outfile) + asimov_paramset, hypo_paramset = get_paramsets(args, define_nuisance()) - if args.run_mcmc: - if args.likelihood is Likelihood.GOLEMFIT: - fitter = gf_utils.setup_fitter(args, asimov_paramset) - else: fitter = None + prefix = '' + outfile = args.datadir + '/fr' + prefix + gen_identifier(args) + print '== {0:<25} = {1}'.format('outfile', outfile) - print 'asimov_paramset', asimov_paramset - print 'llh_paramset', llh_paramset + print 'asimov_paramset', asimov_paramset + print 'hypo_paramset', hypo_paramset - ln_prob = partial( - llh_utils.ln_prob, args=args, fitter=fitter, - asimov_paramset=asimov_paramset, llh_paramset=llh_paramset + if args.run_mcmc: + ln_prob_eval = partial( + ln_prob, + hypo_paramset = hypo_paramset, + args = args, ) - ndim = len(llh_paramset) if args.mcmc_seed_type == MCMCSeedType.UNIFORM: p0 = mcmc_utils.flat_seed( - llh_paramset, nwalkers=args.nwalkers + hypo_paramset, nwalkers=args.nwalkers ) elif args.mcmc_seed_type == MCMCSeedType.GAUSSIAN: p0 = mcmc_utils.gaussian_seed( - llh_paramset, nwalkers=args.nwalkers + hypo_paramset, nwalkers=args.nwalkers ) - if args.save_measured_fr: - n = misc_utils.gen_identifier(args) + '.txt' - f = args.output_measured_fr + n - if os.path.isfile(f): - open(f, 'w').close() - samples = mcmc_utils.mcmc( p0 = p0, - ln_prob = ln_prob, - ndim = ndim, + ln_prob = ln_prob_eval, + ndim = len(hypo_paramset), nwalkers = args.nwalkers, burnin = args.burnin, nsteps = args.nsteps, args = args, - threads = 1 - # TODO(shivesh): broken because you cannot pickle a GolemFitPy object - # threads = misc_utils.thread_factors(args.threads)[0] + threads = args.threads ) - mcmc_utils.save_chains(samples, outfile) + mmxs = map(fr_utils.angles_to_u, samples) + frs = np.array( + [fr_utils.u_to_fr(args.injected_ratio, x) for x in mmxs] + ) + mcmc_utils.save_chains(frs, outfile) + + of = outfile[:5]+outfile[5:].replace('data', 'plots')+'_posterior' plot_utils.chainer_plot( infile = outfile+'.npy', - outfile = outfile[:5]+outfile[5:].replace('data', 'plots'), - outformat = ['pdf'], + outfile = of, + outformat = ['png'], args = args, - llh_paramset = llh_paramset + llh_paramset = hypo_paramset, + fig_text = gen_figtext(args, hypo_paramset) ) + print "DONE!" |