Initial Commit

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()