diff options
Diffstat (limited to 'utils')
| -rw-r--r-- | utils/enums.py | 8 | ||||
| -rw-r--r-- | utils/mcmc.py | 13 | ||||
| -rw-r--r-- | utils/misc.py | 44 | ||||
| -rw-r--r-- | utils/plot.py | 167 |
4 files changed, 218 insertions, 14 deletions
diff --git a/utils/enums.py b/utils/enums.py index 31885de..e154d0f 100644 --- a/utils/enums.py +++ b/utils/enums.py @@ -40,6 +40,14 @@ class Likelihood(Enum): GOLEMFIT = 3 +class ParamTag(Enum): + NUISANCE = 1 + MMANGLES = 2 + SCALE = 3 + SRCANGLES = 4 + NONE = 5 + + class Priors(Enum): UNIFORM = 1 LOGUNIFORM = 2 diff --git a/utils/mcmc.py b/utils/mcmc.py index d91764f..aebe12f 100644 --- a/utils/mcmc.py +++ b/utils/mcmc.py @@ -9,9 +9,6 @@ Useful functions to use an MCMC for the BSM flavour ratio analysis from __future__ import absolute_import, division -import os -import errno - import argparse from functools import partial @@ -23,7 +20,7 @@ from scipy.stats import multivariate_normal from utils import fr as fr_utils from utils.enums import Likelihood -from utils.misc import enum_keys, enum_parse, parse_bool +from utils.misc import enum_keys, enum_parse, make_dir, parse_bool def lnprior(theta, paramset): @@ -115,13 +112,7 @@ def save_chains(chains, outfile): Output file location of chains """ - try: - os.makedirs(outfile[:-len(os.path.basename(outfile))]) - except OSError as exc: # Python >2.5 - if exc.errno == errno.EEXIST and os.path.isdir(outfile[:-len(os.path.basename(outfile))]): - pass - else: - raise + make_dir(outfile) print 'Saving chains to location {0}'.format(outfile+'.npy') np.save(outfile+'.npy', chains) diff --git a/utils/misc.py b/utils/misc.py index 5c3eb2e..c54d25c 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -16,21 +16,23 @@ import multiprocessing import numpy as np -from utils.enums import Likelihood +from utils.enums import Likelihood, ParamTag class Param(object): """Parameter class to store parameters. """ - def __init__(self, name, value, ranges, std=None, tex=None): + def __init__(self, name, value, ranges, std=None, tex=None, tag=None): self._ranges = None self._tex = None + self._tag = None self.name = name self.value = value self.ranges = ranges self.std = std self.tex = tex + self.tag = tag @property def ranges(self): @@ -42,12 +44,23 @@ class Param(object): @property def tex(self): - return r'{0}'.format(self.tex) + return r'{0}'.format(self._tex) @tex.setter def tex(self, t): self._tex = t if t is not None else r'{\rm %s}' % self.name + @property + def tag(self): + return self._tag + + @tex.setter + def tag(self, t): + if t is None: self._tag = ParamTag.NONE + else: + assert t in ParamTag + self._tag = t + class ParamSet(Sequence): """Container class for a set of parameters. @@ -105,6 +118,10 @@ class ParamSet(Sequence): return tuple([obj.name for obj in self._params]) @property + def labels(self): + return tuple([obj.tex for obj in self._params]) + + @property def values(self): return tuple([obj.value for obj in self._params]) @@ -117,12 +134,23 @@ class ParamSet(Sequence): return tuple([obj.std for obj in self._params]) @property + def tags(self): + return tuple([obj.tag for obj in self._params]) + + @property def params(self): return self._params def to_dict(self): return {obj.name: obj.value for obj in self._params} + def from_tag(self, tag): + return tuple([obj for obj in self._params if obj.tag is tag]) + + def idx_from_tag(self, tag): + return tuple([idx for idx, obj in enumerate(self._params) + if obj.tag is tag]) + def gen_outfile_name(args): """Generate a name for the output file based on the input args. @@ -222,6 +250,16 @@ def enum_parse(s, c): return c[s.upper()] +def make_dir(outfile): + try: + os.makedirs(outfile[:-len(os.path.basename(outfile))]) + except OSError as exc: # Python >2.5 + if exc.errno == errno.EEXIST and os.path.isdir(outfile[:-len(os.path.basename(outfile))]): + pass + else: + raise + + def thread_type(t): if t.lower() == 'max': return multiprocessing.cpu_count() diff --git a/utils/plot.py b/utils/plot.py new file mode 100644 index 0000000..fb90b3e --- /dev/null +++ b/utils/plot.py @@ -0,0 +1,167 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 19, 2018 + +""" +Plotting functions for the BSM flavour ratio analysis +""" + +from __future__ import absolute_import, division + +import argparse + +import numpy as np +import matplotlib as mpl +mpl.use('Agg') +from matplotlib import rc + +import getdist +from getdist import plots +from getdist import mcsamples + +from utils import misc as misc_utils +from utils.enums import ParamTag +from utils.fr import angles_to_u, angles_to_fr + +rc('text', usetex=False) +rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18}) + + +def plot_argparse(parser): + """Arguments for plotting.""" + parser.add_argument( + '--plot-angles', type=misc_utils.parse_bool, default='True', + help='Plot MCMC triangle in the angles space' + ) + parser.add_argument( + '--plot-elements', type=misc_utils.parse_bool, default='False', + help='Plot MCMC triangle in the mixing elements space' + ) + + +def flat_angles_to_u(x): + """Convert from angles to mixing elements.""" + return abs(angles_to_u(x)).astype(np.float32).flatten().tolist() + + +def plot_Tchain(Tchain, axes_labels, ranges): + """Plot the Tchain using getdist.""" + Tsample = mcsamples.MCSamples( + samples=Tchain, labels=axes_labels, ranges=ranges + ) + + Tsample.updateSettings({'contours': [0.90, 0.99]}) + Tsample.num_bins_2D=500 + Tsample.fine_bins_2D=500 + Tsample.smooth_scale_2D=0.03 + + g = plots.getSubplotPlotter() + g.settings.num_plot_contours = 2 + g.settings.axes_fontsize = 10 + g.settings.figure_legend_frame = False + g.triangle_plot( + [Tsample], filled=True, + ) + return g + + +def gen_figtext(args): + """Generate the figure text.""" + mr1, mr2, mr3 = args.measured_ratio + if args.fix_source_ratio: + sr1, sr2, sr3 = args.source_ratio + if args.fix_scale: + return 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \ + 'observed flavour ratio = [{3:.2f}, {4:.2f}, ' \ + '{5:.2f}]\nSigma = {6:.3f}\nDimension = {7}\nEnergy = ' \ + '{8} GeV\nScale = {9}'.format( + sr1, sr2, sr3, mr1, mr2, mr3, args.sigma_ratio, + args.dimension, int(args.energy), args.scale + ) + else: + return 'Source flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nIC ' \ + 'observed flavour ratio = [{3:.2f}, {4:.2f}, ' \ + '{5:.2f}]\nSigma = {6:.3f}\nDimension = {7}\nEnergy = {8} ' \ + 'GeV'.format( + sr1, sr3, sr2, mr1, mr2, mr3, args.sigma_ratio, + args.dimension, int(args.energy) + ) + else: + if args.fix_scale: + return 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \ + '{2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} ' \ + 'GeV\nScale = {6}'.format( + mr1, mr2, mr3, args.sigma_ratio, args.dimension, + int(args.energy), args.scale + ) + else: + return 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, ' \ + '{2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} ' \ + 'GeV'.format( + mr1, mr2, mr3, args.sigma_ratio, args.dimension, + int(args.energy) + ) + +def chainer_plot(infile, outfile, outformat, args, mcmc_paramset): + """Make the triangle plot.""" + if not args.plot_angles and not args.plot_elements: + return + + raw = np.load(infile) + print 'raw.shape', raw.shape + + misc_utils.make_dir(outfile) + fig_text = gen_figtext(args) + + axes_labels = mcmc_paramset.labels + ranges = mcmc_paramset.ranges + + if args.plot_angles: + Tchain = raw + g = plot_Tchain(Tchain, axes_labels, ranges) + + if args.fix_mixing and args.fix_source_ratio: + mpl.pyplot.figtext(0.4, 0.7, fig_text, fontsize=4) + else: + mpl.pyplot.figtext(0.5, 0.7, fig_text, fontsize=15) + for of in outformat: + g.export(outfile+'_angles.'+of) + + if args.plot_elements: + nu_index = mcmc_paramset.idx_from_tag(ParamTag.NUISANCE) + fr_index = mcmc_paramset.idx_from_tag(ParamTag.MMANGLES) + sc_index = mcmc_paramset.idx_from_tag(ParamTag.SCALE) + sr_index = mcmc_paramset.idx_from_tag(ParamTag.SRCANGLES) + + nu_elements = raw[:,nu_index] + fr_elements = np.array(map(angles_to_fr, raw[:,fr_index])) + sc_elements = raw[:,sc_index] + sr_elements = np.array(map(flat_angles_to_u, raw[:,sr_index])) + Tchain = np.column_stack( + [nu_elements, fr_elements, sc_elements, sr_elements] + ) + + trns_ranges = ranges[nu_index,] + trns_axes_labels = axes_labels[nu_index,] + if not args.fix_mixing: + trns_axes_labels += \ + [r'\mid \tilde{U}_{e1} \mid' , r'\mid \tilde{U}_{e2} \mid' , r'\mid \tilde{U}_{e3} \mid' , \ + r'\mid \tilde{U}_{\mu1} \mid' , r'\mid \tilde{U}_{\mu2} \mid' , r'\mid \tilde{U}_{\mu3} \mid' , \ + r'\mid \tilde{U}_{\tau1} \mid' , r'\mid \tilde{U}_{\tau2} \mid' , r'\mid \tilde{U}_{\tau3} \mid'] + trns_ranges += [(0, 1)] * 9 + if not args.fix_scale: + trns_axes_labels += [axes_labels[sc_index]] + trns_ranges += [ranges[sc_index]] + if not args.fix_source_ratio: + trns_axes_labels += [r'\phi_e', r'\phi_\mu', r'\phi_\tau'] + trns_ranges += [(0, 1)] * 3 + + g = plot_Tchain(Tchain, trns_axes_labels, trns_ranges) + + if args.fix_mixing and args.fix_source_ratio: + mpl.pyplot.figtext(0.4, 0.7, fig_text, fontsize=4) + else: + mpl.pyplot.figtext(0.5, 0.7, fig_text, fontsize=15) + for of in outformat: + g.export(outfile+'_angles.'+of) |