aboutsummaryrefslogtreecommitdiff
path: root/src/run-fit
blob: 203aa592a1825e810dd6936ca5bf3fd2005ac43e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
#!/usr/bin/env python
from __future__ import print_function, division
import subprocess
from os.path import splitext, split

def run_fit(filename):
    head, tail = split(filename)
    root, ext = splitext(tail)
    output = root + '.txt'
    cmd = ["./fit", filename, "-o", output]
    subprocess.call(cmd)

if __name__ == '__main__':
    import argparse
    from multiprocessing import Pool, cpu_count
    import signal
    import os

    parser = argparse.ArgumentParser("fit multiple zdab files")
    parser.add_argument("-j", "--jobs", type=int, default=None,
                        help="number of jobs")
    parser.add_argument("filenames", nargs="+",
                        help="zdab files")
    args = parser.parse_args()

    jobs = args.jobs

    if jobs is None:
        jobs = cpu_count()

    # see https://stackoverflow.com/questions/11312525/catch-ctrlc-sigint-and-exit-multiprocesses-gracefully-in-python
    handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
    p = Pool(jobs)
    signal.signal(signal.SIGINT, handler)
    try:
        result = p.map_async(run_fit,args.filenames)
        result.get()
    except KeyboardInterrupt:
        print("ctrl-c caught. quitting...")
        p.terminate()
        os.killpg(os.getpgid(0),9)
    else:
        p.close()

    p.join()
ef='#n409'>409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839
#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to do final null hypothesis test that the events in the 20 MeV - 10 GeV
range are consistent with atmospheric neutrino events. To run it just run:

    $ ./chi2 [list of data fit results] --mc [list of atmospheric MC files] --muon-mc [list of muon MC files] --steps [steps]

