# B/RK analyzer observables and histograms (protocode)
#
#
# Mikael Mieskolainen, 2020
# m.mieskolainen@imperial.ac.uk
import bz2
import numpy as np
import iceplot
import icebrk.tools as tools
obs_M = {
# Axis limits
'xlim' : (4.5, 6.5),
'ylim' : None,
'xlabel' : r'System $M$',
'ylabel' : r'Candidates',
'units' : {'x': 'GeV', 'y': ''},
'label' : r'3-body system invariant mass',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0.0, 10.0 + 0.1, 0.1),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : True
}
obs_Pt = {
# Axis limits
'xlim' : (0, 25.0),
'ylim' : None,
'xlabel' : r'System $P_t$',
'ylabel' : r'Candidates',
'units' : {'x': 'GeV', 'y': ''},
'label' : r'System tranverse momentum',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0.0, 25.0 + 1.0, 1.0),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : True
}
obs_q2 = {
# Axis limits
'xlim' : (0, 12.0),
'ylim' : None,
'xlabel' : r'Electron pair $q^2$',
'ylabel' : r'Candidates',
'units' : {'x': r'GeV$^2$', 'y': ''},
'label' : r'Electron system invariant squared',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0, 12 + 0.4, 0.4),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : True
}
obs_pt_l1 = {
# Axis limits
'xlim' : (0, 12.0),
'ylim' : None,
'xlabel' : r'Leading electron $p_t$',
'ylabel' : r'Candidates',
'units' : {'x': r'GeV', 'y': ''},
'label' : r'Leading $e$ transverse momentum',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0, 12 + 0.5, 0.5),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
obs_pt_l2 = {
# Axis limits
'xlim' : (0, 12.0),
'ylim' : None,
'xlabel' : r'Sub-leading electron $p_t$',
'ylabel' : r'Candidates',
'units' : {'x': r'GeV', 'y': ''},
'label' : r'Sub-leading transverse momentum',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0, 12 + 0.5, 0.5),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
obs_pt_k = {
# Axis limits
'xlim' : (0, 12.0),
'ylim' : None,
'xlabel' : r'Kaon $p_t$',
'ylabel' : r'Candidates',
'units' : {'x': r'GeV', 'y': ''},
'label' : r'Kaon transverse momentum',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(0, 12 + 0.5, 0.5),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
# ** MC ONLY **
obs_first_t3i = {
# Axis limits
'xlim' : (-1, 20),
'ylim' : None,
'xlabel' : r'First signal triplet',
'ylabel' : r'Events',
'units' : {'x': r'MC index', 'y': ''},
'label' : r'First signal triplet',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(-1, 20 + 1, 1),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
obs_last_t3i = {
# Axis limits
'xlim' : (-1, 20),
'ylim' : None,
'xlabel' : r'Last signal triplet',
'ylabel' : r'Events',
'units' : {'x': r'MC index', 'y': ''},
'label' : r'Last signal triplet',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(-1, 20 + 0.5, 0.5),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
obs_N_signal_t3 = {
# Axis limits
'xlim' : (-1, 20),
'ylim' : None,
'xlabel' : r'Number of signal triplets',
'ylabel' : r'Events',
'units' : {'x': '', 'y': ''},
'label' : r'Number of signal triplets',
'figsize' : (4,4),
# Histogramming
'bins' : np.arange(-1, 20 + 0.5, 0.5),
'density' : False,
# Function to calculate
'func' : None,
# Disk save
'pickle' : False
}
# Dictionary of all batch observables
obs_all = {
# JAGGED
'M' : obs_M,
'Pt' : obs_Pt,
'q2' : obs_q2,
'pt_l1' : obs_pt_l1,
'pt_l2' : obs_pt_l2,
'pt_k' : obs_pt_k,
# NORMAL
'first_t3i' : obs_first_t3i,
'last_t3i' : obs_last_t3i,
'N_signal_t3' : obs_N_signal_t3,
'N_signal_pfpf_t3' : obs_N_signal_t3,
'N_signal_lowlow_t3' : obs_N_signal_t3
}
[docs]
def calc_batch_observables(l1_p4, l2_p4, k_p4):
"""JAGGED + VECTORIZED (operates on event batch) observables.
Args:
l1_p4:
l2_p4:
k_p4:
Returns:
x: Observables
"""
x = {
'M' : None,
'Pt' : None,
'q2' : None,
'pt_l1' : None,
'pt_l2' : None,
'pt_k' : None
}
x['M'] = (l1_p4['e'] + l2_p4['e'] + k_p4['k']).mass
x['Pt'] = (l1_p4['e'] + l2_p4['e'] + k_p4['k']).pt
x['q2'] = (l1_p4['e'] + l2_p4['e']).mass2
x['pt_l1'] = l1_p4['e'].pt
x['pt_l2'] = l2_p4['e'].pt
x['pt_k'] = k_p4['k'].pt
return x
[docs]
def calc_batch_MC_observables(d, l1_p4, l2_p4, k_p4):
""" MC ONLY batch observables.
Args:
d:
l1_p4:
l2_p4:
k_p4:
Returns:
x
"""
x = {
}
return x
[docs]
def calc_observables(evt_index, d, l1_p4, l2_p4, k_p4, sets, MAXT3):
"""NON-JAGGED (NORMAL) observables.
Args:
l1_p4:
l2_p4:
k_p4:
Returns:
x: Observables
"""
x = {
}
return x
vals = {
'init' : False,
}
vals_pfpf = {
'init' : False,
}
vals_lowlow = {
'init' : False,
}
[docs]
def calc_MC_observables(evt_index, d, l1_p4, l2_p4, k_p4, sets, MAXT3):
"""MC ONLY observables.
Args:
evt_index:
d:
l1_p4:
l2_p4:
k_p4:
sets:
MAXT3:
Returns:
x: Observables
"""
if vals['init'] == False:
vals['init'] = True
for i in range(100):
vals[str(i)] = 0
if vals_pfpf['init'] == False:
vals_pfpf['init'] = True
for i in range(100):
vals_pfpf[str(i)] = 0
if vals_lowlow['init'] == False:
vals_lowlow['init'] = True
for i in range(100):
vals_lowlow[str(i)] = 0
x = {
'first_t3i' : None,
'last_t3i' : None,
'N_signal_t3' : None,
'N_signal_pfpf_t3' : None,
'N_signal_lowlow_t3' : None
}
# Number of signal triplets
x['N_signal_t3'] = np.sum(d['_BToKEE_is_signal'][evt_index])
x['N_signal_lowlow_t3'] = np.sum(d['_BToKEE_is_signal'][evt_index] & d['Electron_isLowPt'][d['BToKEE_l1Idx']][evt_index] & d['Electron_isLowPt'][d['BToKEE_l2Idx']][evt_index])
x['N_signal_pfpf_t3'] = np.sum(d['_BToKEE_is_signal'][evt_index] & d['Electron_isPF'][d['BToKEE_l1Idx']][evt_index] & d['Electron_isPF'][d['BToKEE_l2Idx']][evt_index])
vals[str(x['N_signal_t3'])] += 1
vals_lowlow[str(x['N_signal_lowlow_t3'])] += 1
vals_pfpf[str(x['N_signal_pfpf_t3'])] += 1
# The first signal index
x['first_t3i'] = tools.index_of_first_signal(evt_index, d, sets, MAXT3)
# The last signal index
x['last_t3i'] = tools.index_of_last_signal(evt_index, d, sets, MAXT3)
#print(list(vals.values())[0:22])
#print(list(vals_lowlow.values())[0:22])
#print(list(vals_pfpf.values())[0:22])
return x
[docs]
def pickle_files(iodir, N_algo, label, mode='rb'):
"""Open pickle files.
Args:
iodir:
N_algo:
label:
mode: mode = 'rb' (read binary), 'ab' (append binary), 'wb' (write binary)
Returns:
x: Observables
"""
wfile = {'S': dict(), 'B': dict()}
for ID in wfile.keys():
for i in range(N_algo):
wfile[ID][str(i)] = bz2.BZ2File(iodir + f'/{label}_{ID}_weights_{i}.bz2', mode)
obsfile = {'S': dict(), 'B': dict()}
for ID in obsfile.keys():
for obs in obs_all.keys():
if obs_all[obs]['pickle']:
obsfile[ID][obs] = bz2.BZ2File(iodir + f'/{label}_{ID}_obs_{obs}.bz2', mode)
return obsfile, wfile