def test_cluster_recombination_schemes():
recomb_schemes = {
'E_scheme': (983.28, -0.8676, 36.46),
'pt_scheme': (983.52, -0.8672, 0.00),
'pt2_scheme': (983.52, -0.8679, 0.00),
'Et_scheme': (983.55, -0.8672, 0.00),
'Et2_scheme': (983.55, -0.8679, 0.00),
'BIpt_scheme': (983.52, -0.8671, 0.00),
'BIpt2_scheme': (983.52, -0.8679, 0.00),
'WTA_pt_scheme': (983.52, -0.8684, 0.14),
'WTA_modp_scheme': (983.03, -0.8684, 0.14),
}
for recomb_scheme, jet in recomb_schemes.items():
sequence = cluster(get_event(),
R=0.6,
p=-1,
recomb_scheme=recomb_scheme)
jets = sequence.inclusive_jets()
assert jets[0].pt == approx(jet[0], abs=1e-2)
assert jets[0].eta == approx(jet[1], abs=1e-4)
assert jets[0].mass == approx(jet[2], abs=1e-2)
with pytest.raises(ValueError):
sequence = cluster(get_event(),
R=0.6,
p=-1,
recomb_scheme='invalid_scheme')
def jet_trimmer_1J(event, R0, R1, pt_cut):
# R0 = Clustering radius for the main jets
# R1 = Clustering radius for the subjets in the primary jet
# pt_cut = Threshold for subjets (relative to the primary jet it's a subjet of)
flattened_event = flatten(event)
sequence = cluster(flattened_event, R=R0, p=-1)
# Main jets
jets = sequence.inclusive_jets()
# In case we are missing a leading jet, break early
if len(jets) == 0:
return jets
# Take just the leading jet
jet0 = jets[0]
# Define a cut threshold that the subjets have to meet (i.e. 5% of the original jet pT)
jet0_max = jet0.pt
jet0_cut = jet0_max * pt_cut
# Grab the subjets by clustering with R1
subjets = cluster(jet0.constituents_array(), R=R1, p=1)
subjet_array = subjets.inclusive_jets()
j0 = []
if (subjet_array[0].pt >= jet0_cut):
# j0 = subjet_array[0].constituents_array()
for ij, subjet in enumerate(subjet_array):
# if (ij == 0):
# continue
if subjet.pt < jet0_cut:
# subjet doesn't meet the percentage cut on the original jet pT
continue
if subjet.pt >= jet0_cut:
# Get the subjets pt, eta, phi constituents
subjet_data = subjet.constituents_array()
j0.append(subjet_data)
# j0 = np.append(j0, subjet_data)
else:
j0 = subjet_array[0].constituents_array() * 0
jet = j0[0]
for i, subjet in enumerate(j0):
if i == 0:
continue
jet = np.append(jet, subjet)
sequence = cluster(jet, R=R0, p=-1)
jet = sequence.inclusive_jets()
return jet
def find_new_var_beta_3_kT(reclustered,clustered,pt_cut = 1): new_var = [] for sqn in range(len(reclustered)): tau_1 = np.zeros(3) tau_2 = np.zeros(2) d_0 = 0 jet = clustered[sqn][0] #Now, we need to recluster jet into two kT subjets sequence_Cluster = cluster(jet, R=0.2,p=1) jets_Cluster = sequence_Cluster.inclusive_jets() sub_1 = jets_Cluster[0] try: sub_2 = jets_Cluster[1] except: sub_2 = sub_1 for k in jet: if k.pt < pt_cut: continue d_0 = d_0 + k.pt*0.8 tau_1 = tau_1 + k.pt*np.array([R(k,jet)**0.5,R(k,jet)**1,R(k,jet)**2]) d_2 = min(R(sub_1,k),R(sub_2,k)) tau_2 = tau_2 + k.pt*np.array([d_2**1,d_2**2]) tau_1 = tau_1/d_0 tau_2 = tau_2/d_0 c = 0; d = 0.5; e = -1; a = 0; b = 0; #new_var.append(tau_1[0]**a*tau_1[1]**b*tau_1[2]**c*tau_2[0]**d*tau_2[1]**e) new_var.append(tau_1[0]**2*tau_2[0]**0.5*tau_2[1]**-1) return(new_var)
def pythia_sim(cmd_file, part_name="unnamed", make_plots=False):
# The main simulation. Takes a cmd_file as input. part_name
# is the name of the particle we're simulating decays from.
# Only necessary for titling graphs.
# Returns an array of 2D histograms, mapping eta, phi, with transverse
# energy.
pythia = Pythia(cmd_file, random_state=1)
selection = ((STATUS == 1) & ~HAS_END_VERTEX)
unclustered_particles = []
a = 0
part_tensor = []
for event in pythia(events=num_events):
jets_particle_eta = []
jets_particle_phi = []
jets_particle_energy = []
vectors = event.all(selection)
sequence = cluster(vectors, R=1.0, p=-1, ep=True)
jets = sequence.inclusive_jets()
unclustered_particles.append(sequence.unclustered_particles())
part_data = []
num_b_jets = 0
for i, jet in enumerate(jets):
if np.abs(jet.userinfo[pid]) == 5: num_b_jets += 1
data = np.array((jet.mass, jet.eta, jet.phi, jet.pt))
if is_massless_or_isolated(jet): discarded_data.append(jet)
else: part_data.extend(data)
# Calculate total number of b and non-b jets in the event
# Histogram the particle data
part_data = np.array(part_data)
hist_array = hist_jet_vals(part_data) # is y-axis jet num or event num
return np.array(part_tensor)
def test_jet_area():
sequence = cluster(get_event(), R=0.6, p=-1, area='active')
jets = sequence.inclusive_jets()
for jet in jets:
area, error = jet.area
if len(jet) > 3: # TODO: need better way to test this
assert area > 0
def makeJets(tracks, R, p=-1, mode='old'):
# Cluster AK(R) jets
vectors = np.zeros(
tracks.size,
np.dtype([('pT', 'f8'), ('eta', 'f8'), ('phi', 'f8'), ('mass', 'f8')]))
i = 0
tracks = tracks[tracks.mag2 > 0]
for track in tracks:
vectors[i] = np.array((track.pt, track.eta, track.phi, track.mass),
np.dtype([('pT', 'f8'), ('eta', 'f8'),
('phi', 'f8'), ('mass', 'f8')]))
i += 1
sequence = pyjet.cluster(vectors, R=R, p=p)
jets_list = sequence.inclusive_jets()
if mode == 'old':
return jets_list
jets_pt = np.zeros(len(jets_list))
jets_eta = np.zeros(len(jets_list))
jets_phi = np.zeros(len(jets_list))
jets_E = np.zeros(len(jets_list))
i = 0
for jet in jets_list:
jets_pt[i] = jet.pt
jets_eta[i] = jet.eta
jets_phi[i] = jet.phi
jets_E[i] = jet.e
i += 1
jets = uproot_methods.TLorentzVectorArray.from_ptetaphie(
jets_pt, jets_eta, jets_phi, jets_E)
return jets
def createJetCollections(events_combined, nEvts, truthBit=False, dataLabel="UnlabelledData"):
if truthBit:
alljets,leadpT=createJetsWithTruth(events_combined, nEvts)
else:
leadpT = {}
alljets = {}
leadpT[dataLabel]=[]
alljets[dataLabel]=[]
for i in range(nEvts): #len(events_combined)):
if (i%(nEvts/10)==0):
print(dataLabel,i)
pass
pseudojets_input = np.zeros(
len([x for x in events_combined[i][::3] if x > 0]),
dtype=DTYPE_PTEPM
)
for j in range(700):
if (events_combined[i][j*3]>0):
pseudojets_input[j]['pT'] = events_combined[i][j*3]
pseudojets_input[j]['eta'] = events_combined[i][j*3+1]
pseudojets_input[j]['phi'] = events_combined[i][j*3+2]
pass
pass
sequence = cluster(pseudojets_input, R=1.0, p=-1)
jets = sequence.inclusive_jets(ptmin=20)
leadpT[dataLabel] += [jets[0].pt]
alljets[dataLabel] += [jets]
pass
return alljets, leadpT
def createJetsWithTruth(events_combined, nEvts):
leadpT = {}
alljets = {}
for mytype in ['background','signal']:
leadpT[mytype]=[]
alljets[mytype]=[]
for i in range(nEvts): #len(events_combined)):
if (i%(nEvts/10)==0):
print(mytype,i)
pass
issignal = events_combined[i][2100]
if (mytype=='background' and issignal):
continue
elif (mytype=='signal' and issignal==0):
continue
pseudojets_input = np.zeros(len([x for x in events_combined[i][::3] if x > 0]), dtype=DTYPE_PTEPM)
for j in range(700):
if (events_combined[i][j*3]>0):
pseudojets_input[j]['pT'] = events_combined[i][j*3]
pseudojets_input[j]['eta'] = events_combined[i][j*3+1]
pseudojets_input[j]['phi'] = events_combined[i][j*3+2]
pass
pass
sequence = cluster(pseudojets_input, R=1.0, p=-1)
jets = sequence.inclusive_jets(ptmin=20)
leadpT[mytype] += [jets[0].pt]
alljets[mytype] += [jets]
pass
return alljets,leadpT
def makeJets_Comparisons(jet_algo, R, particles):
eta_min, eta_max = -4., 4.
bins = 200
cone = R * 10
eta = np.linspace(eta_min, eta_max,
bins + 1)[:-1] + (eta_max - eta_min) / (2 * bins)
phi = np.linspace(-np.pi, np.pi, bins + 1)[:-1] + (np.pi / bins)
X, Y = np.meshgrid(phi, eta)
event = np.zeros(len(particles), dtype=DTYPE_PTEPM)
event['pT'] = particles['pT'] #[p.pt for p in particles]
event['eta'] = particles['eta'] # [p.eta for p in particles]
event['phi'] = particles['phi'] # [p.phi for p in particles]
ghosts = np.zeros(eta.shape[0] * phi.shape[0], dtype=DTYPE_PTEPM)
ghosts['pT'] = 1e-8
ghosts['eta'] = Y.ravel()
ghosts['phi'] = X.ravel()
# add ghosts to the event
event = np.concatenate([event, ghosts], axis=0)
# p = -1 (ak), 0 (CA), 1 (kt)
sequence = cluster(event, R=R, p=jet_algo)
jets = sequence.inclusive_jets(ptmin=30)
# print(jets)
return jets
def cluster_to_jets(events_combined):
alljets = {}
for mytype in ['background', 'signal']:
alljets[mytype] = []
#for i in range(np.shape(events_combined)[1]):
for i in range(NUMBER_OF_EVENTS):
if (i % (NUMBER_OF_EVENTS / 10) == 0):
print(mytype, i)
pass
issignal = events_combined[i][2100]
if (mytype == 'background' and issignal):
continue
elif (mytype == 'signal' and issignal == 0):
continue
pseudojets_input = np.zeros(len(
[x for x in events_combined[i][::3] if x > 0]),
dtype=DTYPE_PTEPM)
for j in range(700):
if (events_combined[i][j * 3] > 0):
pseudojets_input[j]['pT'] = events_combined[i][j * 3]
pseudojets_input[j]['eta'] = events_combined[i][j * 3 + 1]
pseudojets_input[j]['phi'] = events_combined[i][j * 3 + 2]
pass
pass
sequence = cluster(pseudojets_input, R=1.0, p=-1)
jets = sequence.inclusive_jets(ptmin=20)
alljets[mytype] += [jets]
pass
return alljets
def runDefaultJetClustering(self, event, eventTypes=[]):
for evtType in eventTypes:
if "signal" in evtType or "background" in evtType:
isSignal = event[2100]
if (evtType == 'background'
and isSignal) or (evtType == 'signal'
and isSignal == 0):
continue
self.isSignal = isSignal
self.type = evtType
#actual jet clustering, taken from LHCOlympics page
pseudojets_input = np.zeros(len([x for x in event[::3] if x > 0]),
dtype=DTYPE_PTEPM)
for j in range(700):
if (event[j * 3] > 0):
pseudojets_input[j]['pT'] = event[j * 3]
pseudojets_input[j]['eta'] = event[j * 3 + 1]
pseudojets_input[j]['phi'] = event[j * 3 + 2]
pass
pass
sequence = cluster(pseudojets_input, R=1.0, p=-1)
jets = sequence.inclusive_jets(ptmin=20)
self.calculateSubjettiness(jets)
self.allJets, self.allConstituents = self.convertJetsToDType(jets)
pass
pass
def __enter__(self):
num_const=self.data.shape[1]/3
truth_flag=False
if not (num_const)%1==0:
truth_flag=True
for N,event in self.data.iterrows():
if truth_flag:
truth_label=event[self.data.shape[1]-1]
pseudojets=np.zeros(int(num_const)*3, dtype=DTYPE_PTEPM)
cut=int(num_const)
for j in range(int(num_const)):
if event[j*3]==0.0:
cut=j
break
pseudojets[j]['pT' ] = event[j*3]
pseudojets[j]['eta'] = event[j*3+1]
pseudojets[j]['phi'] = event[j*3+2]
#...cluster jets:
pseudojets = np.sort(pseudojets[:cut], order='pT')
sequence = cluster(pseudojets, R=self.R, p=self.p)
jets = sequence.inclusive_jets(ptmin=self.ptmin)
label = truth_label
yield N, jets, label
if N+1==self.stop: break
def cluster_lhco_testdata_part(h5file, R, p, evrange):
f = pd.read_hdf(h5file, start=evrange[0], stop=evrange[1])
events_combined = f.T
leadpT = {}
subleadpT = {}
leadM = {
} # take the two leading pT jets, look at mass of heaviest and lightest jets
subleadM = {}
alljets = {}
for mytype in ['background', 'signal']:
leadpT[mytype] = []
subleadpT[mytype] = []
leadM[mytype] = []
subleadM[mytype] = []
leadMpT[mytype] = []
subleadMpT[mytype] = []
alljets[mytype] = []
for i in range(evrange[0], evrange[1]): #len(events_combined)):
if (i % 10000 == 0):
print(mytype, i)
pass
issignal = events_combined[i][2100]
if (mytype == 'background' and issignal):
continue
elif (mytype == 'signal' and issignal == 0):
continue
pseudojets_input = np.zeros(len(
[x for x in events_combined[i][::3] if x > 0]),
dtype=DTYPE_PTEPM)
for j in range(700):
if (events_combined[i][j * 3] > 0):
pseudojets_input[j]['pT'] = events_combined[i][j * 3]
pseudojets_input[j]['eta'] = events_combined[i][j * 3 + 1]
pseudojets_input[j]['phi'] = events_combined[i][j * 3 + 2]
pass
pass
sequence = cluster(pseudojets_input, R=1.0, p=p)
jets = sequence.inclusive_jets(ptmin=20)
leadpT[mytype] += [jets[0].pt]
if len(jets) > 1:
subleadpT[mytype] += [jets[1].pt]
jets = sorted(jets,
key=lambda PseudoJet: PseudoJet.pt,
reverse=True)[0:2]
jets = sorted(jets,
key=lambda PseudoJet: PseudoJet.mass,
reverse=True)
leadM[mytype] += [jets[0].mass]
if len(jets) > 1:
subleadM[mytype] += [jets[1].mass]
alljets[mytype] += [jets]
return alljets, leadpT, subleadpT, leadM, subleadM
def calc_KtDeltaR(jet):
particles = jet.constituents_array()
if particles.shape[0] < 2:
return 0.0
subjets = pyjet.cluster(particles, R=0.4, p=1).exclusive_jets(2)
return CalcDeltaR(subjets[0], subjets[1])
def jet_clustering(event, R0, p = -1):
# R0 = Clustering radius for the main jets
flattened_event = flatten(event)
## p = -1, 0, 1 => anti-kt, C/A, kt Algorithm
sequence = cluster(flattened_event, R=R0, p= p)
# List of jets
jets = sequence.inclusive_jets()
return jets
def test_cluster():
sequence = cluster(get_event(), R=0.6, p=-1)
jets = sequence.inclusive_jets()
assert len(jets) == 91
assert jets[0].pt == approx(983.28, abs=2)
assert isinstance(jets[0].parents, tuple)
len(jets[0].parents) == 2
jets[0].parents[0].child.pt == jets[0].pt
jets[0].parents[0].child == jets[0]
# too few parameters specified for jet definition
with pytest.raises(RuntimeError):
cluster(get_event())
# hashable
hash(sequence)
hash(jets[0])
def test_jet_area():
if not USING_EXTERNAL_FASTJET:
raise SkipTest("using internal fastjet")
sequence = cluster(get_event(), R=0.6, p=-1, area='active')
jets = sequence.inclusive_jets()
for jet in jets:
area, error = jet.area
if len(jet) > 3: # TODO: need better way to test this
assert_true(area > 0)
def w_pf(event, R=0.8, p=1):
from pyjet import cluster
jet = cluster(event, R=R, p=p).exclusive_jets(1) # check if correct
delRs = np.sqrt((event['phi'][:, None]-np.array([jet[0].phi]))**2 +
(event['eta'][:, None]-np.array([jet[0].eta]))**2)
delRs = np.squeeze(delRs)
w_pf_num = event['pT']*delRs
return w_pf_num.sum()/(event['pT'].sum())
def pythia_sim(cmd_file, part_name="unnamed", make_plots=False):
# The main simulation. Takes a cmd_file as input. part_name
# is the name of the particle we're simulating decays from.
# Only necessary for titling graphs.
# Returns an array of 2D histograms, mapping eta, phi, with transverse
# energy.
pythia = Pythia(cmd_file, random_state=1)
selection = ((STATUS == 1) & ~HAS_END_VERTEX)
unclustered_particles = []
a = 0
part_tensor = []
sj_data_per_event = []
for event in pythia(events=num_events):
lead_jet_invalid = False
sub_jet_data = [] # There are multiple jets in each event
jets_particle_eta = []
jets_particle_phi = []
jets_particle_energy = []
vectors = event.all(selection)
sequence = cluster(vectors, R=1.0, p=-1, ep=True)
jets = sequence.inclusive_jets()
unclustered_particles.append(sequence.unclustered_particles())
part_data = []
for i, jet in enumerate(jets):
jet_data = (
jet.mass, jet.eta, jet.phi, jet.pt,
len(jet.constituents_array()), 2*jet.mass/jet.pt
)
part_data = return_particle_data(jet)
if is_massless_or_isolated(jet):
discarded_data.append(jet)
if i == 0: lead_jet_invalid = True
else:
jets_particle_eta.extend(part_data[0])
jets_particle_phi.extend(part_data[1])
jets_particle_energy.extend(part_data[2])
if i < 3:
sub_jet_data.append(jet_data)
lead_jet_valid = not lead_jet_invalid
if lead_jet_valid:
sj_data_per_event.append(np.array(sub_jet_data))
plt.figure()
part_tensor.append(plt.hist2d(jets_particle_eta, jets_particle_phi,
weights=jets_particle_energy, normed=True,
range=[(-5,5),(-1*np.pi,np.pi)],
bins=(20,32), cmap='binary')[0]) # We're only taking the
plt.close() # Zeroth element, which is the raw data of the 2D Histogram
if make_plots:
plt.xlabel("$\eta$")
plt.ylabel("$\phi$")
plt.title("Particles from "+part_name)
cbar = plt.colorbar()
cbar.set_label('Tranverse Energy of Each Particle ($GeV$)')
plt.savefig("hists/Jets_Particles_"+part_name+str(a)+".png")
a += 1
return np.array(part_tensor), np.array(sj_data_per_event)
def get_lund_history(jet, R, p):
# re-clustering the jet
clustered_jet = cluster(jet.constituents_array(), R=R, p=p)
splittings = []
# each level in the clustering history is represented by a list of subjets
# each subjet has [4-momentum,plane_id]
lund_history = []
# looping over the levels
for t in range(0, clustered_jet.n_exclusive_jets(0)):
# list momenta of all subjets at current splitting
j_obs = [[j.mass, j.pt, j.eta, j.phi]
for j in clustered_jet.exclusive_jets(t)]
j_obs.sort(key=lambda i: i[1], reverse=True)
j_obs_id = [[j_obs[i], i] for i in range(len(j_obs))]
if len(j_obs) <= 2:
lund_history.append(j_obs_id)
if len(j_obs) > 2:
# list momenta,id of all subjets at previous splitting
pj_obs_id = [j for j in lund_history[-1]]
pj_obs = [j[0] for j in lund_history[-1]]
# work out which subjet split, and label it p_obs
p_obs = [
pj_obs[i] for i in range(len(pj_obs)) if pj_obs[i] not in j_obs
]
p_obs_id = [
pj_obs_id[i] for i in range(len(pj_obs))
if pj_obs[i] not in j_obs
]
# work out which subjets didn't split, label them np_obs
np_obs = [
pj_obs[i] for i in range(len(pj_obs)) if pj_obs[i] in j_obs
]
np_obs_id = [
pj_obs_id[i] for i in range(len(pj_obs)) if pj_obs[i] in j_obs
]
# work out what it split into, and put them in d_obs
d_obs = [
j_obs[i] for i in range(len(j_obs)) if j_obs[i] not in pj_obs
]
d_obs.sort(key=lambda i: i[1], reverse=True)
pid = p_obs_id[0][1]
d_obs_id = [[d_obs[i], i + pid] for i in range(len(d_obs))]
if len(d_obs) == 2 and len(p_obs) == 1:
lund_history.append(np_obs_id + d_obs_id)
return lund_history
def pythia_sim(cmd_file, part_name="", make_plots=False):
# The main simulation. Takes a cmd_file as input. part_name
# is the name of the particle we're simulating decays from.
# Only necessary for titling graphs.
# Returns an array of 2D histograms, mapping eta, phi, with transverse
# energy.
if part_name == "":
for char in cmd_file:
if char == ".":
break
else:
part_name += char
pythia = Pythia(cmd_file, random_state=1)
selection = ((STATUS == 1) & ~HAS_END_VERTEX)
unclustered_particles = []
part_tensor = []
a = 0
for event in pythia(events=num_events):
jets_particle_eta = []
jets_particle_phi = []
jets_particle_energy = []
vectors = event.all(selection)
sequence = cluster(vectors, R=1.0, p=-1, ep=True)
jets = sequence.inclusive_jets()
unclustered_particles.append(sequence.unclustered_particles())
part_data = []
for i, jet in enumerate(jets):
part_data = return_particle_data(jet)
if is_massless_or_isolated(jet):
discarded_data.append(jet)
else:
jets_particle_eta.extend(part_data[0])
jets_particle_phi.extend(part_data[1])
jets_particle_energy.extend(part_data[2])
plt.figure()
part_tensor.append(
plt.hist2d(jets_particle_eta,
jets_particle_phi,
weights=jets_particle_energy,
density=True,
range=[(-5, 5), (-1 * np.pi, np.pi)],
bins=(20, 32),
cmap='binary')[0]) # We're only taking the
if not make_plots:
plt.close(
) # Zeroth element, which is the raw data of the 2D Histogram
if make_plots:
plt.xlabel("$\eta$")
plt.ylabel("$\phi$")
plt.title("Particles from " + part_name)
cbar = plt.colorbar()
cbar.set_label('Tranverse Energy of Each Particle ($GeV$)')
plt.savefig("hists/Jets_Particles_" + part_name + str(a) + ".png")
plt.close()
a += 1
return np.array(part_tensor)
def pythia_sim(cmd_file):
# The main simulation. Takes a cmd_file as input.
pythia = Pythia(cmd_file, random_state=1)
selection = ((STATUS == 1) & ~HAS_END_VERTEX)
unclustered_particles = []
for event in pythia(events=num_events):
vectors = event.all(selection)
sequence = cluster(vectors, R=0.4, p=-1, ep=True)
jets = sequence.inclusive_jets()
print(NoneType_statistics(jets, debug_data))
def _cluster_jets(self, jet_p4, n_jet, max_R):
pseudojets_input = np.zeros(len(jet_p4), dtype=DTYPE_PTEPM)
for i, p4 in enumerate(jet_p4):
pseudojets_input[i]['pT'] = p4.pt
pseudojets_input[i]['eta'] = p4.eta
pseudojets_input[i]['phi'] = p4.phi
pseudojets_input[i]['mass'] = p4.mass
sequence = cluster(pseudojets_input, R=max_R, p=1)
jets = sequence.exclusive_jets(n_jet) ## decluster
return jets
def findJetsSingleEvent(self, particles, jetR):
jets = cluster(particles, R=float(jetR),
p=-1).inclusive_jets() # p=-1 --> anti-kT
jetList = []
for jet in jets:
if not self.cut_jet(jet, jetR):
l = [jet.pt, jet.eta, jet.phi, jet.mass] + list(
jet.constituents_array())
jetList.append(l)
return jetList
def nsubjettiness(event, n, R=0.8, p=1):
from pyjet import cluster
subjets = cluster(event, R=R, p=p).exclusive_jets(n)
delRs = np.sqrt((event['phi'][:, None]-np.array([[J.phi for J in subjets]]))**2 +
(event['eta'][:, None]-np.array([[J.eta for J in subjets]]))**2)
taun = event['pT']*np.min(delRs, axis=1)
if event['pT'].sum() > 0:
return taun.sum()/(event['pT'].sum()*R)
else:
return None
def _internal_compute(self, tag, x, weight, mask, pt0, m0=None, mjj=None):
p4 = self.compute_p4(x[:, :, :4])
p4 *= weight[:, :, None]
p4 = p4.astype(np.float64)
n_batch = p4.shape[0]
pt = x[:, :, 0] * weight
for i in range(n_batch):
particle_mask = pt[i, :] > 0
particle_mask = np.logical_and(
particle_mask,
np.logical_and(
~np.isnan(p4[i]).sum(-1),
~np.isinf(p4[i]).sum(-1)
)
)
evt = p4[i][np.logical_and(
mask[i].astype(bool),
particle_mask
)]
evt = np.core.records.fromarrays(
evt.T,
names='E, px, py, pz',
formats='f8, f8, f8, f8'
)
seq = pyjet.cluster(evt, R=0.4, p=-1, ep=True)
jets = seq.inclusive_jets()
if len(jets) > 0:
self.dists['pt'][tag].append(jets[0].pt - pt0[i])
self.truth_pt[tag].append(pt0[i])
self.dists_2['pt'][tag][0].append(pt0[i])
self.dists_2['pt'][tag][1].append(jets[0].pt)
if m0 is not None:
self.dists['m'].append(jets[0].mass - m0[i])
self.dists_2['m'][tag][0].append(m0[i])
self.dists_2['m'][tag][1].append(jets[0].mass)
if mjj is not None:
if len(jets) > 1:
j0, j1 = jets[:2]
mjj_pred = np.sqrt(
(j0.e + j1.e) ** 2
- (j0.px + j1.px) ** 2
- (j0.py + j1.py) ** 2
- (j0.pz + j1.pz) ** 2
)
else:
mjj_pred = 0
if mjj[i] > 0:
self.dists['mjj'][tag].append(mjj_pred - mjj[i])
self.dists_2['mjj'][tag][0].append(mjj[i])
self.dists_2['mjj'][tag][1].append(mjj_pred)
def tn(jet, n): #t1 t2 t3 t21 t32
assert n >= 0
if n == 0:
return t0(jet)
particles = jet.constituents_array()
if len(particles) < n:
return -1
subjets = pyjet.cluster(particles, R=1.0, p=1).exclusive_jets(n)
subjets_array = [subjet.constituents_array() for subjet in subjets]
wta_axes = [a[np.argmax(a['pT'])] for a in subjets_array]
wta_axes = np.array(wta_axes, dtype=subjets_array[0].dtype)
return np.sum(particles['pT'] *
CalcDeltaRArray(particles, wta_axes).min(axis=0)) / t0(jet)
def makeJets(tracks, R):
# Cluster AK15 jets
vectors = np.zeros(tracks.size, np.dtype([('pT', 'f8'), ('eta', 'f8'),
('phi', 'f8'), ('mass', 'f8')]))
i = 0
for track in tracks:
vectors[i] = np.array((track.pt, track.eta, track.phi, track.mass),
np.dtype([('pT', 'f8'), ('eta', 'f8'),
('phi', 'f8'), ('mass', 'f8')]))
i += 1
sequence = pyjet.cluster(vectors, R=R, p=-1)
jetsAK15 = sequence.inclusive_jets()
return jetsAK15
def _cluster(epxpypyz, return_np=True, **kwargs):
import pyjet
kwargs.setdefault('algo', 'antikt')
kwargs.setdefault('R', .3)
kwargs.setdefault('ep', True)
jets = pyjet.cluster(np_to_structuredarray(epxpypyz),
algo='antikt',
R=0.3,
ep=True).inclusive_jets()
if return_np:
return np.array([[j.e, j.px, j.py, j.pz] for j in jets])
else:
return jets
def tau21(jet, subR=0.2):
'''Input: jet from the jet clustering result '''
jet_substruct_features = {}
seq = fj.cluster(jet, R=subR, algo='kt')
cnsts = jet.constituents()
cndts1 = seq.exclusive_jets(1)
tau1 = subjettiness(cndts1, cnsts)
if (len(cnsts) > 1):
cndts2 = seq.exclusive_jets(2)
tau2 = subjettiness(cndts2, cnsts)
else:
tau2 = 0
return tau2 / tau1
