diff options
Diffstat (limited to 'chainer_plot.py')
| -rwxr-xr-x | chainer_plot.py | 271 |
1 files changed, 0 insertions, 271 deletions
diff --git a/chainer_plot.py b/chainer_plot.py deleted file mode 100755 index 3facb30..0000000 --- a/chainer_plot.py +++ /dev/null @@ -1,271 +0,0 @@ -#! /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 - -from utils.fr import angles_to_u, angles_to_fr - -rc('text', usetex=False) -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)'] - labels = [r'convNorm', r'promptNorm', 'muonNorm', 'astroNorm', 'astroDeltaGamma'] + labels - print 'labels', 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)] - ranges = [(0, 5), (0, 5), (0, 5), (0, 5), (0, 5)] + ranges - print 'ranges', ranges - - def flat_angles_to_u(x): - return abs(angles_to_u(x)).astype(np.float32).flatten().tolist() - - raw = np.load(infile) - print 'raw.shape', raw.shape - if not angles: - nuisance, raw = raw[:,5:], raw[:,-5:] - if fix_mixing: - fr_elements = np.array(map(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(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(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]) - Tchain = np.column_stack([nuisance, Tchain]) - else: - Tchain = raw - print 'Tchain.shape', Tchain.shape - - 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() |
