<jupyter_start><jupyter_text>OTUS Paper Plots | $p p > Z > e^+ e^-$ This notebook creates the plots which are found in the paper: https://arxiv.org/abs/2101.08944. It also performs the statistical tests which are included in the supplementary material.  Load Required Libraries<jupyter_code>%matplotlib inline
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import matplotlib.cm as cm
from scipy import stats

# Add utilityFunctions/ to easily use utility .py files
import sys
sys.path.append("../utilityFunctions")<jupyter_output><empty_output><jupyter_text>Meta Parameters<jupyter_code>save_figs = True  # Global switch to turn on saving figures to files
N_slices  = 1000  # Global number of slices to calculate SW distance with

data_directory    = '../data/'
results_dataPath  = '../experiments/ppzee/otus_results-dataset=ppzee_test.npz'
dataset_name      = 'ppzee'
img_dir           = dataset_name+'Images/'

np.random.seed(4) # Set numpy random seed<jupyter_output><empty_output><jupyter_text>Get Truth Datasets and Model Results The data was split into training, validation, and two testing sets. We load only the testing datasets here. We also load in the model results calculated on one of the testing datasets.  Note that we have two separate testing sets to ensure the independence of samples when performing statistical tests.Specifically, we have 4 initial datasets: x_test1, z_test1, x_test2, and z_test2.Test set 1 has been used to obtain network results. These results will be compared to the independent test set 2.For testing the encoder's performance we will use the following convention:- z_pred_te1 = E(x_test1, $\epsilon$)In plots this will be denoted as $x \rightarrow \tilde{z}$.For testing the decoder's performance we will use the following convention:- x_pred_truth_te1 = D(z_test1, $\epsilon$)In plots this will be denoted as $z \rightarrow \tilde{x}'$.For testing the overall performance of the full encoder-decoder chain we will use the following convention:- x_pred_te1 = D(E(x_test1, $\epsilon$), $\epsilon$)In plots this will be denoted as $x \rightarrow \tilde{z} \rightarrow \tilde{x}$.Note that $\epsilon$ denotes the random noise input of the networks. Note that all $\epsilon$ are different random samples. This happens internally and thus won't be shown in what follows.  Load in Truth Data (Testing Dataset)<jupyter_code>from func_utils import get_dataset, standardize

#-- Load Testing Truth Data --#
truthDataset = get_dataset(dataset_name+'_test', data_dir=data_directory)
zData, xData = truthDataset['z_data'], truthDataset['x_data']

#-- Remove MET from xData --#
MET = False
if MET == False:
    xData = xData[:, :-4]
    
#-- Split data into the two testing sets --#
f_test1 = int(0.5*zData.shape[0])

x_test1  = xData[0:f_test1]
z_test1  = zData[0:f_test1]

x_test2  = xData[f_test1:]
z_test2  = zData[f_test1:]

#-- Make sure testing sets have the same number of events --#
assert(x_test1.shape[0] == x_test2.shape[0])
assert(z_test1.shape[0] == z_test2.shape[0])
print('z_test1 Shape, x_test1 Shape: ', z_test1.shape, x_test1.shape)
print('z_test2 Shape, x_test2 Shape: ', z_test2.shape, x_test2.shape)<jupyter_output>z_test1 Shape, x_test1 Shape:  (80000, 8) (80000, 8)
z_test2 Shape, x_test2 Shape:  (80000, 8) (80000, 8)<jupyter_text>Get Results (Calculated on Testing Dataset)<jupyter_code>#-- Load Testing Results Data--#
resultsDataset = np.load(results_dataPath)

zPred      = resultsDataset['x_encoded']        # E(x)
xPred      = resultsDataset['x_reconstructed']  # D(E(x))
xPredTruth = resultsDataset['z_decoded']        # D(z)

#-- Get predicted z for x_test1, E(x_test1) --#
z_pred_te1 = zPred[0:f_test1]

#-- Get predicted x for full autoencoder chain, D(E(x_test1)) --#
x_pred_te1 = xPred[0:f_test1]

