make fig2 script

1 files changed, 184 insertions, 73 deletions

diff --git a/fig2.py b/fig2.py
index 1a0d19d..8bf108b 100755
--- a/fig2.py
+++ b/fig2.py
@@ -18,45 +18,60 @@ import numpy as np
 from utils import fr as fr_utils
 from utils import misc as misc_utils
 from utils import plot as plot_utils
-from utils.enums import str_enum
-from utils.enums import ParamTag, PriorsCateg
-from utils.param import Param, ParamSet
+from utils.plot import PLANCK_SCALE
 
 from matplotlib import pyplot as plt
+from matplotlib.patches import Circle
+from matplotlib.legend_handler import HandlerPatch
+import matplotlib.gridspec as gridspec
 
 
-def define_nuisance():
-    """Define the nuisance parameters."""
-    nuisance = []
-    tag = ParamTag.NUISANCE
-    lg_prior = PriorsCateg.LIMITEDGAUSS
-    nuisance.extend([
-        Param(name='convNorm',        value=1.,  seed=[0.5, 2. ], ranges=[0.1, 10.], std=0.4, prior=lg_prior, tag=tag),
-        Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 20.], std=2.4, prior=lg_prior, tag=tag),
-        Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0. , 10.], std=0.1, tag=tag),
-        Param(name='astroNorm',       value=6.9, seed=[0.,  5. ], ranges=[0. , 20.], std=1.5, tag=tag),
-        Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[-5., 5. ], std=0.1, tag=tag)
-    ])
-    return ParamSet(nuisance)
-
-
-def get_paramsets(args, nuisance_paramset):
-    paramset = []
-    gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)]
-    paramset.extend(gf_nuisance)
-    tag = ParamTag.BESTFIT
-    paramset.extend([
-        Param(name='astroFlavorAngle1', value=0, ranges=[0., 1.], std=0.2, tag=tag),
-        Param(name='astroFlavorAngle2', value=0, ranges=[-1., 1.], std=0.2, tag=tag),
-    ])
-    paramset = ParamSet(paramset)
-    return paramset
+DIM = 6
+NBINS = 25
+
+
+class HandlerCircle(HandlerPatch):
+    def create_artists(self, legend, orig_handle, xdescent, ydescent, width,
+                       height, fontsize, trans):
+        r = 10
+        x = r + width//2 + 10
+        y = height//2 - 3
+
+        # create
+        p = Circle(xy=(x, y), radius=r)
+
+        # update with data from original object
+        self.update_prop(p, orig_handle, legend)
+
+        # move xy to legend
+        p.set_transform(trans)
+        return [p]
+
+
+def cmap_discretize(cmap, N):
+    colors_i = np.concatenate((np.linspace(0, 1., N), (0.,0.,0.,0.)))
+    colors_rgba = cmap(colors_i)
+    indices = np.linspace(0, 1., N+1)
+    cdict = {}
+    for ki, key in enumerate(('red','green','blue')):
+        cdict[key] = [
+            (indices[i], colors_rgba[i-1,ki], colors_rgba[i,ki]) for i in xrange(N+1)
+        ]
+    # Return colormap object.
+    return mpl.colors.LinearSegmentedColormap(cmap.name + "_%d"%N, cdict, 1024)
+
+
+def alp(x):
+    y = list(x)
+    y[-1] = 0.7
+    return y
 
 
 def process_args(args):
     """Process the input args."""
     pass
 
