#!/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[id][1:] + bins[id][:-1])/2
plt.hist(bincenters, bins=bins[id], 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)
major = major[major <= bins[id][-1]]
minor = minor[minor <= bins[id][-1]]
plt.gca().set_xticks(major)
plt.gca().set_xticks(minor,minor=True)
plt.gca().set_xlim(10,10000)
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[id][:,np.newaxis],ke,resolution)*df.weight.values
else:
cdf = fast_cdf(bins[id][:,np.newaxis],ke,resolution)
if 'flux_weight' in df.columns:
cdf *= df.flux_weight.values
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[id])[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 = dict(tuple(mc.groupby('id')))
for id in (20,22,2020,2022,2222):
if id not in df_dict:
df_dict[id] = mc.iloc[:0]
df_dict_muon = dict(tuple(muons.groupby('id')))
for id in (20,22,2020,2022,2222):
if id not in df_dict_muon:
df_dict_muon[id] = muons.iloc[:0]
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
weights_dict = dict(tuple(weights.groupby('universe')))
nlls = []
for universe in range(nuniverses):
data_mc_with_weights = pd.merge(data_mc,weights_dict[universe],how='left',on=['run','unique_id'])
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','unique_id'])
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
from sddm.renormalize import *
parser = argparse.ArgumentParser("run null hypothesis test")
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")
parser.add_argument("--mcpl", nargs='+', required=True, help="GENIE MCPL files")
args = parser.parse_args()
setup_matplotlib(args.save)
import matplotlib.pyplot as plt
if args.run_list is not None:
run_list = np.genfromtxt(args.run_list)
# 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'):
if args.run_list and run not in run_list:
continue
evs.append(get_events(df.filename.values, merge_fits=True, nhit_thresh=args.nhit_thresh))
ev = pd.concat(evs).reset_index()
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_mc for ev_mc in ev_mcs if len(ev_mc) > 0]).reset_index()
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)
# The next two things we have to do are reweight the Monte Carlo data since
# I accidentally simulated the muon neutrino flux instead of the tau
# neutrino flux and load in the GENIE systematics weights.
#
# Both of these are a bit tricky because of the fact that I had to
# resimulate some MC events since they failed to simulate (there was a
# packer error and occasionally a geometry error that was causing ~10% of
# the MC events to fail). Since I had to resimulate them, it's not possible
# to connect the GENIE weights to the MC events by just using the event
# number.
#
# Therefore, I decided to use a "unique_id" field which I compute by
# hashing the position of the event. This unique_id along with the run
# should completely specify a unique mapping between the events.
# Add the "flux_weight" column to the ev_mc data since I stupidly simulated
# the muon neutrino flux for the tau neutrino flux in GENIE. Doh!
mcpl = load_mcpl_files(args.mcpl)
ev_mc = renormalize_data(ev_mc.reset_index(),mcpl)
# Merge weights with MCPL dataframe to get the unique id column in the
# weights dataframe since that is what we use to merge with the Monte
# Carlo.
weights = pd.merge(weights,mcpl[['run','evn','unique_id']],on=['run','evn'],how='left')
# There are a handful of weights which turn out to be slightly negative for
# some reason. For example:
#
# run evn universe weight
# 10970 25 597 -0.000055
# 11389 87 729 -0.021397
# 11701 204 2 -0.000268
# 11919 120 82 -0.002245
# 11976 163 48 -0.000306
# 11976 163 710 -0.000022
# 12131 76 175 -0.000513
# 12207 70 255 -0.002925
# 12207 70 282 -0.014856
# 12207 70 368 -0.030593
# 12207 70 453 -0.019011
# 12207 70 520 -0.020748
# 12207 70 834 -0.028754
# 12207 70 942 -0.020309
# 12233 230 567 -0.000143
# 12618 168 235 -0.000020
# 13428 128 42 -0.083639
# 14264 23 995 -0.017637
# 15034 69 624 -0.000143
# 15752 154 957 -0.006827
weights = weights[weights.weight > 0]
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 (~ev.run.isin(ev_mc.run)).any():
print_warning("Error! The following runs have no Monte Carlo: %s" % \
np.unique(ev.run[~ev.run.isin(ev_mc.run)].values))
sys.exit(1)
if not args.pull and not args.coverage:
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]
bins = {20:np.logspace(np.log10(20),np.log10(10e3),21),
22:np.logspace(np.log10(20),np.log10(10e3),21)[:-5],
2020:np.logspace(np.log10(20),np.log10(10e3),21),
2022:np.logspace(np.log10(20),np.log10(10e3),21)[:-5],
2222:np.logspace(np.log10(20),np.log10(10e3),21)[:-5]}
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','unique_id'])
data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == 0],how='left',on=['run','unique_id'])
data_mc_with_weights.weight *= data_mc_with_weights.flux_weight
data_atm_mc_with_weights.weight *= data_atm_mc_with_weights.flux_weight
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','unique_id'])
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','unique_id'])
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,weights='flux_weight',replace=True), muon.sample(n=n_muon,replace=True)))
data_atm = pd.concat((data_atm_mc.sample(n=n_atm,weights='flux_weight',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,weights,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','unique_id'])
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','unique_id'])
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()
