diff options
Diffstat (limited to 'chainer_plot.py')
| -rwxr-xr-x | chainer_plot.py | 265 |
1 files changed, 265 insertions, 0 deletions
diff --git a/chainer_plot.py b/chainer_plot.py new file mode 100755 index 0000000..26dc164 --- /dev/null +++ b/chainer_plot.py @@ -0,0 +1,265 @@ +#! /usr/bin/env python +""" +From an MCMC chains file, make a triangle plot. +""" + +from __future__ import absolute_import, division + +import sys, os +import errno + +import numpy as np +import matplotlib as mpl +mpl.use('Agg') +from matplotlib import rc, rcParams + +import getdist +from getdist import plots +from getdist import mcsamples + +import mcmc_scan + +rc('text', usetex=True) +rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18}) +cols = ['#29A2C6','#FF6D31','#FFCB18','#73B66B','#EF597B', '#333333'] + +font = {'family' : 'serif', + 'weight' : 'bold', + 'size' : 18} + + +def plot(infile, angles, outfile, measured_ratio, sigma_ratio, fix_sfr, + fix_mixing, fix_scale, source_ratio, scale, dimension, energy, scale_bounds): + """Make the triangle plot""" + if not angles: + if fix_mixing: + labels = [r'{\rm log}_{10}\Lambda', r'\phi_e', r'\phi_\mu', r'\phi_\tau'] + elif fix_sfr: + if fix_scale: + 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'] + else: + 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', \ + r'{\rm log}_{10}(\Lambda)'] + else: + if fix_scale: + 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', \ + r'{\rm log}_{10}\Lambda', r'\phi_e', r'\phi_\mu', r'\phi_\tau'] + else: + 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', \ + r'\phi_e', r'\phi_\mu', r'\phi_\tau'] + else: + if fix_sfr: + if fix_mixing: + assert 0 + labels=[r'\tilde{s}_{12}^2', r'{\rm log}_{10}\Lambda'] + elif fix_scale: + labels=[r'\tilde{s}_{12}^2', r'\tilde{c}_{13}^4', + r'\tilde{s}_{23}^2', r'\tilde{\delta_{CP}}'] + else: + labels=[r'\tilde{s}_{12}^2', r'\tilde{c}_{13}^4', + r'\tilde{s}_{23}^2', r'\tilde{\delta_{CP}}', + r'{\rm log}_{10}\Lambda'] + else: + if fix_mixing: + labels=[r'{\rm log}_{10}\Lambda', r'sin^4(\phi)', r'cos(2\psi)'] + elif fix_scale: + labels=[r'\tilde{s}_{12}^2', r'\tilde{c}_{13}^4', + r'\tilde{s}_{23}^2', r'\tilde{\delta_{CP}}', + r'sin^4(\phi)', r'cos(2\psi)'] + else: + labels=[r'\tilde{s}_{12}^2', r'\tilde{c}_{13}^4', + r'\tilde{s}_{23}^2', r'\tilde{\delta_{CP}}', + r'{\rm log}_{10}\Lambda', r'sin^4(\phi)', r'cos(2\psi)'] + print labels + + if not fix_scale: + s2 = np.log10(scale_bounds) + + if not angles: + if fix_mixing: + ranges = [s2, (0, 1), (0, 1), (0, 1)] + elif fix_sfr: + if fix_scale: + ranges = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1)] + else: + ranges = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), s2] + else: + if fix_scale: + ranges = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1)] + else: + ranges = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), s2, (0, 1), (0, 1), (0, 1)] + else: + if fix_sfr: + if fix_mixing: + ranges = [(0, 1), s2] + elif fix_scale: + ranges = [(0, 1), (0, 1), (0, 1), (0, 2*np.pi)] + else: + ranges = [(0, 1), (0, 1), (0, 1), (0, 2*np.pi), s2] + else: + if fix_mixing: + ranges = [s2, (0, 1), (-1, 1)] + elif fix_scale: + ranges = [(0, 1), (0, 1), (0, 1), (0, 2*np.pi), (0, 1), (-1, 1)] + else: + ranges = [(0, 1), (0, 1), (0, 1), (0, 2*np.pi), s2, (0, 1), (-1, 1)] + + def flat_angles_to_u(x): + return abs(mcmc_scan.angles_to_u(x)).astype(np.float32).flatten().tolist() + + raw = np.load(infile) + print 'raw.shape', raw.shape + if not angles: + if fix_mixing: + fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + sc_elements = raw[:,:-2] + Tchain = np.column_stack([sc_elements, fr_elements]) + elif fix_sfr: + if fix_scale: + Tchain = np.array(map(flat_angles_to_u, raw)) + else: + sc_elements = raw[:,-1:] + m_elements = np.array(map(flat_angles_to_u, raw[:,:-1])) + Tchain = np.column_stack([m_elements, sc_elements]) + else: + if fix_scale: + fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + m_elements = np.array(map(flat_angles_to_u, raw[:,:-2])) + Tchain = np.column_stack([m_elements, fr_elements]) + else: + fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + sc_elements = raw[:,-3:-2] + m_elements = np.array(map(flat_angles_to_u, raw[:,:-3])) + Tchain = np.column_stack([m_elements, sc_elements, fr_elements]) + else: + Tchain = raw + + if fix_sfr: + if fix_scale: + label = '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( + source_ratio[0], source_ratio[1], source_ratio[2], + measured_ratio[0], measured_ratio[1], measured_ratio[2], sigma_ratio, + dimension, int(energy), scale + ) + else: + label = '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( + source_ratio[0], source_ratio[1], source_ratio[2], + measured_ratio[0], measured_ratio[1], measured_ratio[2], sigma_ratio, + dimension, int(energy) + ) + else: + if fix_scale: + label = 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} GeV\nScale = {6}'.format( + measured_ratio[0], measured_ratio[1], measured_ratio[2], sigma_ratio, + dimension, int(energy), scale + ) + else: + label = 'IC observed flavour ratio = [{0:.2f}, {1:.2f}, {2:.2f}]\nSigma = {3:.3f}\nDimension = {4}\nEnergy = {5} GeV'.format( + measured_ratio[0], measured_ratio[1], measured_ratio[2], sigma_ratio, + dimension, int(energy) + ) + + Tsample = mcsamples.MCSamples( + samples=Tchain, labels=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, + ) + if fix_mixing and fix_sfr: + mpl.pyplot.figtext(0.4, 0.7, label, fontsize=4) + else: + mpl.pyplot.figtext(0.5, 0.7, label, fontsize=15) + print 'outfile = {0}'.format(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 + g.export(outfile) + + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + '--infile', type=str, required=True, + help='Path to MCMC chains' + ) + parser.add_argument( + '--angles', default=False, action='store_true', + help='Plot in terms of mixing angles' + ) + parser.add_argument( + '--outfile', type=str, default='./untitled.pdf', + help='Path to output plot' + ) + parser.add_argument( + '--bestfit-ratio', type=int, nargs=3, required=False, + help='Set the bestfit flavour ratio' + ) + parser.add_argument( + '--sigma-ratio', type=float, required=False, + help='Set the 1 sigma for the flavour ratio' + ) + parser.add_argument( + '--fix-sfr', action='store_true', + help='Fix the source flavour ratio' + ) + parser.add_argument( + '--fix-mixing', action='store_true', + help='Fix the new physics mixing values to a single term, s_12^2' + ) + parser.add_argument( + '--source-ratio', type=int, nargs=3, default=[2, 1, 0], + help='Set the source flavour ratio for the case when you want to fix it' + ) + parser.add_argument( + '--scale', type=float, required=False, + help='Fix the scale to this value' + ) + parser.add_argument( + '--dimension', type=int, default=3, help='Dimension' + ) + parser.add_argument( + '--energy', type=float, default=1000, help='Energy' + ) + parser.add_argument( + '--scale-bounds', type=float, nargs=2, + help='Upper and lower limits to plot the new physics scale' + ) + args = parser.parse_args() + return args + + +def main(): + args = vars(parse_args()) + plot(**args) + + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() |