diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/contour.py | 260 | ||||
| -rw-r--r-- | scripts/fr.py | 250 | ||||
| -rw-r--r-- | scripts/mc_texture.py | 233 | ||||
| -rw-r--r-- | scripts/mc_unitary.py | 203 | ||||
| -rw-r--r-- | scripts/mc_x.py | 203 | ||||
| -rw-r--r-- | scripts/plot_sens.py | 295 | ||||
| -rw-r--r-- | scripts/sens.py | 296 |
7 files changed, 1740 insertions, 0 deletions
diff --git a/scripts/contour.py b/scripts/contour.py new file mode 100644 index 0000000..6f08617 --- /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 golemflavor import fr as fr_utils +from golemflavor import gf as gf_utils +from golemflavor import llh as llh_utils +from golemflavor import misc as misc_utils +from golemflavor import mcmc as mcmc_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import str_enum +from golemflavor.enums import DataType, Likelihood, MCMCSeedType, ParamTag +from golemflavor.enums import PriorsCateg +from golemflavor.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() diff --git a/scripts/fr.py b/scripts/fr.py new file mode 100644 index 0000000..e32d7d0 --- /dev/null +++ b/scripts/fr.py @@ -0,0 +1,250 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +HESE BSM flavour ratio MCMC analysis script +""" + +from __future__ import absolute_import, division + +import os +import argparse +from functools import partial + +import numpy as np + +from golemflavor import fr as fr_utils +from golemflavor import gf as gf_utils +from golemflavor import llh as llh_utils +from golemflavor import mcmc as mcmc_utils +from golemflavor import misc as misc_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import DataType, Likelihood, MCMCSeedType +from golemflavor.enums import ParamTag, PriorsCateg, Texture +from golemflavor.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), + 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 + ) + ]) + 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 = [] + llh_paramset = [] + + gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)] + + llh_paramset.extend( + [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)] + ) + llh_paramset.extend(gf_nuisance) + + for parm in llh_paramset: + parm.value = args.__getattribute__(parm.name) + + boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension] + tag = ParamTag.SCALE + llh_paramset.append( + Param( + name='logLam', value=np.mean(boundaries), ranges=boundaries, std=3, + tex=r'{\rm log}_{10}\left (\Lambda^{-1}' + \ + misc_utils.get_units(args.dimension)+r'\right )', + tag=tag + ) + ) + llh_paramset = ParamSet(llh_paramset) + + tag = ParamTag.BESTFIT + 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]) + + 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.""" + args.source_ratio = fr_utils.normalise_fr(args.source_ratio) + if args.data is not DataType.REAL: + args.injected_ratio = fr_utils.normalise_fr(args.injected_ratio) + + args.binning = np.logspace( + np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+1 + ) + + args.likelihood = Likelihood.GOLEMFIT + + args.mcmc_threads = misc_utils.thread_factors(args.threads)[1] + args.threads = misc_utils.thread_factors(args.threads)[0] + + if args.texture is Texture.NONE: + raise ValueError('Must assume a BSM texture') + + +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='25', + 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' + ) + fr_utils.fr_argparse(parser) + gf_utils.gf_argparse(parser) + llh_utils.llh_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 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, llh_paramset = get_paramsets(args, define_nuisance()) + outfile = args.datadir + '/{0}/{1}/chains_'.format( + *map(misc_utils.parse_enum, [args.stat_method, args.data]) + ) + misc_utils.gen_identifier(args) + print '== {0:<25} = {1}'.format('outfile', outfile) + + if args.run_mcmc: + gf_utils.setup_fitter(args, asimov_paramset) + + print 'asimov_paramset', asimov_paramset + print 'llh_paramset', llh_paramset + + ln_prob = partial( + llh_utils.ln_prob, + args=args, + asimov_paramset=asimov_paramset, + llh_paramset=llh_paramset + ) + + if args.mcmc_seed_type == MCMCSeedType.UNIFORM: + p0 = mcmc_utils.flat_seed( + llh_paramset, nwalkers=args.nwalkers + ) + elif args.mcmc_seed_type == MCMCSeedType.GAUSSIAN: + p0 = mcmc_utils.gaussian_seed( + llh_paramset, nwalkers=args.nwalkers + ) + + samples = mcmc_utils.mcmc( + p0 = p0, + ln_prob = ln_prob, + ndim = len(llh_paramset), + nwalkers = args.nwalkers, + burnin = args.burnin, + nsteps = args.nsteps, + args = args, + threads = args.mcmc_threads + ) + mcmc_utils.save_chains(samples, outfile) + + raw = np.load(outfile+'.npy') + raw[:,4] *= 1E23 + raw[:,5] *= 1E21 + ranges = list(llh_paramset.ranges) + ranges[4] = [x*1E23 for x in ranges[4]] + ranges[5] = [x*1E21 for x in ranges[5]] + + labels = [ + r'${\rm sin}^2\theta_{12}$', + r'${\rm cos}^4\theta_{13}$', + r'${\rm sin}^2\theta_{23}$', + r'$\delta$', + r'$\Delta m_{21}^2\left[10^{-5}\,{\rm eV}^2\right]$', + r'$\Delta m_{31}^2\left[10^{-3}\,{\rm eV}^2\right]$', + r'$N_{\rm conv}$', + r'$N_{\rm prompt}$', + r'$N_{\rm muon}$', + r'$N_{\rm astro}$', + r'$\gamma_{\rm astro}$', + r'${\rm log}_{10}\left[\Lambda^{-1}_{'+ \ + r'{0}'.format(args.dimension)+r'}'+ \ + misc_utils.get_units(args.dimension)+r'\right]$' + ] + + plot_utils.chainer_plot( + infile = raw, + outfile = outfile[:5]+outfile[5:].replace('data', 'plots'), + outformat = ['pdf'], + args = args, + llh_paramset = llh_paramset, + labels = labels, + ranges = ranges + ) + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() diff --git a/scripts/mc_texture.py b/scripts/mc_texture.py new file mode 100644 index 0000000..2813dc0 --- /dev/null +++ b/scripts/mc_texture.py @@ -0,0 +1,233 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : April 25, 2019 + +""" +Sample points for a specific scenario +""" + +from __future__ import absolute_import, division + +import argparse +from copy import deepcopy +from functools import partial + +import numpy as np + +from golemflavor import fr as fr_utils +from golemflavor import llh as llh_utils +from golemflavor import mcmc as mcmc_utils +from golemflavor import misc as misc_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import MCMCSeedType, ParamTag, PriorsCateg, Texture +from golemflavor.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), + 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 + ) + ]) + 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 = [] + + llh_paramset.extend( + [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)] + ) + + for parm in llh_paramset: + parm.value = args.__getattribute__(parm.name) + + boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension] + tag = ParamTag.SCALE + llh_paramset.append( + Param( + name='logLam', value=np.mean(boundaries), ranges=boundaries, std=3, + tex=r'{\rm log}_{10}\left (\Lambda^{-1}' + \ + misc_utils.get_units(args.dimension)+r'\right )', + tag=tag + ) + ) + llh_paramset = ParamSet(llh_paramset) + + tag = ParamTag.BESTFIT + flavour_angles = fr_utils.fr_to_angles([1, 1, 1]) + 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.texture is Texture.NONE: + raise ValueError('Must assume a BSM texture') + args.source_ratio = fr_utils.normalise_fr(args.source_ratio) + + args.binning = np.logspace( + np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+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( + '--seed', type=misc_utils.seed_parse, default='25', + 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( + '--spectral-index', type=float, default='-2', + help='Astro spectral index' + ) + parser.add_argument( + '--datadir', type=str, default='./untitled', + help='Path to store chains' + ) + fr_utils.fr_argparse(parser) + 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 gen_identifier(args): + f = '_DIM{0}'.format(args.dimension) + f += '_SRC_' + misc_utils.solve_ratio(args.source_ratio) + f += '_{0}'.format(misc_utils.str_enum(args.texture)) + return f + + +def triangle_llh(theta, args, llh_paramset): + """Log likelihood function for a given theta.""" + if len(theta) != len(llh_paramset): + raise AssertionError( + 'Dimensions of scan is not the same as the input ' + 'params\ntheta={0}\nparamset]{1}'.format(theta, llh_paramset) + ) + for idx, param in enumerate(llh_paramset): + param.value = theta[idx] + + return 1. # Flat LLH + + +def ln_prob(theta, args, llh_paramset): + dc_llh_paramset = deepcopy(llh_paramset) + lp = llh_utils.lnprior(theta, paramset=dc_llh_paramset) + if not np.isfinite(lp): + return -np.inf + return lp + triangle_llh( + theta, + args = args, + llh_paramset = dc_llh_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, llh_paramset = get_paramsets(args, define_nuisance()) + + prefix = '' + outfile = args.datadir + '/mc_texture' + prefix + gen_identifier(args) + print '== {0:<25} = {1}'.format('outfile', outfile) + + print 'asimov_paramset', asimov_paramset + print 'llh_paramset', llh_paramset + + if args.run_mcmc: + ln_prob_eval = partial( + ln_prob, + llh_paramset = llh_paramset, + args = args, + ) + + if args.mcmc_seed_type == MCMCSeedType.UNIFORM: + p0 = mcmc_utils.flat_seed( + llh_paramset, nwalkers=args.nwalkers + ) + elif args.mcmc_seed_type == MCMCSeedType.GAUSSIAN: + p0 = mcmc_utils.gaussian_seed( + llh_paramset, nwalkers=args.nwalkers + ) + + samples = mcmc_utils.mcmc( + p0 = p0, + ln_prob = ln_prob_eval, + ndim = len(llh_paramset), + nwalkers = args.nwalkers, + burnin = args.burnin, + nsteps = args.nsteps, + args = args, + threads = args.threads + ) + + frs = np.array( + map(lambda x: fr_utils.flux_averaged_BSMu( + x, args, args.spectral_index, llh_paramset + ), samples), + dtype=float + ) + frs_scale = np.vstack((frs.T, samples[:-1].T)).T + mcmc_utils.save_chains(frs_scale, outfile) + + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() diff --git a/scripts/mc_unitary.py b/scripts/mc_unitary.py new file mode 100644 index 0000000..bbd5a28 --- /dev/null +++ b/scripts/mc_unitary.py @@ -0,0 +1,203 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Sample points only assuming unitarity +""" + +from __future__ import absolute_import, division + +import argparse +from copy import deepcopy +from functools import partial + +import numpy as np + +from golemflavor import fr as fr_utils +from golemflavor import llh as llh_utils +from golemflavor import mcmc as mcmc_utils +from golemflavor import misc as misc_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import MCMCSeedType, ParamTag, PriorsCateg +from golemflavor.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', 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), + ]) + 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.source_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: + parser.add_argument( + '--'+parm.name, type=float, default=parm.value, + help=parm.name+' to inject' + ) + + +def process_args(args): + """Process the input args.""" + args.source_ratio = fr_utils.normalise_fr(args.source_ratio) + + +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( + '--source-ratio', type=float, nargs=3, default=[1, 2, 0], + help='Set the source flavour ratio' + ) + 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, '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 gen_identifier(args): + f = '_SRC_{0}'.format(misc_utils.solve_ratio(args.source_ratio)) + 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) + 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_unitary' + 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: + 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 + ) + + mmxs = map(fr_utils.angles_to_u, samples) + frs = np.array( + [fr_utils.u_to_fr(args.source_ratio, x) for x in mmxs] + ) + mcmc_utils.save_chains(frs, outfile) + + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() diff --git a/scripts/mc_x.py b/scripts/mc_x.py new file mode 100644 index 0000000..04319c4 --- /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 golemflavor import fr as fr_utils +from golemflavor import llh as llh_utils +from golemflavor import mcmc as mcmc_utils +from golemflavor import misc as misc_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import MCMCSeedType, ParamTag, PriorsCateg +from golemflavor.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() + diff --git a/scripts/plot_sens.py b/scripts/plot_sens.py new file mode 100644 index 0000000..18e0c03 --- /dev/null +++ b/scripts/plot_sens.py @@ -0,0 +1,295 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : April 28, 2018 + +""" +HESE BSM flavour ratio analysis plotting script +""" + +from __future__ import absolute_import, division + +import os +import argparse +from functools import partial +from copy import deepcopy + +import numpy as np +import numpy.ma as ma + +from golemflavor import fr as fr_utils +from golemflavor import llh as llh_utils +from golemflavor import plot as plot_utils +from golemflavor.enums import DataType, Texture +from golemflavor.misc import enum_parse, parse_bool, parse_enum, print_args +from golemflavor.misc import gen_identifier, SortingHelpFormatter +from golemflavor.param import Param, ParamSet + + +MASK_X = (0.3, 0.7) + + +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) + + # Anon points + anon = [] + if args.dimensions[0] == 3: + anon.append([0.825, 0.845]) + anon.append([0.865, 0.875]) + anon.append([0.875, 0.885]) + anon.append([0.905, 0.915]) + anon.append([0.925, 0.935]) + if args.dimensions[0] == 4: + anon.append([0.165, 0.175]) + anon.append([0.805, 0.825]) + anon.append([0.835, 0.845]) + anon.append([0.855, 0.885]) + anon.append([0.965, 0.975]) + if args.dimensions[0] == 5: + anon.append([0.895, 0.905]) + anon.append([0.955, 0.965]) + if args.dimensions[0] == 6: + anon.append([0.115, 0.125]) + anon.append([0.855, 0.865]) + if args.dimensions[0] == 7: + # anon.append([0.815, 0.835]) + anon.append([0.875, 0.885]) + if args.dimensions[0] == 8: + anon.append([0.915, 0.935]) + anon.append([0.875, 0.895]) + anon.append([0.845, 0.855]) + + if args.source_ratios is not None: + if args.x_segments is not None: + raise ValueError('Cannot do both --source-ratios and --x-segments') + if len(args.source_ratios) % 3 != 0: + raise ValueError( + 'Invalid source ratios input {0}'.format(args.source_ratios) + ) + + srs = [args.source_ratios[3*x:3*x+3] + for x in range(int(len(args.source_ratios)/3))] + args.source_ratios = map(fr_utils.normalise_fr, srs) + elif args.x_segments is not None: + x_array = np.linspace(0, 1, args.x_segments) + sources = [] + for x in x_array: + if x >= MASK_X[0] and x <= MASK_X[1]: continue + skip = False + for a in anon: + if (a[1] > x) & (x > a[0]): + print 'Skipping src', x + skip = True + break + if skip: continue + sources.append([x, 1-x, 0]) + args.source_ratios = sources + else: + raise ValueError('Must supply either --source-ratios or --x-segments') + + args.dimensions = np.sort(args.dimensions) + + +def parse_args(args=None): + """Parse command line arguments""" + parser = argparse.ArgumentParser( + description="HESE BSM flavour ratio analysis plotting script", + formatter_class=SortingHelpFormatter, + ) + parser.add_argument( + '--datadir', type=str, + default='/data/user/smandalia/flavour_ratio/data/sensitivity', + help='Path to data directory' + ) + parser.add_argument( + '--segments', type=int, default=10, + help='Number of new physics scales to evaluate' + ) + parser.add_argument( + '--split-jobs', type=parse_bool, default='True', + help='Did the jobs get split' + ) + parser.add_argument( + '--dimensions', type=int, nargs='*', default=[3, 6], + help='Set the new physics dimensions to consider' + ) + parser.add_argument( + '--source-ratios', type=int, nargs='*', default=None, + required=False, help='Set the source flavour ratios' + ) + parser.add_argument( + '--x-segments', type=int, default=None, + required=False, help='Number of segments in x' + ) + parser.add_argument( + '--texture', type=partial(enum_parse, c=Texture), + default='none', choices=Texture, help='Set the BSM mixing texture' + ) + parser.add_argument( + '--data', default='asimov', type=partial(enum_parse, c=DataType), + choices=DataType, help='select datatype' + ) + parser.add_argument( + '--plot-x', type=parse_bool, default='False', + help='Make sensitivity plot x vs limit' + ) + parser.add_argument( + '--plot-table', type=parse_bool, default='False', + help='Make sensitivity table plot' + ) + parser.add_argument( + '--plot-statistic', type=parse_bool, default='False', + help='Plot MultiNest evidence or LLH value' + ) + llh_utils.llh_argparse(parser) + if args is None: return parser.parse_args() + else: return parser.parse_args(args.split()) + + +def main(): + args = parse_args() + process_args(args) + print_args(args) + + dims = len(args.dimensions) + srcs = len(args.source_ratios) + if args.texture is Texture.NONE: + textures = [Texture.OEU, Texture.OET, Texture.OUT] + if args.plot_table: + textures = [Texture.OET, Texture.OUT] + else: + textures = [args.texture] + texs = len(textures) + + prefix = '' + # prefix = 'noprior' + + # Initialise data structure. + statistic_arr = np.full((dims, srcs, texs, args.segments, 2), np.nan) + + print 'Loading data' + argsc = deepcopy(args) + for idim, dim in enumerate(args.dimensions): + argsc.dimension = dim + + datadir = args.datadir + '/DIM{0}'.format(dim) + # Array of scales to scan over. + boundaries = fr_utils.SCALE_BOUNDARIES[argsc.dimension] + eval_scales = np.linspace( + boundaries[0], boundaries[1], args.segments-1 + ) + eval_scales = np.concatenate([[-100.], eval_scales]) + + for isrc, src in enumerate(args.source_ratios): + argsc.source_ratio = src + + for itex, texture in enumerate(textures): + argsc.texture = texture + + base_infile = datadir + '/{0}/{1}/'.format( + *map(parse_enum, [args.stat_method, args.data]) + ) + r'{0}/fr_stat'.format(prefix) + gen_identifier(argsc) + + print '== {0:<25} = {1}'.format('base_infile', base_infile) + + if args.split_jobs: + for idx_sc, scale in enumerate(eval_scales): + infile = base_infile + '_scale_{0:.0E}'.format( + np.power(10, scale) + ) + try: + print 'Loading from {0}'.format(infile+'.npy') + statistic_arr[idim][isrc][itex][idx_sc] = \ + np.load(infile+'.npy')[0] + except: + print 'Unable to load file {0}'.format( + infile+'.npy' + ) + # raise + continue + else: + print 'Loading from {0}'.format(base_infile+'.npy') + try: + statistic_arr[idim][isrc][itex] = \ + np.load(base_infile+'.npy') + except: + print 'Unable to load file {0}'.format( + base_infile+'.npy' + ) + continue + + data = ma.masked_invalid(statistic_arr) + + print 'data', data + if args.plot_statistic: + print 'Plotting statistic' + + for idim, dim in enumerate(args.dimensions): + argsc.dimension = dim + + # Array of scales to scan over. + boundaries = fr_utils.SCALE_BOUNDARIES[argsc.dimension] + eval_scales = np.linspace( + boundaries[0], boundaries[1], args.segments-1 + ) + eval_scales = np.concatenate([[-100.], eval_scales]) + + for isrc, src in enumerate(args.source_ratios): + argsc.source_ratio = src + for itex, texture in enumerate(textures): + argsc.texture = texture + + base_infile = args.datadir + '/DIM{0}/{1}/{2}/'.format( + dim, *map(parse_enum, [args.stat_method, args.data]) + ) + r'{0}/fr_stat'.format(prefix) + gen_identifier(argsc) + basename = os.path.dirname(base_infile) + outfile = basename[:5]+basename[5:].replace('data', 'plots') + outfile += '/' + os.path.basename(base_infile) + + label = r'$\text{Texture}=' + plot_utils.texture_label(texture)[1:] + plot_utils.plot_statistic( + data = data[idim][isrc][itex], + outfile = outfile, + outformat = ['png'], + args = argsc, + scale_param = scale, + label = label + ) + + basename = args.datadir[:5]+args.datadir[5:].replace('data', 'plots') + baseoutfile = basename + '/{0}/{1}/'.format( + *map(parse_enum, [args.stat_method, args.data]) + ) + r'{0}'.format(prefix) + + argsc = deepcopy(args) + if args.plot_x: + for idim, dim in enumerate(args.dimensions): + print '|||| DIM = {0}'.format(dim) + argsc.dimension = dim + plot_utils.plot_x( + data = data[idim], + outfile = baseoutfile + '/hese_x_DIM{0}'.format(dim), + outformat = ['png', 'pdf'], + args = argsc, + normalise = True + ) + + if args.plot_table: + plot_utils.plot_table_sens( + data = data, + outfile = baseoutfile + '/hese_table', + outformat = ['png', 'pdf'], + args = args, + show_lvatmo = True + ) + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() diff --git a/scripts/sens.py b/scripts/sens.py new file mode 100644 index 0000000..6368159 --- /dev/null +++ b/scripts/sens.py @@ -0,0 +1,296 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +HESE BSM flavour ratio analysis script +""" + +from __future__ import absolute_import, division + +import os +import argparse +from functools import partial + +import glob + +import numpy as np +import numpy.ma as ma +from scipy.optimize import minimize + +from golemflavor import fr as fr_utils +from golemflavor import gf as gf_utils +from golemflavor import llh as llh_utils +from golemflavor import misc as misc_utils +from golemflavor import mn as mn_utils +from golemflavor.enums import str_enum +from golemflavor.enums import DataType, Likelihood, ParamTag +from golemflavor.enums import PriorsCateg, StatCateg, Texture +from golemflavor.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), + 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 + ) + ]) + 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=8.0, 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 = [] + llh_paramset = [] + + gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)] + + llh_paramset.extend( + [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)] + ) + llh_paramset.extend(gf_nuisance) + + for parm in llh_paramset: + parm.value = args.__getattribute__(parm.name) + + boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension] + tag = ParamTag.SCALE + llh_paramset.append( + Param( + name='logLam', value=np.mean(boundaries), ranges=boundaries, std=3, + tex=r'{\rm log}_{10}\left (\Lambda^{-1}' + \ + misc_utils.get_units(args.dimension)+r'\right )', + tag=tag + ) + ) + llh_paramset = ParamSet(llh_paramset) + + tag = ParamTag.BESTFIT + 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]) + + 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.""" + args.source_ratio = fr_utils.normalise_fr(args.source_ratio) + if args.data is not DataType.REAL: + args.injected_ratio = fr_utils.normalise_fr(args.injected_ratio) + + args.binning = np.logspace( + np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+1 + ) + + if args.eval_segment.lower() == 'all': + args.eval_segment = None + else: + args.eval_segment = int(args.eval_segment) + + if args.stat_method is StatCateg.BAYESIAN: + args.likelihood = Likelihood.GOLEMFIT + elif args.stat_method is StatCateg.FREQUENTIST: + args.likelihood = Likelihood.GF_FREQ + + if args.texture is Texture.NONE: + raise ValueError('Must assume a BSM texture') + + +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='25', + 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' + ) + parser.add_argument( + '--segments', type=int, default=10, + help='Number of new physics scales to evaluate' + ) + parser.add_argument( + '--eval-segment', type=str, default='all', + help='Which point to evalaute' + ) + parser.add_argument( + '--overwrite', type=misc_utils.parse_bool, default='False', + help='Overwrite chains' + ) + fr_utils.fr_argparse(parser) + gf_utils.gf_argparse(parser) + llh_utils.llh_argparse(parser) + mn_utils.mn_argparse(parser) + nuisance_argparse(parser) + 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) + + if args.seed is not None: + np.random.seed(args.seed) + + asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance()) + + # Scale and BSM mixings will be fixed. + scale_prm = llh_paramset.from_tag(ParamTag.SCALE)[0] + base_mn_pset = llh_paramset.from_tag(ParamTag.SCALE, invert=True) + + # Array of scales to scan over. + boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension] + eval_scales = np.linspace(boundaries[0], boundaries[1], args.segments-1) + eval_scales = np.concatenate([[-100.], eval_scales]) + + # Evaluate just one point (job), or all points. + if args.eval_segment is None: + eval_dim = args.segments + else: eval_dim = 1 + + outfile = args.datadir + '/{0}/{1}/fr_stat'.format( + *map(misc_utils.parse_enum, [args.stat_method, args.data]) + ) + misc_utils.gen_identifier(args) + + if not args.overwrite and os.path.isfile(outfile+'.npy'): + print 'FILE EXISTS {0}'.format(outfile+'.npy') + print 'Exiting...' + return + + # Setup Golemfit. + if args.run_mn: + gf_utils.setup_fitter(args, asimov_paramset) + + # Initialise data structure. + stat_arr = np.full((eval_dim, 2), np.nan) + + for idx_sc, scale in enumerate(eval_scales): + if args.eval_segment is not None: + if idx_sc == args.eval_segment: + outfile += '_scale_{0:.0E}'.format(np.power(10, scale)) + else: continue + print '|||| SCALE = {0:.0E}'.format(np.power(10, scale)) + + if not args.overwrite and os.path.isfile(outfile+'.npy'): + print 'FILE EXISTS {0}'.format(outfile+'.npy') + t = np.load(outfile+'.npy') + if np.any(~np.isfinite(t)): + print 'nan found, rerunning...' + pass + else: + print 'Exiting...' + return + + # Lower scale boundary for first (NULL) point and set the scale param. + reset_range = None + if scale < scale_prm.ranges[0]: + reset_range = scale_prm.ranges + scale_prm.ranges = (scale, scale_prm.ranges[1]) + scale_prm.value = scale + + identifier = 'b{0}_{1}_{2}_sca{3}'.format( + args.eval_segment, args.segments, str_enum(args.texture), scale + ) + llh = '{0}'.format(args.likelihood).split('.')[1] + data = '{0}'.format(args.data).split('.')[1] + src_string = misc_utils.solve_ratio(args.source_ratio) + prefix = args.mn_output + '/DIM{0}/{1}/{2}/s{3}/{4}'.format( + args.dimension, data, llh, src_string, identifier + ) + try: + stat = mn_utils.mn_evidence( + mn_paramset = base_mn_pset, + llh_paramset = llh_paramset, + asimov_paramset = asimov_paramset, + args = args, + prefix = prefix + ) + except: + print 'Failed run' + raise + print '## Evidence = {0}'.format(stat) + + if args.eval_segment is not None: + stat_arr[0] = np.array([scale, stat]) + else: + stat_arr[idx_sc] = np.array([scale, stat]) + + # Cleanup. + if reset_range is not None: + scale_prm.ranges = reset_range + + if args.run_mn and not args.debug: + try: + for f in glob.glob(prefix + '*'): + print 'cleaning file {0}'.format(f) + os.remove(f) + except: + print 'got error trying to cleanup, continuing' + pass + + misc_utils.make_dir(outfile) + print 'Saving to {0}'.format(outfile+'.npy') + np.save(outfile+'.npy', stat_arr) + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() |