#-- Get predicted x for true z_test1, D(z_test1) --#
x_pred_truth_te1 = xPredTruth[0:f_test1]<jupyter_output><empty_output><jupyter_text>Define Statistical Tests That We Will Use For each one-dimensional distribution we will perform the following statistical tests:- $\chi^2$- 2-sample, 2-sided Kolmogorov-Smirov testWe will also report the 1D Wasserstein distance.z-space plots will compare- z_test2 with z_pred_te1 In plots this will be denoted as $z$ vs $\tilde{z}$. x-space plots will compare- x_test2 with x_pred_te1 In plots this will be denoted as $x$ vs $\tilde{x}$. - x_test2 with x_pred_truth_te1 In plots this will be denoted as $x$ vs $\tilde{x}'$.  Import statistical test functions<jupyter_code>from func_utils import runStatAnalysis
from func_utils import sliced_wd_numpy as avSWDistance
# ReducedChiSq(), W1Dist() #?, KSTest(), runStatAnalysis(), generateTheta() #?, avSWDistance() #?<jupyter_output><empty_output><jupyter_text>Calculate SW Distance Results (Global Metric) Over N_Slices  z-space results<jupyter_code>print(r"SW distance $z$ vs $\tilde{z}$: ", avSWDistance(z_test2, z_pred_te1, N_slices))<jupyter_output>SW distance $z$ vs $\tilde{z}$:  0.9842470146462973<jupyter_text>x-space results<jupyter_code>print(r"SW distance $x$ vs $\tilde{x}$: ", avSWDistance(x_test2, x_pred_te1, N_slices))
print(r"SW distance $x$ vs $\tilde{x}^\prime$: ", avSWDistance(x_test2, x_pred_truth_te1, N_slices))<jupyter_output>SW distance $x$ vs $\tilde{x}$:  1.3339858997722187
SW distance $x$ vs $\tilde{x}^\prime$:  3.0273768282707865<jupyter_text>Run Statistical Analysis for Principal Axes We run the statistical analysis on the principal axis plots that are shown in the paper. These plots are for the positron's ($e^+$) $p_Y$, $p_Z$, and $E$ in both z-space and x-space.  We first find the correct binning such that no bin has less than 5 entries (necessary for $\chi^2$ test). We store this binning in a list to use in later plotting.  z-space Find Correct Binning<jupyter_code># Want to find binsY such that n has no bins with <5 entries
binsY = [-100, -60] + [-50+5*i for i in range(20+1)] + [60, 100]
n, _, _ = plt.hist(z_pred_te1[:,5], bins=binsY, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
# Want to find binsZ such that n has no bins with <5 entries
binsZ = [-400] + [-250+20*i for i in range(25+1)] + [400]
n, _, _ = plt.hist(z_pred_te1[:,6], bins=binsZ, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
# Want to find binsE such that n has no bins with <5 entries
binsE = [0] + [20+10*i for i in range(25+1)] + [400]
n, _, _ = plt.hist(z_pred_te1[:,7], bins=binsE, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot<jupyter_output>Number of entries in each bin:  [  312.  2945.  9314. 17720. 11873.  7709.  5697.  4377.  3387.  2713.
  2190.  1886.  1558.  1369.  1209.  1050.   879.   742.   667.   601.
   506.   365.   304.   213.   136.   107.   165.]
Do any bins have less than 5 entries?  False<jupyter_text>x-space Find Correct Binning<jupyter_code># Check if above z-space bins will work for x-space too
n, _, _ = plt.hist(x_pred_te1[:,5], bins=binsY, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
n, _, _ = plt.hist(x_pred_te1[:,6], bins=binsZ, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
n, _, _ = plt.hist(x_pred_te1[:,7], bins=binsE, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot<jupyter_output>Number of entries in each bin:  [  186.  3253. 10166. 16879. 11707.  7594.  5799.  4472.  3498.  2660.
  2171.  1864.  1620.  1368.  1187.  1008.   892.   738.   644.   587.
   458.   351.   273.   219.   153.    99.   147.]
Do any bins have less than 5 entries?  False<jupyter_text>Run statistical analysis for z-space principal axis data<jupyter_code>#-- Define suitable binning for principal axis data (obtained from above) --#
bins_py = binsY 
bins_pz = binsZ        
bins_E  = binsE            

zbinsList = [bins_py, bins_pz, bins_E]
#-- z-space principal axis data --#
zStatsList = []

for j in range(3):
    
    k = (j+1) + 4
              
    zStatsList.append(runStatAnalysis(z_pred_te1[:,k], z_test2[:,k], zbinsList[j]))

#-- Demonstration of how to access statistical test information --#
print(len(zStatsList))
print('W1D, KSstat, KSpval, RC2, dof, chi2pval')
print(zStatsList[0])
print("")
print('W1D: %.2E'% zStatsList[0][0])
print('(KSstat, pval): (%.2E, %.3f)' % (zStatsList[0][1], zStatsList[0][2]))
print('(RC2, dof, pval): (%.3f, %d, %.3f)' % (zStatsList[0][3], zStatsList[0][4], zStatsList[0][5]))<jupyter_output>3
W1D, KSstat, KSpval, RC2, dof, chi2pval
(1.3387174247094598, 0.016137500000000027, 1.791226386726743e-09, 50.58309020650058, 23, 0.0)

W1D: 1.34E+00
(KSstat, pval): (1.61E-02, 0.000)
(RC2, dof, pval): (50.583, 23, 0.000)<jupyter_text>Run statistical analysis for x-space principal axis  data<jupyter_code>xbinsList = zbinsList 
#-- x-space principal axis data --#
xStatsList = []
xStatsListTruth = []

for j in range(3):
    
    k = (j+1) + 4
              
    xStatsList.append(runStatAnalysis(x_pred_te1[:,k], x_test2[:,k], xbinsList[j]))
        
    xStatsListTruth.append(runStatAnalysis(x_pred_truth_te1[:,k], x_test2[:,k], xbinsList[j]))

#-- Demonstration of how to access statistical test information --#
print(len(xStatsList))
print('W1D, KSstat, KSpval, RC2, dof, chi2pval')
print(xStatsList[0])
print("")
print('W1D: %.2E'% xStatsList[0][0])
print('(KSstat, pval): (%.2E, %.3f)' % (xStatsList[0][1], xStatsList[0][2]))
print('(RC2, dof, pval): (%.3f, %d, %.3f)' % (xStatsList[0][3], xStatsList[0][4], xStatsList[0][5]))<jupyter_output>3
W1D, KSstat, KSpval, RC2, dof, chi2pval
(0.42236631680525877, 0.003475000000000006, 0.7195419941463636, 1.3905810968076997, 23, 0.10048984629634139)

W1D: 4.22E-01
(KSstat, pval): (3.48E-03, 0.720)
(RC2, dof, pval): (1.391, 23, 0.100)<jupyter_text>Make Principal Axes Plots  Load Plotting Functions<jupyter_code>from plot_utils import ratioPlotSingle, ratioPlotDouble, transportPlot<jupyter_output><empty_output><jupyter_text>Create Principal Axes Plots z-space<jupyter_code>#-- Plot this for the positron's py --#
ratioPlotSingle(z_test2[:,5], z_pred_te1[:,5], zbinsList[0], (-100,100), (0,7000), (0., 6.05), 
                   r'Positron $p_y$ [GeV/c]', legend=True, save_figs=save_figs, img_dir=img_dir) #zStatsList[0],
#-- Plot this for the positron's pz --#
ratioPlotSingle(z_test2[:,6], z_pred_te1[:,6], zbinsList[1], (-400,400), (0,7000), (.6,1.45), 
                   r'Positron $p_z$ [GeV/c]', legend=True, save_figs=save_figs, img_dir=img_dir) #zStatsList[1], 

#-- Plot this for the positron's E --#
ratioPlotSingle(z_test2[:,7], z_pred_te1[:,7], zbinsList[2], (0,400), (0,7000), (.6,1.65), 
                   r'Positron $E$ [GeV]', legend=True, save_figs=save_figs, img_dir=img_dir) #zStatsList[2],<jupyter_output><empty_output><jupyter_text>Create Principal Axes Plots x-space<jupyter_code>#-- Plot this for the positron's py --#
ratioPlotDouble(x_test2[:,5], x_pred_te1[:,5], x_pred_truth_te1[:,5], xbinsList[0], 
                   (-100,100), (0,7000), (.25,1.95), 
                   r'Positron $p_y$ [GeV/c]', legend=True, save_figs=save_figs, img_dir=img_dir)
#-- Plot this for the positron's pz --#
ratioPlotDouble(x_test2[:,6], x_pred_te1[:,6], x_pred_truth_te1[:,6], xbinsList[1], 
                   (-400,400), (0,7000), (0.25,1.45), 
                   r'Positron $p_z$ [GeV/c]', legend=True, save_figs=save_figs, img_dir=img_dir) #xStatsList[1], xStatsListTruth[1],
#-- Plot this for the positron's E --#
ratioPlotDouble(x_test2[:,7], x_pred_te1[:,7], x_pred_truth_te1[:,7], xbinsList[2], 
                   (0,400), (0,7000), (.5,1.35), 
                   r'Positron $E$ [GeV]', legend=True, save_figs=save_figs, img_dir=img_dir) #xStatsList[2], xStatsListTruth[2],<jupyter_output><empty_output><jupyter_text>Calculate Invariant Mass of Z Boson<jupyter_code>from func_utils import calcInvM
#-- Calculate Invariant Mass of Z boson --#

# z-space (z_pred_te1 vs z_test2)
zpred_IM_Z = calcInvM(z_pred_te1[:,0:4] + z_pred_te1[:,4:8])
ztest_IM_Z = calcInvM(z_test2[:,0:4] + z_test2[:,4:8])

# x-space
xpred_IM_Z      = calcInvM(x_pred_te1[:,0:4] + x_pred_te1[:,4:8])
xtest_IM_Z      = calcInvM(x_test2[:,0:4] + x_test2[:,4:8])
xpredtruth_IM_Z = calcInvM(x_pred_truth_te1[:,0:4] + x_pred_truth_te1[:,4:8])<jupyter_output><empty_output><jupyter_text>Perform Statistical Analysis  z-space Find Correct Binning and Compute Statistics<jupyter_code>bins = [70 + 1*i for i in range(40+1)]

n, _, _ = plt.hist(zpred_IM_Z, bins=bins, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
zIMstatsList = []

bins_Z = bins

zIMstatsList.append(runStatAnalysis(zpred_IM_Z, ztest_IM_Z, bins_Z))<jupyter_output><empty_output><jupyter_text>x-space Find Correct Binning and Compute Statistics<jupyter_code># Want to find bins such that n has no bins with <5 entries
# Try using binning for z-space case
bins = bins_Z
n, _, _ = plt.hist(xpred_IM_Z, bins=bins, histtype='step', density=False)
print("Number of entries in each bin: ",n)
print("Do any bins have less than 5 entries? ",np.any(n < 5))

plt.close() # Don't show plot
xIMstatsList = []
xIMstatsListTruth = []

xIMstatsList.append(runStatAnalysis(xpred_IM_Z, xtest_IM_Z, bins_Z))
xIMstatsListTruth.append(runStatAnalysis(xpredtruth_IM_Z, xtest_IM_Z, bins_Z))<jupyter_output><empty_output><jupyter_text>Plot z-space Invariant mass<jupyter_code>ratioPlotSingle(ztest_IM_Z, zpred_IM_Z, bins_Z, (70,110), (0,60000), (0,5.75), 
                   r'$M_{Z}$ [GeV/$c^2$]', legend=True, save_figs=save_figs, img_dir=img_dir) #zIMstatsList[0],<jupyter_output><empty_output><jupyter_text>Plot x-space Invariant mass<jupyter_code>ratioPlotDouble(xtest_IM_Z, xpred_IM_Z, xpredtruth_IM_Z, bins_Z, 
                   (70,110), (0,12000), (0.25,2.25),  
                   r'$M_{Z}$ [GeV/$c^2$]', legend=True, save_figs=save_figs, img_dir=img_dir) #xIMstatsList[0],  xIMstatsListTruth[0],<jupyter_output><empty_output><jupyter_text>Create Transport Plots  Decoder $\tilde{z} \rightarrow \tilde{x}$<jupyter_code>#-- Make decoder transport plot from learned latent space--#

# Positron energy
Z = z_pred_te1[:,7].reshape(-1,1)
X = x_pred_te1[:,7].reshape(-1,1)

n = Z.shape[1]
ZX = [np.concatenate((Z[:,i].reshape(-1,1), X[:,i].reshape(-1,1)), axis=1) for i in range(n)]

transportPlot(ZX[0], zlim=(0,400), xlim=(0,400), nz=20, nx=20, name='decTransport_zpred_xpred_te.png', dec=True, title=r'$\hspace{4}$ OTUS encoder-decoder: $\tilde{z} \rightarrow \tilde{x}$', save_figs=save_figs, img_dir=img_dir)<jupyter_output><empty_output><jupyter_text>Delphes $z \rightarrow x$<jupyter_code>#-- Make delphes transport plot --#

# Electron energy
Z = z_test1[:,7].reshape(-1,1)
X = x_test1[:,7].reshape(-1,1)

n = Z.shape[1]
ZX = [np.concatenate((Z[:,i].reshape(-1,1), X[:,i].reshape(-1,1)), axis=1) for i in range(n)]

transportPlot(ZX[0], zlim=(0,400), xlim=(0,400), nz=20, nx=20, name='delphesTransport_ztest_xtest.png', dec=True, title=r'Ground truth: $z \rightarrow x$', save_figs=save_figs, img_dir=img_dir)<jupyter_output><empty_output><jupyter_text>Decoder $z \rightarrow \tilde{x}^\prime$<jupyter_code>#-- Make decoder transport plot from truth --#

# Electron energy
Z = z_test1[:,7].reshape(-1,1)
X = x_pred_truth_te1[:,7].reshape(-1,1)

n = Z.shape[1]
ZX = [np.concatenate((Z[:,i].reshape(-1,1), X[:,i].reshape(-1,1)), axis=1) for i in range(n)]

transportPlot(ZX[0], zlim=(0,400), xlim=(0,400), nz=20, nx=20, name='decTransport_ztruth_xpredtruth_te.png', dec=True, title=r'OTUS decoder: $z \rightarrow \tilde{x}^\prime$', save_figs=save_figs, img_dir=img_dir)<jupyter_output><empty_output><jupyter_text>Statistics Tables We document the results of the above statistical tests in LaTeX table friendly format.<jupyter_code>from func_utils import statTableSingle, statTableDouble<jupyter_output><empty_output><jupyter_text>Supplementary Table 1 (z-space)<jupyter_code># Fig. 2a: fracDiff_zpred_z_Positronp_y.png
statTableSingle(statList12=zStatsList[0], figLabel='2a', cline=True)

# Fig. 2b: fracDiff_zpred_z_Positronp_z.png
statTableSingle(statList12=zStatsList[1], figLabel='2b', cline=True)

# Fig. 2c: fracDiff_zpred_z_PositronE.png
statTableSingle(statList12=zStatsList[2], figLabel='2c', cline=False)

# Fig. 4a: fracDiff_zpred_z_M_Z.png
statTableSingle(statList12=zIMstatsList[0], figLabel='4a', cline=False)<jupyter_output>\footnotesize{\textbf{Fig. 2a}} & \footnotesize{$1.34 \times 10^{+00}$} & \footnotesize{($50.583$, $23$)} & \footnotesize{$1.61 \times 10^{-02}$}\\\cline{2-4}
\footnotesize{\textbf{Fig. 2b}} & \footnotesize{$1.59 \times 10^{+00}$} & \footnotesize{($1.325$, $26$)} & \footnotesize{$4.90 \times 10^{-03}$}\\\cline{2-4}
\footnotesize{\textbf{Fig. 2c}} & \footnotesize{$1.29 \times 10^{+00}$} & \footnotesize{($8.814$, $26$)} & \footnotesize{$1.47 \times 10^{-02}$}\\\hline
\footnotesize{\textbf{Fig. 4a}} & \footnotesize{$2.73 \times 10^{+01}$} & \footnotesize{($822.762$, $39$)} & \footnotesize{$2.46 \times 10^{-01}$}\\\hline<jupyter_text>Supplementary Table 2 (x-space)<jupyter_code># Fig. 2d: fracDiff_xpred_xpredtruth_x_Positronp_y.png
statTableDouble(statList12=xStatsList[0], statList13=xStatsListTruth[0], figLabel='2d', cline=True)

# Fig. 2e: fracDiff_xpred_xpredtruth_x_Positronp_z.png
statTableDouble(statList12=xStatsList[1], statList13=xStatsListTruth[1], figLabel='2e', cline=True)

# Fig. 2f: fracDiff_xpred_xpredtruth_x_PositronE.png
statTableDouble(statList12=xStatsList[2], statList13=xStatsListTruth[2], figLabel='2f', cline=False)

# Fig. 4b: fracDiff_xpred_xpredtruth_x_M_Z.png
statTableDouble(statList12=xIMstatsList[0], statList13=xIMstatsListTruth[0], figLabel='4b', cline=False)<jupyter_output>\footnotesize{\textbf{Fig. 2d}} & \footnotesize{$4.22 \times 10^{-01}$} & \footnotesize{($1.391$, $23$)} & \footnotesize{$3.48 \times 10^{-03}$} & \footnotesize{$1.05 \times 10^{+00}$} & \footnotesize{($37.560$, $23$)} & \footnotesize{$1.22 \times 10^{-02}$}\\\cline{2-7}
\footnotesize{\textbf{Fig. 2e}} & \footnotesize{$3.71 \times 10^{+00}$} & \footnotesize{($1.523$, $26$)} & \footnotesize{$1.03 \times 10^{-02}$} & \footnotesize{$9.53 \times 10^{+00}$} & \footnotesize{($4.775$, $26$)} & \footnotesize{$7.49 \times 10^{-03}$}\\\cline{2-7}
\footnotesize{\textbf{Fig. 2f}} & \footnotesize{$6.64 \times 10^{-01}$} & \footnotesize{($0.489$, $26$)} & \footnotesize{$3.19 \times 10^{-03}$} & \footnotesize{$3.64 \times 10^{+00}$} & \footnotesize{($9.370$, $26$)} & \footnotesize{$2.00 \times 10^{-02}$}\\\hline
\footnotesize{\textbf{Fig. 4b}} & \footnotesize{$7.28 \times 10^{-01}$} & \footnotesize{($5.055$, $39$)} & \footnotesize{$2.61 \times 10^{-02}$} & \footnotesize{$7.15 \times 10^{-01}$} & \foo[...]
{'fraction_non_alphanumeric': 0.16056688989254012, 'fraction_numerical': 0.06826558687639517, 'mean_word_length': 5.127226463104326, 'pattern_counts': {'":': 0, '<': 104, '<?xml version=': 0, '>': 94, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '25', 'dir': 'jupyter-structured-clean-dedup', 'id': '664995', 'license': 'MIT', 'n_tokens_mistral': 8271, 'n_tokens_neox': 7468, 'n_words': 1761, 'path': 'PaperPlots/ppzeePaperPlots.ipynb', 'repo_name': 'yiboyang/otus', 'size': '19263'}