Newer
Older
LbLcsFormFactors / plot-sample.py
@Elena Graverini Elena Graverini on 21 Mar 2017 3 KB Initial commit
# -*- coding: utf-8 -*-
# @Author: Elena Graverini
# @Date:   2017-03-16 11:20:45
# @Last Modified by:   Elena Graverini
# @Last Modified time: 2017-03-21 11:38:40
# import matplotlib
from __future__ import division
import sys, os, gc
import matplotlib.pyplot as plt
from eosdata import MCMCDataFile  # , Plotter1D
import pickle


def main():
    input_file = "constantFF/samples-mu-2.hdf5"
    save_name = 'test-output.pdf'
    if len(sys.argv) > 1:
        input_file = str(sys.argv[1])
    else:
        print('Warning: no input / output file names specified.\nUsing %s and %s.' % (input_file, save_name))
    if len(sys.argv) > 2:
        save_name = str(sys.argv[2])
    pklfilename = save_name.replace('.pdf', '.pkl')
    suffix = 'prescaled_10'
    if not os.path.isfile(pklfilename):
        datafile = MCMCDataFile(input_file)
        # print(datafile.parameters[0])
        gc.collect()
        data = datafile.data()[:, 0]
        d = {}
        reduced_data = data[::10]
        d['reduced_data'] = reduced_data
        n, bins, patches = plt.hist(reduced_data, bins=100, normed=True, alpha=0.3)
        bin_centers = (bins[:-1] + bins[1:]) / 2
        d['bin_contents'] = n
        d['bin_centers'] = bin_centers
        pickle.dump(d, open(pklfilename, 'wb'))
    else:
        d = pickle.load(open(pklfilename, 'rb'))
        plt.scatter(d['bin_centers'], d['bin_contents'])
        suffix += '_scatter'
    plt.savefig(save_name.replace('.pdf', '_%s.pdf' % suffix))
    # # plt.figure()
    # p1d = Plotter1D(datafile, save_name)
    # p1d.histogram(0, xmin=None, xmax=None, kde=None)


def superimpose(files, labels, outname):
    colors = ['#0078FF', '#FF6600', '#0AAFB6', '#FF3333', '#0000FF', '#00CC00', '#BF8040', '#FF33CC', '#FF7733']
    ds = []
    for f in files:
        ds.append(pickle.load(open(f, 'rb')))
    for i in range(len(files)):
        plt.plot(ds[i]['bin_centers'], ds[i]['bin_contents'], label=labels[i], color=colors[i], linewidth=2.0)
    plt.legend(loc='best')
    plt.savefig(outname)
    plt.close("all")


def makeFileFromMarzia(x, y, fname):
    ds = {
        'bin_centers': x,
        'bin_contents': y,
    }
    if not os.path.isfile(fname):
        os.system('mkdir -p %s' % os.path.dirname(fname))
        pickle.dump(ds, open(fname, 'wb'))


if __name__ == '__main__':
    # main()
    marzia_f = 'Marzia/Marzia_constantFF.pkl'
    mx = [0.2, 0.4, 0.80, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5]
    my = [0.0102553, 0.0110076, 0.0123026, 0.0128953, 0.0142323, 0.015366, 0.0162975, 0.0170258, 0.0175493, 0.0178644, 0.0179659, 0.0178455, 0.017491, 0.016884, 0.0159959, 0.0147803, 0.0131556, 0.0109562, 0.00773188]
    makeFileFromMarzia(mx, my, marzia_f)
    files_mu = ['constantFF/mcmc_gamma_muons.pkl', '1pole/mcmc_gamma_muons_1pole.pkl', '1pole_x_sqrtlambda/mcmc_gamma_muonss_1pole_x_sqrtlambda.pkl', marzia_f]
    out_mu = 'comparison_mu.pdf'
    files_tau = ['constantFF/mcmc_gamma_taus.pkl', '1pole/mcmc_gamma_taus_1pole.pkl', '1pole_x_sqrtlambda/mcmc_gamma_taus_1pole_x_sqrtlambda.pkl']
    out_tau = 'comparison_tau.pdf'
    labels = ['Constant FF', 'One pole', r'One pole $\times\; \sqrt{\lambda}$', 'MB: constant FF']
    superimpose(files_mu, labels, out_mu)
    superimpose(files_tau, labels, out_tau)