After running you will get a plot showing the best fit results of the MC and
data along with p-values for each of the possible particle combinations (single
electron, single muon, double electron, etc.).
"""
from __future__ import print_function, division
import numpy as np
from scipy.stats import iqr, poisson
from matplotlib.lines import Line2D
from scipy.stats import iqr, norm, beta, percentileofscore
from scipy.special import spence
from sddm.stats import *
from sddm.dc import estimate_errors, EPSILON, truncnorm_scaled
import emcee
from sddm import printoptions
from sddm.utils import fast_cdf, correct_energy_bias
import nlopt
from itertools import chain

# Likelihood Fit Parameters
# 0 - Atmospheric Neutrino Flux Scale
# 1 - Electron energy bias
# 2 - Electron energy resolution
# 3 - Muon energy bias
# 4 - Muon energy resolution
# 5 - External Muon scale

FIT_PARS = [
    'Atmospheric Neutrino Flux Scale',
    'Electron energy bias',
    'Electron energy resolution',
    'Muon energy bias',
    'Muon energy resolution',
    'External Muon scale']

# Uncertainty on the energy scale
#
# - the muon energy scale and resolution terms come directly from measurements
#   on stopping muons, so those are known well.
# - for electrons, we only have Michel electrons at the low end of our energy
#   range, and therefore we don't really have any good way of constraining the
#   energy scale or resolution. However, if we assume that the ~7% energy bias
#   in the muons is from the single PE distribution (it seems likely to me that
#   that is a major part of the bias), then the energy scale should be roughly
#   the same. Since the Michel electron distributions are consistent, we leave
#   the mean value at 0, but to be conservative, we set the error to 10%.
# - The energy resolution for muons was pretty much spot on, and so we expect
#   the same from electrons. In addition, the Michel spectrum is consistent so
#   at that energy level we don't see anything which leads us to expect a major
#   difference. To be conservative, and because I don't think it really affects
#   the analysis at all, I'll leave the uncertainty here at 10% anyways.
PRIORS = [
    1.0,   # Atmospheric Neutrino Scale
    0.0,   # Electron energy scale
    0.0,   # Electron energy resolution
    0.053, # Muon energy scale
    0.0,   # Muon energy resolution
    0.0,   # Muon scale
]

PRIOR_UNCERTAINTIES = [
    0.2  , # Atmospheric Neutrino Scale
    0.1,   # Electron energy scale
    0.1,   # Electron energy resolution
    0.01,  # Muon energy scale
    0.013, # Muon energy resolution
    10.0,  # Muon scale
]

# Lower bounds for the fit parameters
PRIORS_LOW = [
    EPSILON,
    -10,
    EPSILON,
    -10,
    EPSILON,
    0
]

# Upper bounds for the fit parameters
PRIORS_HIGH = [
    10,
    10,
    10,
    10,
    10,
    1e9
]

particle_id = {20: 'e', 22: r'\mu'}

def plot_hist2(hists, bins, color=None):
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)

        bincenters = (bins[1:] + bins[:-1])/2
        plt.hist(bincenters, bins=bins, histtype='step', weights=hists[id],color=color)
        plt.gca().set_xscale("log")
        major = np.array([10,100,1000,10000])
        minor = np.unique(list(chain(*list(range(i,i*10,i) for i in major[:-1]))))
        minor = np.setdiff1d(minor,major)
        plt.gca().set_xticks(major)
        plt.gca().set_xticks(minor,minor=True)
        plt.xlabel("Energy (MeV)")
        plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')

    if len(hists):
        plt.tight_layout()

def get_mc_hists(data,x,bins,scale=1.0,reweight=False):
    """
    Returns the expected Monte Carlo histograms for the atmospheric neutrino
    background.

    Args:
        - data: pandas dataframe of the Monte Carlo events
        - x: fit parameters
        - bins: histogram bins
        - scale: multiply histograms by an overall scale factor

    This function does two basic things:

        1. apply the energy bias and resolution corrections
        2. histogram the results

    Returns a dictionary mapping particle id combo -> histogram.
    """
    df_dict = {}
    for id in (20,22,2020,2022,2222):
        df_dict[id] = data[data.id == id]

    return get_mc_hists_fast(df_dict,x,bins,scale,reweight)

def get_mc_hists_fast(df_dict,x,bins,scale=1.0,reweight=False):
    """
    Same as get_mc_hists() but the first argument is a dictionary mapping
    particle id -> dataframe. This is much faster than selecting the events
    from the dataframe every time.
    """
    mc_hists = {}

    for id in (20,22,2020,2022,2222):
        df = df_dict[id]

        if id == 20:
            ke = df.energy1.values*(1+x[1])
            resolution = df.energy1.values*max(EPSILON,x[2])
        elif id == 2020:
            ke = df.energy1.values*(1+x[1]) + df.energy2.values*(1+x[1])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[2]))**2 + (df.energy2.values*max(EPSILON,x[2]))**2)
        elif id == 22:
            ke = df.energy1.values*(1+x[3])
            resolution = df.energy1.values*max(EPSILON,x[4])
        elif id == 2222:
            ke = df.energy1.values*(1+x[3]) + df.energy2.values*(1+x[3])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[4]))**2 + (df.energy2.values*max(EPSILON,x[4]))**2)
        elif id == 2022:
            ke = df.energy1.values*(1+x[1]) + df.energy2.values*(1+x[3])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[2]))**2 + (df.energy2.values*max(EPSILON,x[4]))**2)

        if reweight:
            cdf = fast_cdf(bins[:,np.newaxis],ke,resolution)*df.weight.values
        else:
            cdf = fast_cdf(bins[:,np.newaxis],ke,resolution)

        mc_hists[id] = np.sum(cdf[1:] - cdf[:-1],axis=-1)
        mc_hists[id] *= scale
    return mc_hists

def get_data_hists(data,bins,scale=1.0):
    """
    Returns the data histogrammed into `bins`.
    """
    data_hists = {}
    for id in (20,22,2020,2022,2222):
        data_hists[id] = np.histogram(data[data.id == id].ke.values,bins=bins)[0]*scale
    return data_hists

def make_nll(data, muons, mc, atmo_scale_factor, muon_scale_factor, bins, reweight=False, print_nll=False):
    df_dict = {}
    for id in (20,22,2020,2022,2222):
        df_dict[id] = mc[mc.id == id]

    df_dict_muon = {}
    for id in (20,22,2020,2022,2222):
        df_dict_muon[id] = muons[muons.id == id]

    data_hists = get_data_hists(data,bins)

    def nll(x, grad=None):
        if (x < PRIORS_LOW).any() or (x > PRIORS_HIGH).any():
            return np.inf

        # Get the Monte Carlo histograms. We need to do this within the
        # likelihood function since we apply the energy resolution parameters
        # to the Monte Carlo.
        mc_hists = get_mc_hists_fast(df_dict,x,bins,scale=1/atmo_scale_factor,reweight=reweight)
        muon_hists = get_mc_hists_fast(df_dict_muon,x,bins,scale=1/muon_scale_factor)

        # Calculate the negative log of the likelihood of observing the data
        # given the fit parameters

        nll = 0
        for id in data_hists:
            oi = data_hists[id]
            ei = mc_hists[id]*x[0] + muon_hists[id]*x[5] + EPSILON
            N = ei.sum()
            nll -= -N - np.sum(gammaln(oi+1)) + np.sum(oi*np.log(ei))

        # Add the priors
        nll -= norm.logpdf(x,PRIORS,PRIOR_UNCERTAINTIES).sum()

        if print_nll:
            # Print the result
            print("nll = %.2f" % nll)

        return nll
    return nll

def get_mc_hists_posterior(data_mc,muon_hists,data_hists,atmo_scale_factor,muon_scale_factor,x,bins):
    """
    Returns the posterior on the Monte Carlo histograms.

    Basically this function just histograms the Monte Carlo data. However,
    there is one extra thing it does. In general when doing a fit, the Monte
    Carlo histograms have some uncertainty since you can never simulate an
    infinite number of statistics. I don't think I've ever really seen anyone
    properly treat this. Since the uncertainty on the central value in each bin
    is just given by the Dirichlet distribution, we treat the problem of
    finding the best value of the posterior as a problem in which you're prior
    is equal to the expected number of events from the Monte Carlo, and then
    you actually see the data.  Since the likelihood on the true mean in each
    bin is a multinomial, the posterior is also a dirichlet where the alpha
    parameters are given by a sum of the prior and observed counts.

    All that is a long way of saying we calculate the posterior as the sum of
    the Monte Carlo events (unscaled) and the observed events. In the limit of
    infinite statistics, this is just equal to the Monte Carlo predicted
    histogram, but deals with the fact that we don't have infinite statistics,
    and so a single outlier event isn't necessarily a problem with the model.

    Returns a dictionary mapping particle id combo -> histogram.
    """
    mc_hists = get_mc_hists(data_mc,x,bins,reweight=True)
    for id in (20,22,2020,2022,2222):
        mc_hists[id] = get_mc_hist_posterior(mc_hists[id],data_hists[id],norm=x[0]/atmo_scale_factor)
        # FIXME: does the orering of when we add the muons matter here?
        mc_hists[id] += muon_hists[id]*x[5]/muon_scale_factor
    return mc_hists

def get_multinomial_prob(data, data_muon, data_mc, atmo_scale_factor, muon_scale_factor, id, x_samples, bins, percentile=50.0, size=10000):
    """
    Returns the p-value that the histogram of the data is drawn from the MC
    histogram.

    The p-value is calculated by first sampling the posterior of the fit
    parameters `size` times. For each iteration we calculate a p-value. We then
    return the `percentile` percentile of all the p-values. This approach is
    similar to both the supremum and posterior predictive methods of
    calculating a p-value. For more information on different methods of
    calculating p-values see https://cds.cern.ch/record/1099967/files/p23.pdf.

    Arguments:

        data: 1D array of KE values
        data_mc: 1D array of MC KE values
        x_samples: MCMC samples of the floated parameters in the fit
        bins: bins used to bin the mc histogram
        size: number of values to compute
    """
    df_dict_muon = {}
    for _id in (20,22,2020,2022,2222):
        df_dict_muon[_id] = data_muon[data_muon.id == _id]

    data_hists = get_data_hists(data,bins)

    ps = []
    for i in range(size):
        x = x_samples[np.random.randint(x_samples.shape[0])]

        muon_hists = get_mc_hists_fast(df_dict_muon,x,bins)

        mc = get_mc_hists_posterior(data_mc,muon_hists,data_hists,atmo_scale_factor,muon_scale_factor,x,bins)[id]
        N = mc.sum()
        # Fix a bug in scipy(). See https://github.com/scipy/scipy/issues/8235 (I think).
        mc = mc + 1e-10
        p = mc/mc.sum()
        chi2_data = nllr(data_hists[id],mc)
        # To draw the multinomial samples we first draw the expected number of
        # events from a Poisson distribution and then loop over the counts and
        # unique values. The reason we do this is that you can't call
        # multinomial.rvs with a multidimensional `n` array, and looping over every
        # single entry takes forever
        ns = np.random.poisson(N,size=1000)
        samples = []
        for n, count in zip(*np.unique(ns,return_counts=True)):
            samples.append(multinomial.rvs(n,p,size=count))
        samples = np.concatenate(samples)
        # Calculate the negative log likelihood ratio for the data simulated under
        # the null hypothesis
        chi2_samples = nllr(samples,mc)
        ps.append(np.count_nonzero(chi2_samples >= chi2_data)/len(chi2_samples))
    return np.percentile(ps,percentile)

def get_prob(data,muon,mc,atmo_scale_factor,muon_scale_factor,samples,bins,size):
    prob = {}
    for id in (20,22,2020,2022,2222):
        prob[id] = get_multinomial_prob(data,muon,mc,atmo_scale_factor,muon_scale_factor,id,samples,bins,size=size)
        print(id, prob[id])
    return prob

def do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,steps,print_nll=False,walkers=100,thin=10,refit=True):
    """
    Run the fit and return the minimum along with samples from running an MCMC
    starting near the minimum.

    Args:
        - data: pandas dataframe representing the data to fit
        - muon: pandas dataframe representing the expected background from
                external muons
        - data_mc: pandas dataframe representing the expected background from
                   atmospheric neutrino events
        - weights: pandas dataframe with the GENIE weights
        - bins: an array of bins to use for the fit
        - steps: the number of MCMC steps to run

    Returns a tuple (xopt, universe, samples) where samples is an array of
    shape (steps, number of parameters).
    """
    nll = make_nll(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,print_nll=print_nll)

    pos = np.empty((walkers, len(PRIORS)),dtype=np.double)
    for i in range(pos.shape[0]):
        pos[i] = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

    nwalkers, ndim = pos.shape

    # We use the KDEMove here because I think it should sample the likelihood
    # better. Because we have energy scale parameters and we are doing a binned
    # likelihood, the likelihood is discontinuous. There can also be several
    # local minima. The author of emcee recommends using the KDEMove with a lot
    # of workers to try and properly sample a multimodal distribution. In
    # addition, I've found that the autocorrelation time for the KDEMove is
    # much better than the other moves.
    sampler = emcee.EnsembleSampler(nwalkers, ndim, lambda x: -nll(x), moves=emcee.moves.KDEMove())
    with np.errstate(invalid='ignore'):
        sampler.run_mcmc(pos, steps)

    print("Mean acceptance fraction: {0:.3f}".format(np.mean(sampler.acceptance_fraction)))

    try:
        print("autocorrelation time: ", sampler.get_autocorr_time(quiet=True))
    except Exception as e:
        print(e)

    samples = sampler.get_chain(flat=True,thin=thin)

    # Now, we use nlopt to find the best set of parameters. We start at the
    # best starting point from the MCMC and then run the SBPLX routine.
    x0 = sampler.get_chain(flat=True)[sampler.get_log_prob(flat=True).argmax()]
    opt = nlopt.opt(nlopt.LN_SBPLX, len(x0))
    opt.set_min_objective(nll)
    low = np.array(PRIORS_LOW)
    high = np.array(PRIORS_HIGH)
    opt.set_lower_bounds(low)
    opt.set_upper_bounds(high)
    opt.set_ftol_abs(1e-10)
    opt.set_initial_step([0.01]*len(x0))
    xopt = opt.optimize(x0)

    # Get the total number of "universes" simulated in the GENIE reweight tool
    nuniverses = weights['universe'].max()+1

    nlls = []
    for universe in range(nuniverses):
        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
        data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)

        nll = make_nll(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll)
        nlls.append(nll(xopt))

    universe = np.argmin(nlls)

    if refit:
        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
        data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)

        # Create a new negative log likelihood function with the weighted Monte Carlo.
        nll = make_nll(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll)

        # Now, we refit with the Monte Carlo weighted by the most likely GENIE
        # systematics.
        pos = np.empty((walkers, len(PRIORS)),dtype=np.double)
        for i in range(pos.shape[0]):
            pos[i] = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

        nwalkers, ndim = pos.shape

        # We use the KDEMove here because I think it should sample the likelihood
        # better. Because we have energy scale parameters and we are doing a binned
        # likelihood, the likelihood is discontinuous. There can also be several
        # local minima. The author of emcee recommends using the KDEMove with a lot
        # of workers to try and properly sample a multimodal distribution. In
        # addition, I've found that the autocorrelation time for the KDEMove is
        # much better than the other moves.
        sampler = emcee.EnsembleSampler(nwalkers, ndim, lambda x: -nll(x), moves=emcee.moves.KDEMove())
        with np.errstate(invalid='ignore'):
            sampler.run_mcmc(pos, steps)

        print("Mean acceptance fraction: {0:.3f}".format(np.mean(sampler.acceptance_fraction)))

        try:
            print("autocorrelation time: ", sampler.get_autocorr_time(quiet=True))
        except Exception as e:
            print(e)

        samples = sampler.get_chain(flat=True,thin=thin)

        # Now, we use nlopt to find the best set of parameters. We start at the
        # best starting point from the MCMC and then run the SBPLX routine.
        x0 = sampler.get_chain(flat=True)[sampler.get_log_prob(flat=True).argmax()]
        opt = nlopt.opt(nlopt.LN_SBPLX, len(x0))
        opt.set_min_objective(nll)
        low = np.array(PRIORS_LOW)
        high = np.array(PRIORS_HIGH)
        opt.set_lower_bounds(low)
        opt.set_upper_bounds(high)
        opt.set_ftol_abs(1e-10)
        opt.set_initial_step([0.01]*len(x0))
        xopt = opt.optimize(x0)

    return xopt, universe, samples

if __name__ == '__main__':
    import argparse
    import numpy as np
    import pandas as pd
    import sys
    import h5py
    from sddm.plot_energy import *
    from sddm.plot import *
    from sddm import setup_matplotlib
    import nlopt

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("filenames", nargs='+', help="input files")
    parser.add_argument("--save", action='store_true', default=False, help="save corner plots for backgrounds")
    parser.add_argument("--mc", nargs='+', required=True, help="atmospheric MC files")
    parser.add_argument("--muon-mc", nargs='+', required=True, help="muon MC files")
    parser.add_argument("--nhit-thresh", type=int, default=None, help="nhit threshold to apply to events before processing (should only be used for testing to speed things up)")
    parser.add_argument("--steps", type=int, default=1000, help="number of steps in the MCMC chain")
    parser.add_argument("--multinomial-prob-size", type=int, default=10000, help="number of p values to compute")
    parser.add_argument("--coverage", type=int, default=0, help="plot p value coverage")
    parser.add_argument("--pull", type=int, default=0, help="plot pull plots")
    parser.add_argument("--weights", nargs='+', required=True, help="GENIE reweight HDF5 files")
    parser.add_argument("--print-nll", action='store_true', default=False, help="print nll values")
    parser.add_argument("--walkers", type=int, default=100, help="number of walkers")
    parser.add_argument("--thin", type=int, default=10, help="number of steps to thin")
    parser.add_argument("--run-list", default=None, help="run list")
    args = parser.parse_args()

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    # Loop over runs to prevent using too much memory
    evs = []
    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.filenames],ignore_index=True)
    for run, df in rhdr.groupby('run'):
        evs.append(get_events(df.filename.values, merge_fits=True, nhit_thresh=args.nhit_thresh))
    ev = pd.concat(evs)

    if args.run_list is not None:
        runs = np.genfromtxt(args.run_list)
        ev = ev[ev.run.isin(runs)]

    ev = correct_energy_bias(ev)

    # Note: We loop over the MC filenames here instead of just passing the
    # whole list to get_events() because I had to rerun some of the MC events
    # using SNOMAN and so most of the runs actually have two different files
    # and otherwise the GTIDs will clash
    ev_mcs = []
    for filename in args.mc:
        ev_mcs.append(get_events([filename], merge_fits=True, nhit_thresh=args.nhit_thresh, mc=True))
    ev_mc = pd.concat(ev_mcs)
    muon_mc = get_events(args.muon_mc, merge_fits=True, nhit_thresh=args.nhit_thresh, mc=True)
    weights = pd.concat([read_hdf(filename, "weights") for filename in args.weights],ignore_index=True)

    ev_mc = correct_energy_bias(ev_mc)
    muon_mc = correct_energy_bias(muon_mc)

    # Set all prompt events in the MC to be muons
    muon_mc.loc[muon_mc.prompt & muon_mc.filename.str.contains("cosmic"),'muon'] = True

    ev = ev.reset_index()
    ev_mc = ev_mc.reset_index()
    muon_mc = muon_mc.reset_index()

    # 00-orphan cut
    ev = ev[(ev.gtid & 0xff) != 0]
    ev_mc = ev_mc[(ev_mc.gtid & 0xff) != 0]
    muon_mc = muon_mc[(muon_mc.gtid & 0xff) != 0]

    # remove events 200 microseconds after a muon
    ev = ev.groupby('run',group_keys=False).apply(muon_follower_cut)

    # Get rid of events which don't have a successful fit
    ev = ev[~np.isnan(ev.fmin)]
    ev_mc = ev_mc[~np.isnan(ev_mc.fmin)]
    muon_mc = muon_mc[~np.isnan(muon_mc.fmin)]

    # require (r < av radius)
    ev = ev[ev.r < AV_RADIUS]
    ev_mc = ev_mc[ev_mc.r < AV_RADIUS]
    muon_mc = muon_mc[muon_mc.r < AV_RADIUS]

    # require psi < 6
    ev = ev[ev.psi < 6]
    ev_mc = ev_mc[ev_mc.psi < 6]
    muon_mc = muon_mc[muon_mc.psi < 6]

    data = ev[ev.signal & ev.prompt & ~ev.atm]
    data_atm = ev[ev.signal & ev.prompt & ev.atm]

    # Right now we use the muon Monte Carlo in the fit. If you want to use the
    # actual data, you can comment the next two lines and then uncomment the
    # two after that.
    muon = muon_mc[muon_mc.muon & muon_mc.prompt & ~muon_mc.atm]
    muon_atm = muon_mc[muon_mc.muon & muon_mc.prompt & muon_mc.atm]
    #muon = ev[ev.muon & ev.prompt & ~ev.atm]
    #muon_atm = ev[ev.muon & ev.prompt & ev.atm]

    if (~rhdr.run.isin(ev_mc.run)).any():
        print_warning("Error! The following runs have no Monte Carlo: %s" % \
            rhdr.run[~rhdr.run.isin(ev_mc.run)].values)
        sys.exit(1)

    ev_mc = ev_mc[ev_mc.run.isin(rhdr.run)]

    data_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ~ev_mc.atm]
    data_atm_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ev_mc.atm]

    data_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ~ev_mc.atm]

    bins = np.logspace(np.log10(20),np.log10(10e3),21)

    atmo_scale_factor = 100.0
    muon_scale_factor = len(muon) + len(muon_atm)

    if args.coverage:
        p_values = {id: [] for id in (20,22,2020,2022,2222)}
        p_values_atm = {id: [] for id in (20,22,2020,2022,2222)}

        # Set the random seed so we get reproducible results here
        np.random.seed(0)

        xtrue = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == 0],how='left',on=['run','evn'])
        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == 0],how='left',on=['run','evn'])

        for i in range(args.coverage):
            # Calculate expected number of events
            N = len(data_mc)*xtrue[0]/atmo_scale_factor
            N_atm = len(data_atm_mc)*xtrue[0]/atmo_scale_factor
            N_muon = len(muon)*xtrue[5]/muon_scale_factor
            N_muon_atm = len(muon_atm)*xtrue[5]/muon_scale_factor

            # Calculate observed number of events
            n = np.random.poisson(N)
            n_atm = np.random.poisson(N_atm)
            n_muon = np.random.poisson(N_muon)
            n_muon_atm = np.random.poisson(N_muon_atm)

            # Sample data from Monte Carlo
            data = pd.concat((data_mc_with_weights.sample(n=n,replace=True,weights='weight'), muon.sample(n=n_muon,replace=True)))
            data_atm = pd.concat((data_atm_mc_with_weights.sample(n=n_atm,replace=True,weights='weight'), muon_atm.sample(n=n_muon_atm,replace=True)))

            # Smear the energies by the additional energy resolution
            data.loc[data.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id1 == 20))*xtrue[2])
            data.loc[data.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id1 == 22))*xtrue[4])
            data.loc[data.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id2 == 20))*xtrue[2])
            data.loc[data.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id2 == 22))*xtrue[4])
            data['ke'] = data['energy1'].fillna(0) + data['energy2'].fillna(0) + data['energy3'].fillna(0)

            data_atm.loc[data_atm.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id1 == 20))*xtrue[2])
            data_atm.loc[data_atm.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id1 == 22))*xtrue[4])
            data_atm.loc[data_atm.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id2 == 20))*xtrue[2])
            data_atm.loc[data_atm.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id2 == 22))*xtrue[4])
            data_atm['ke'] = data_atm['energy1'].fillna(0) + data_atm['energy2'].fillna(0) + data_atm['energy3'].fillna(0)

            xopt, universe, samples = do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)

            data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
            data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)

            data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
            data_atm_mc_with_weights.weight = data_atm_mc_with_weights.weight.fillna(1.0)

            prob = get_prob(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
            prob_atm = get_prob(data_atm,muon_atm,data_atm_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)

            for id in (20,22,2020,2022,2222):
                p_values[id].append(prob[id])
                p_values_atm[id].append(prob_atm[id])

        fig = plt.figure()
        for id in (20,22,2020,2022,2222):
            if id == 20:
                plt.subplot(2,3,1)
            elif id == 22:
                plt.subplot(2,3,2)
            elif id == 2020:
                plt.subplot(2,3,4)
            elif id == 2022:
                plt.subplot(2,3,5)
            elif id == 2222:
                plt.subplot(2,3,6)
            plt.hist(p_values[id],bins=np.linspace(0,1,11),histtype='step')
            plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')
        despine(fig,trim=True)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_p_value_coverage_plot.pdf")
            fig.savefig("chi2_p_value_coverage_plot.eps")
        else:
            plt.suptitle("P-value Coverage without Neutron Follower")

        fig = plt.figure()
        for id in (20,22,2020,2022,2222):
            if id == 20:
                plt.subplot(2,3,1)
            elif id == 22:
                plt.subplot(2,3,2)
            elif id == 2020:
                plt.subplot(2,3,4)
            elif id == 2022:
                plt.subplot(2,3,5)
            elif id == 2222:
                plt.subplot(2,3,6)
            plt.hist(p_values_atm[id],bins=np.linspace(0,1,11),histtype='step')
            plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')
        despine(fig,trim=True)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_p_value_coverage_plot_atm.pdf")
            fig.savefig("chi2_p_value_coverage_plot_atm.eps")
        else:
            plt.suptitle("P-value Coverage with Neutron Follower")

        sys.exit(0)

    if args.pull:
        pull = [[] for i in range(len(FIT_PARS))]

        # Set the random seed so we get reproducible results here
        np.random.seed(0)

        for i in range(args.pull):
            xtrue = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

            # Calculate expected number of events
            N = len(data_mc)*xtrue[0]/atmo_scale_factor
            N_atm = len(data_atm_mc)*xtrue[0]/atmo_scale_factor
            N_muon = len(muon)*xtrue[5]/muon_scale_factor
            N_muon_atm = len(muon_atm)*xtrue[5]/muon_scale_factor

            # Calculate observed number of events
            n = np.random.poisson(N)
            n_atm = np.random.poisson(N_atm)
            n_muon = np.random.poisson(N_muon)
            n_muon_atm = np.random.poisson(N_muon_atm)

            # Sample data from Monte Carlo
            data = pd.concat((data_mc.sample(n=n,replace=True), muon.sample(n=n_muon,replace=True)))
            data_atm = pd.concat((data_atm_mc.sample(n=n_atm,replace=True), muon_atm.sample(n=n_muon_atm,replace=True)))

            # Smear the energies by the additional energy resolution
            data.loc[data.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id1 == 20))*xtrue[2])
            data.loc[data.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id1 == 22))*xtrue[4])
            data.loc[data.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id2 == 20))*xtrue[2])
            data.loc[data.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id2 == 22))*xtrue[4])
            data['ke'] = data['energy1'].fillna(0) + data['energy2'].fillna(0) + data['energy3'].fillna(0)

            data_atm.loc[data_atm.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id1 == 20))*xtrue[2])
            data_atm.loc[data_atm.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id1 == 22))*xtrue[4])
            data_atm.loc[data_atm.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id2 == 20))*xtrue[2])
            data_atm.loc[data_atm.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id2 == 22))*xtrue[4])
            data_atm['ke'] = data_atm['energy1'].fillna(0) + data_atm['energy2'].fillna(0) + data_atm['energy3'].fillna(0)

            xopt, universe, samples = do_fit(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin,refit=False)

            for i in range(len(FIT_PARS)):
                # The "pull plots" we make here are actually produced via a
                # procedure called "Simulation Based Calibration".
                #
                # See https://arxiv.org/abs/1804.06788.
                pull[i].append(percentileofscore(samples[:,i],xtrue[i]))

        fig = plt.figure()
        axes = []
        for i, name in enumerate(FIT_PARS):
            axes.append(plt.subplot(3,2,i+1))
            n, bins, patches = plt.hist(pull[i],bins=np.linspace(0,100,11),histtype='step')
            expected = len(pull[i])/(len(bins)-1)
            plt.axhline(expected,color='k',ls='--',alpha=0.25)
            plt.axhspan(poisson.ppf(0.005,expected), poisson.ppf(0.995,expected), facecolor='0.5', alpha=0.25)
            plt.title(name)
        for ax in axes:
            despine(ax=ax,left=True,trim=True)
            ax.get_yaxis().set_visible(False)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_pull_plot.pdf")
            fig.savefig("chi2_pull_plot.eps")
        else:
            plt.show()

        sys.exit(0)

    xopt, universe, samples = do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)

    data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
    data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)

    data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
    data_atm_mc_with_weights.weight = data_atm_mc_with_weights.weight.fillna(1.0)

    prob = get_prob(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
    prob_atm = get_prob(data_atm,muon_atm,data_atm_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)

    plt.figure()
    for i in range(len(FIT_PARS)):
        plt.subplot(3,2,i+1)
        plt.hist(samples[:,i],bins=100,histtype='step')
        plt.xlabel(FIT_PARS[i].title())
        despine(ax=plt.gca(),left=True,trim=True)
        plt.gca().get_yaxis().set_visible(False)
    plt.tight_layout()

    if args.save:
        plt.savefig("chi2_fit_posterior.pdf")
        plt.savefig("chi2_fit_posterior.eps")
    else:
        plt.suptitle("Fit Posteriors")

    handles = [Line2D([0], [0], color='C0'),
               Line2D([0], [0], color='C1'),
               Line2D([0], [0], color='C2')]
    labels = ('Data','Monte Carlo','External Muons')

    fig = plt.figure()
    hists = get_data_hists(data,bins)
    hists_muon = get_mc_hists(muon,xopt,bins,scale=xopt[5]/muon_scale_factor)
    hists_mc = get_mc_hists(data_mc_with_weights,xopt,bins,scale=xopt[0]/atmo_scale_factor,reweight=True)
    plot_hist2(hists,bins=bins,color='C0')
    plot_hist2(hists_mc,bins=bins,color='C1')
    plot_hist2(hists_muon,bins=bins,color='C2')
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)
        plt.text(0.95,0.95,"p = %.2f" % prob[id],horizontalalignment='right',verticalalignment='top',transform=plt.gca().transAxes)
    fig.legend(handles,labels,loc='upper right')

    despine(fig,trim=True)
    if args.save:
        plt.savefig("chi2_prompt.pdf")
        plt.savefig("chi2_prompt.eps")
    else:
        plt.suptitle("Without Neutron Follower")
    fig = plt.figure()
    hists = get_data_hists(data_atm,bins)
    hists_muon = get_mc_hists(muon_atm,xopt,bins,scale=xopt[5]/muon_scale_factor)
    hists_mc = get_mc_hists(data_atm_mc_with_weights,xopt,bins,scale=xopt[0]/atmo_scale_factor,reweight=True)
    plot_hist2(hists,bins=bins,color='C0')
    plot_hist2(hists_mc,bins=bins,color='C1')
    plot_hist2(hists_muon,bins=bins,color='C2')
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)
        plt.text(0.95,0.95,"p = %.2f" % prob_atm[id],horizontalalignment='right',verticalalignment='top',transform=plt.gca().transAxes)
    fig.legend(handles,labels,loc='upper right')

    despine(fig,trim=True)
    if args.save:
        plt.savefig("chi2_atm.pdf")
        plt.savefig("chi2_atm.eps")
    else:
        plt.suptitle("With Neutron Follower")
        plt.show()