+
 def parse_args(args=None):
     """Parse command line arguments"""
     parser = argparse.ArgumentParser(
@@ -76,60 +91,156 @@ def main():
     process_args(args)
     misc_utils.print_args(args)
 
-    paramset = get_paramsets(args, define_nuisance())
-    n_params = len(paramset)
-    print 'n_params', n_params
-
     prefix = ''
-    contour_infile = args.datadir + '/contour' + prefix + '_REAL.npy'
-    contour_chains = np.load(contour_infile)
 
-    flavour_angles = contour_chains[:,-2:]
-    flavour_ratios = np.array(
-        map(fr_utils.angles_to_fr, flavour_angles)
+    # Load HESE contour.
+    contour_infile = args.datadir + '/contour' + prefix + '/contour_REAL.npy'
+    contour_angles = np.load(contour_infile)[:,-2:]
+    contour_frs = np.array(map(fr_utils.angles_to_fr, contour_angles))
+
+    # Load mc_unitary.
+    mcu_basefile = args.datadir + '/mc_unitary' + prefix + '/mc_unitary_SRC_'
+    mcu_frs_120 = np.load(mcu_basefile + '1_2_0.npy')
+    mcu_frs_100 = np.load(mcu_basefile + '1_0_0.npy')
+    mcu_frs_010 = np.load(mcu_basefile + '0_1_0.npy')
+
+    # Load mc_texture.
+    mct_basefile = args.datadir + '/mc_texture' + prefix + '/mc_texture_SRC_'
+    mct_chains_010_OET = np.load(mct_basefile + '0_1_0_OET.npy')
+    mct_chains_100_OUT = np.load(mct_basefile + '1_0_0_OUT.npy')
+
+    # Caculate min/max scales.
+    min_scale = 1E100
+    max_scale = 1E-100
+
+    for i in (mct_chains_010_OET, mct_chains_100_OUT):
+        min_scale = min_scale if min_scale < np.min(i[:,-1]) else np.min(i[:,-1])
+        max_scale = max_scale if max_scale > np.max(i[:,-1]) else np.max(i[:,-1])
+    min_scale -= np.log10(PLANCK_SCALE[DIM])
+    max_scale -= np.log10(PLANCK_SCALE[DIM])
+
+    cmap_green = cmap_discretize(mpl.colors.LinearSegmentedColormap.from_list(
+        "", ["lime", "gold", "darkorange"]
+    ), 10)
+    cmap_blue = cmap_discretize(mpl.colors.LinearSegmentedColormap.from_list(
+        "", ["blue", "fuchsia", "darkmagenta"]
+    ), 10)
+
+    norm = mpl.colors.Normalize(vmin=min_scale, vmax=max_scale)
+    color_010 = map(alp, map(cmap_green, map(
+        norm, mct_chains_010_OET[:,-1]-np.log10(PLANCK_SCALE[DIM])
+    )))
+    color_100 = map(alp, map(cmap_blue, map(
+        norm, mct_chains_100_OUT[:,-1]-np.log10(PLANCK_SCALE[DIM])
+    )))
+
+    fontsize = 23
+    ax_labels = [r'$f_{e}^{\oplus}$', r'$f_{\mu}^{\oplus}$', r'$f_{\tau}^{\oplus}$']
+
+    # Figure
+    fig = plt.figure(figsize=(10, 10))
+    gs = gridspec.GridSpec(2, 1, height_ratios=[40, 1])
+    gs.update(hspace=0.3, wspace=0.15)
+
+    # Axis
+    ax = fig.add_subplot(gs[0])
+    tax = plot_utils.get_tax(ax, scale=NBINS, ax_labels=ax_labels)
+
+    # Plot HESE contour
+    coverages = {68: 'grey', 90: 'darkgrey'}
+    for cov in coverages.iterkeys():
+        plot_utils.flavour_contour(
+            frs = contour_frs,
+            ax = ax,
+            nbins = NBINS,
+            coverage = cov,
+            linewidth = 2.3,
+            color = coverages[cov],
+            alpha=0.6,
+            zorder=0
+        )
+    ax.text(0.34*NBINS, 0.143*NBINS, r'$68\%$', fontsize=fontsize, rotation=3)
+    ax.text(0.34*NBINS, 0.038*NBINS, r'$90\%$', fontsize=fontsize, rotation=0)
+
+    # Plot unitary contours
+    mcu_kwargs = {mcu_frs_120:{'color':'red',   'alpha':0.2, 'zorder':3},
+                  mcu_frs_100:{'color':'black', 'alpha':0.1, 'zorder':2},
+                  mcu_frs_010:{'color':'black', 'alpha':0.1, 'zorder':2}}
+    for frs, kwargs in mcu_kwargs.itervals():
+        plot_utils.flavour_contour(
+            frs = frs,
+            ax = ax,
+            fill = True,
+            nbins = NBINS,
+            coverage = 100,
+            linewidth = 1,
+            **kwargs
+        )
+
+    # Plot BSM points
+    tax.scatter(
+        mct_chains_010_OET[:,:-1]*NBINS, marker='o', s=0.03, color=color_010,
+        zorder=5
     )
-
-    nbins = 25
-
-    ax_labels = [r'$f_{e}$', r'$f_{\mu}$', r'$f_{\tau}$']
-
-    fig = plt.figure(figsize=(8, 8))
-    ax = fig.add_subplot(111)
-    tax = plot_utils.get_tax(ax, scale=nbins, ax_labels=ax_labels)
-
-    plot_utils.flavour_contour(
-        frs = flavour_ratios,
-        ax = ax,
-        nbins = nbins,
-        coverage = 99,
-        linewidth = 2,
-        color = 'green'
+    tax.scatter(
+        mct_chains_100_OUT[:,:-1]*NBINS, marker='o', s=0.03, color=color_100,
+        zorder=5
     )
 
-    plot_utils.flavour_contour(
-        frs = flavour_ratios,
-        ax = ax,
-        nbins = nbins,
-        coverage = 90,
-        linewidth = 2,
-        color = 'blue'
+    # Legend
+    legend_elements = []
+    legend_elements.append(
+        Circle((0., 0.), 0.1, facecolor='lime', alpha=0.7, edgecolor='k',
+               linewidth=2., label=r'$\left (0:1:0\right )\:w/\:{\rm New\:Physics}$')
     )
-
-    plot_utils.flavour_contour(
-        frs = flavour_ratios,
-        ax = ax,
-        nbins = nbins,
-        coverage = 68,
-        linewidth = 2,
-        color = 'red'
+    legend_elements.append(
+        Circle((0., 0.), 0.1, facecolor='blue', alpha=0.7, edgecolor='k',
+               linewidth=2., label=r'$\left (1:0:0\right )\:w/\:{\rm New\:Physics}$')
+    )
+    legend_elements.append(
+        Circle((0., 0.), 0.1, facecolor='red', alpha=0.7, edgecolor='k',
+               linewidth=2., label=r'$\left (1:2:0\right )$')
     )
+    legend_elements.append(
+        Circle((0., 0.), 0.1, facecolor='grey', alpha=0.7, edgecolor='k',
+               linewidth=2., label=r'$\left (0:1:0\right ) + \left (1:0:0\right )$')
+    )
+    legend = plt.legend(handles=legend_elements, loc=(0.65, 0.8),
+                        title='Source composition',
+                        fontsize=fontsize,
+                        handler_map={Circle: HandlerCircle()})
+    plt.setp(legend.get_title(), fontsize=fontsize)
+    legend.get_frame().set_linestyle('-')
+
+    # Colorbar
+    gs00 = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=gs[1])
+    ax0 = fig.add_subplot(gs00[0])
+    cb = mpl.colorbar.ColorbarBase(
+        ax0, cmap=cmap_010, norm=norm, orientation='horizontal'
+    )
+    cb.ax.tick_params(labelsize=fontsize-5)
+    ax0.text(
+        0.5, 2, r'$\mathcal{O}_{e\tau}\:texture$', fontsize=fontsize,
+        rotation=0, verticalalignment='center', horizontalalignment='center'
+    )
+
+
+    ax1 = fig.add_subplot(gs00[1])
+    cb = mpl.colorbar.ColorbarBase(ax1, cmap=cmap_100, norm=norm, orientation='horizontal')
+    cb.ax.tick_params(labelsize=fontsize-5)
+    ax1.text(0.5, 2, r'$\mathcal{O}_{\mu\tau}\:texture$', fontsize=fontsize,
+             rotation=0, verticalalignment='center', horizontalalignment='center')
 
-    ax.legend()
+    fig.text(0.5, 0.038, r'${\rm New\:Physics\:Scale}\:[\:{\rm log}_{10} (\Lambda_6\:/\:{\rm M}^{2}_{\rm Planck})\: ]$',
+             fontsize=fontsize+5, horizontalalignment='center')
 
+    outformat = ['png']
     outfile = args.datadir[:5]+args.datadir[5:].replace('data', 'plots')
-    outfile += '/fig2' + prefix + '.png'
-    print 'Saving plot as {0}'.format(outfile)
-    fig.savefig(outfile, bbox_inches='tight', dpi=150)
+    outfile += '/fig2' + prefix
+    make_dir(outfile)
+    for of in outformat:
+        print 'Saving plot as {0}'.format(outfile + '.' + of)
+        fig.savefig(outfile + '.' + of, bbox_inches='tight', dpi=150)
 
     print "DONE!"