From 932a8691e16eb904e3eec61daae08d72c2039f10 Mon Sep 17 00:00:00 2001 From: shivesh Date: Sun, 6 May 2018 21:22:47 -0500 Subject: Sun May 6 21:22:46 CDT 2018 --- plot_sens.py | 368 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 351 insertions(+), 17 deletions(-) (limited to 'plot_sens.py') diff --git a/plot_sens.py b/plot_sens.py index e51d55d..0c862eb 100755 --- a/plot_sens.py +++ b/plot_sens.py @@ -10,21 +10,355 @@ 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 matplotlib as mpl -mpl.use('Agg') -from matplotlib import rc -from matplotlib import pyplot as plt -from matplotlib.offsetbox import AnchoredText - -rc('text', usetex=False) -rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18}) - -FRS = [ - (1, 1, 1, 1, 2, 0), - (1, 1, 1, 0, 1, 0), - # (1, 1, 1, 1, 0, 0), - # (1, 1, 1, 0, 0, 1), -] - -DIMENSION = [3, 6] +import numpy.ma as ma + +from utils import fr as fr_utils +from utils import likelihood as llh_utils +from utils import misc as misc_utils +from utils import plot as plot_utils +from utils.enums import EnergyDependance, Likelihood, ParamTag +from utils.enums import PriorsCateg, SensitivityCateg, StatCateg +from utils.param import Param, ParamSet, get_paramsets + +from utils import multinest as mn_utils + + +def define_nuisance(): + """Define the nuisance parameters.""" + tag = ParamTag.SM_ANGLES + g_prior = PriorsCateg.GAUSSIAN + hg_prior = PriorsCateg.HALFGAUSS + e = 1e-9 + nuisance = [ + 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.995, 1-e], ranges=[0., 1.], std=0.00147, tex=r'c_{13}^4', prior=hg_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='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. , 50.], std=0.3, tag=tag), + 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) + ]) + return ParamSet(nuisance) + + +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.fix_mixing and args.fix_scale: + raise NotImplementedError('Fixed mixing and scale not implemented') + if args.fix_mixing and args.fix_mixing_almost: + raise NotImplementedError( + '--fix-mixing and --fix-mixing-almost cannot be used together' + ) + if args.fix_scale: + raise NotImplementedError( + '--fix-scale not implemented' + ) + + args.measured_ratio = fr_utils.normalise_fr(args.measured_ratio) + if args.fix_source_ratio: + assert len(args.source_ratios) % 3 == 0 + 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) + + 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 args.split_jobs and args.run_method is SensitivityCateg.FULL: + raise NotImplementedError( + 'split_jobs and run_method not implemented' + ) + + args.dimensions = np.sort(args.dimensions) + + args_copy = deepcopy(args) + scale_regions = [] + for dim in args.dimensions: + args_copy.dimension = dim + _, scale_region = fr_utils.estimate_scale(args_copy) + scale_regions.append(scale_region) + args.scale_region = [np.min(scale_regions), np.max(scale_regions)] + args.scale = np.power(10., np.average(np.log10(args.scale_region))) + + +def parse_args(args=None): + """Parse command line arguments""" + parser = argparse.ArgumentParser( + description="HESE BSM flavour ratio analysis plotting script", + formatter_class=misc_utils.SortingHelpFormatter, + ) + parser.add_argument( + '--infile', type=str, default='./untitled', + help='Path to input dir' + ) + parser.add_argument( + '--run-method', default='full', + type=partial(misc_utils.enum_parse, c=SensitivityCateg), + choices=SensitivityCateg, + help='Choose which type of sensivity plot to make' + ) + parser.add_argument( + '--stat-method', default='bayesian', + type=partial(misc_utils.enum_parse, c=StatCateg), choices=StatCateg, + help='Statistical method to employ' + ) + parser.add_argument( + '--sens-bins', type=int, default=10, + help='Number of bins for the Bayes factor plot' + ) + parser.add_argument( + '--split-jobs', type=misc_utils.parse_bool, default='False', + help='Did the jobs get split' + ) + parser.add_argument( + '--plot', type=misc_utils.parse_bool, default='True', + help='Make sensitivity plots' + ) + parser.add_argument( + '--plot-statistic', type=misc_utils.parse_bool, default='False', + help='Plot MultiNest evidence or LLH value' + ) + fr_utils.fr_argparse(parser) + llh_utils.likelihood_argparse(parser) + mn_utils.mn_argparse(parser) + nuisance_argparse(parser) + misc_utils.remove_option(parser, 'dimension') + misc_utils.remove_option(parser, 'source_ratio') + misc_utils.remove_option(parser, 'scale') + misc_utils.remove_option(parser, 'scale_region') + 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=[2, 1, 0], + help='Set the source flavour ratios for the case when you want to fix it' + ) + 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) + + asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance()) + + scale = llh_paramset.from_tag(ParamTag.SCALE)[0] + mmangles = llh_paramset.from_tag(ParamTag.MMANGLES) + if args.run_method is SensitivityCateg.FULL: + st_paramset_arr = [llh_paramset.from_tag(ParamTag.SCALE, invert=True)] + elif args.run_method in [SensitivityCateg.FIXED_ANGLE, SensitivityCateg.CORR_ANGLE]: + nscale_pmset = llh_paramset.from_tag([ParamTag.SCALE, ParamTag.MMANGLES], invert=True) + st_paramset_arr = [nscale_pmset] * 3 + elif args.run_method in [SensitivityCateg.FIXED_ONE_ANGLE, SensitivityCateg.CORR_ONE_ANGLE]: + nscale_pmset = llh_paramset.from_tag(ParamTag.SCALE, invert=True) + st_paramset_arr = [] + for x in xrange(3): + st_paramset_arr.append( + ParamSet([prms for prms in nscale_pmset + if mmangles[x].name != prms.name]) + ) + + corr_angles_categ = [SensitivityCateg.CORR_ANGLE, SensitivityCateg.CORR_ONE_ANGLE] + fixed_angle_categ = [SensitivityCateg.FIXED_ANGLE, SensitivityCateg.FIXED_ONE_ANGLE] + + if args.run_method in corr_angles_categ: + scan_angles = np.linspace(0+1e-9, 1-1e-9, args.sens_bins) + else: scan_angles = np.array([0]) + print 'scan_angles', scan_angles + + dims = len(args.dimensions) + srcs = len(args.source_ratios) + if args.run_method is SensitivityCateg.FULL: + statistic_arr = np.full((dims, srcs, args.sens_bins, 2), np.nan) + elif args.run_method in fixed_angle_categ: + statistic_arr = np.full((dims, srcs, len(st_paramset_arr), args.sens_bins, 2), np.nan) + elif args.run_method in corr_angles_categ: + statistic_arr = np.full( + (dims, srcs, len(st_paramset_arr), args.sens_bins, args.sens_bins, 3), np.nan + ) + + print 'Loading data' + for idim, dim in enumerate(args.dimensions): + argsc = deepcopy(args) + argsc.dimension = dim + _, scale_region = fr_utils.estimate_scale(argsc) + argsc.scale_region = scale_region + scan_scales = np.linspace( + np.log10(scale_region[0]), np.log10(scale_region[1]), args.sens_bins + ) + + for isrc, src in enumerate(args.source_ratios): + argsc.source_ratio = src + infile = args.infile + if args.likelihood is Likelihood.GOLEMFIT: + infile += '/golemfit/' + elif args.likelihood is Likelihood.GAUSSIAN: + infile += '/gaussian/' + infile += '/DIM{0}/fix_ifr/'.format(dim) + if args.likelihood is Likelihood.GAUSSIAN: + infile += '{0}/'.format(str(args.sigma_ratio).replace('.', '_')) + infile += 'fr_stat/{0}/{1}/fr_stat'.format(args.stat_method, args.run_method) \ + + misc_utils.gen_identifier(argsc) + print '== {0:<25} = {1}'.format('infile', infile) + + if args.split_jobs: + for idx_an, an in enumerate(scan_angles): + for idx_sc, sc in enumerate(scan_scales): + filename = infile + '_scale_{0:.0E}'.format(np.power(10, sc)) + try: + if args.run_method in fixed_angle_categ: + print 'Loading from {0}'.format(filename+'.npy') + statistic_arr[idim][isrc][:,idx_sc] = np.load(filename+'.npy')[:,0] + if args.run_method in corr_angles_categ: + filename += '_angle_{0:<04.2}'.format(an) + print 'Loading from {0}'.format(filename+'.npy') + statistic_arr[idim][isrc][:,idx_an,idx_sc] = np.load(filename+'.npy')[:,0,0] + except: + print 'Unable to load file {0}'.format(filename+'.npy') + continue + else: + print 'Loading from {0}'.format(infile+'.npy') + try: + statistic_arr[idim][isrc] = np.load(infile+'.npy') + except: + print 'Unable to load file {0}'.format(infile+'.npy') + continue + + print 'statistic_arr', statistic_arr + + data = ma.masked_invalid(statistic_arr) + if args.plot_statistic: + print 'Plotting statistic' + + argsc = deepcopy(args) + base_infile = args.infile + if args.likelihood is Likelihood.GOLEMFIT: + base_infile += '/golemfit/' + elif args.likelihood is Likelihood.GAUSSIAN: + base_infile += '/gaussian/' + for idim, dim in enumerate(args.dimensions): + argsc.dimension = dim + _, scale_region = fr_utils.estimate_scale(argsc) + argsc.scale_region = scale_region + infile = base_infile + '/DIM{0}/fix_ifr/'.format(dim) + if args.likelihood is Likelihood.GAUSSIAN: + infile += '{0}/'.format(str(args.sigma_ratio).replace('.', '_')) + infile += 'fr_stat/' + + for isrc, src in enumerate(args.source_ratios): + argsc.source_ratio = src + finfile = infile +'/{0}/{1}/fr_stat'.format(args.stat_method, args.run_method) \ + + misc_utils.gen_identifier(argsc) + basename = os.path.dirname(finfile) + '/statrun/' + os.path.basename(finfile) + baseoutfile = basename[:5]+basename[5:].replace('data', 'plots') + if args.run_method is SensitivityCateg.FULL: + outfile = baseoutfile + plot_utils.plot_statistic( + data = data[idim][isrc], + outfile = outfile, + outformat = ['png'], + args = argsc, + scale_param = scale, + ) + if args.run_method in fixed_angle_categ: + for idx_scen in xrange(len(st_paramset_arr)): + print '|||| SCENARIO = {0}'.format(idx_scen) + outfile = baseoutfile + '_SCEN{0}'.format(idx_scen) + if idx_scen == 0: label = r'$\mathcal{O}_{12}=\frac{\pi}{4}$' + elif idx_scen == 1: label = r'$\mathcal{O}_{13}=\frac{\pi}{4}$' + elif idx_scen == 2: label = r'$\mathcal{O}_{23}=\frac{\pi}{4}$' + plot_utils.plot_statistic( + data = data[idim][isrc][idx_scen], + outfile = outfile, + outformat = ['png'], + args = argsc, + scale_param = scale, + label = label + ) + elif args.run_method in corr_angles_categ: + for idx_scen in xrange(len(st_paramset_arr)): + print '|||| SCENARIO = {0}'.format(idx_scen) + basescenoutfile = baseoutfile + '_SCEN{0}'.format(idx_scen) + if idx_scen == 0: label = r'$\mathcal{O}_{12}=' + elif idx_scen == 1: label = r'$\mathcal{O}_{13}=' + elif idx_scen == 2: label = r'$\mathcal{O}_{23}=' + for idx_an, an in enumerate(scan_angles): + print '|||| ANGLE = {0:<04.2}'.format(float(an)) + outfile = basescenoutfile + '_ANGLE{0}'.format(idx_an) + _label = label + r'{0:<04.2}$'.format(an) + plot_utils.plot_statistic( + data = data[idim][isrc][idx_scen][idx_an][:,1:], + outfile = outfile, + outformat = ['png'], + args = argsc, + scale_param = scale, + label = _label + ) + + if args.plot: + print 'Plotting sensitivities' + + basename = args.infile + '{0}/{1}'.format(args.likelihood, args.stat_method) + baseoutfile = basename[:5]+basename[5:].replace('data', 'plots') + + if args.run_method is SensitivityCateg.FULL: + plot_utils.plot_sens_full( + data = data, + outfile = basename + '/FULL', + outformat = ['png'], + args = args, + ) + elif args.run_method in fixed_angle_categ: + plot_utils.plot_sens_fixed_angle( + data = data, + outfile = basename + '/FIXED_ANGLE', + outformat = ['png'], + args = args, + ) + elif args.run_method in corr_angles_categ: + plot_utils.plot_sens_corr_angle( + data = data, + outfile = basename + '/CORR_ANGLE', + outformat = ['png'], + args = args, + ) + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() -- cgit v1.2.3