Skip to content

Commit ee0e559

Browse files
cmsbuildcmsbuild
authored
Merge pull request #48674 from kk428/new_branch_PR
LST: Modifying LST to create efficiency plots for jets
2 parents a9f1be3 + ca3846f commit ee0e559

File tree

10 files changed

+724
-60
lines changed

10 files changed

+724
-60
lines changed

‎RecoTracker/LSTCore/standalone/analysis/jets/myjets.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
###############################################
2+
#
3+
# Library of functions used when dealing
4+
# with jets, especially reformat_jets.py
5+
#
6+
###############################################
7+
8+
import fastjet
9+
from pyjet import cluster
10+
from particle import Particle
11+
import numpy as np
12+
import awkward as ak
13+
import vector
14+
15+
16+
# Takes an entry from a tree and extracts lists of key parameters.
17+
def getLists(entry, hardSc = False, pTcut=True):
18+
# This applies to the entire event
19+
pdgidList = entry.sim_pdgId
20+
evLen = len(pdgidList)
21+
22+
pTList = np.ones(evLen, dtype=np.float64)*-999
23+
etaList = np.ones(evLen, dtype=np.float64)*-999
24+
phiList = np.ones(evLen, dtype=np.float64)*-999
25+
massList = np.ones(evLen, dtype=np.float64)*-999
26+
27+
# Putting the data in the right format
28+
for j in range(evLen):#range(len(simEvnp)):
29+
if (hardSc and entry.sim_event[j] !=0 ):
30+
# print(f"event != 0 {entry.sim_event[j]}")
31+
continue
32+
if(pTcut and entry.sim_q[j] != 0 and entry.sim_pt[j] < 0.75 ):
33+
continue
34+
vtxX = entry.simvtx_x[entry.sim_parentVtxIdx[j]]
35+
vtxY = entry.simvtx_y[entry.sim_parentVtxIdx[j]]
36+
if(np.sqrt(vtxX**2 + vtxY**2)>10):
37+
continue
38+
pdgid = pdgidList[j]
39+
massList[j] = Particle.from_pdgid(pdgid).mass/1000.
40+
41+
pTList[j] = entry.sim_pt[j]
42+
etaList[j] = entry.sim_eta[j]
43+
phiList[j] = entry.sim_phi[j]
44+
45+
massList = massList[massList != -999]
46+
pTList = pTList[pTList != -999]
47+
etaList = etaList[etaList != -999]
48+
phiList = phiList[phiList != -999]
49+
50+
return pTList, etaList, phiList, massList
51+
52+
53+
# Takes an entry from a tree and extracts particles from it. Uses
54+
# those particles to create jets, which it returns.
55+
def createJets(pTList, etaList, phiList, massList):
56+
57+
jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4)
58+
59+
length = np.size(pTList)
60+
fjs = []
61+
62+
for j in range(length):
63+
f=fastjet.PseudoJet(pseudojet_from_pt_eta_phi_m(pTList[j], etaList[j], phiList[j], massList[j]))
64+
fjs.append(f)
65+
66+
cluster = fastjet.ClusterSequence(fjs, jetdef)
67+
jets = cluster.inclusive_jets(ptmin=20)
68+
69+
return cluster, jets
70+
71+
def plotOneJet(jet, name):
72+
const = jet.constituents_array()
73+
plt.scatter(const["eta"], const["phi"], c='green')
74+
plt.scatter([jet.eta], [jet.phi], c="red")
75+
plt.xlabel(r'$\eta$')
76+
plt.ylabel(r'$\phi$')
77+
plt.savefig(name)
78+
plt.clf()
79+
80+
def matchArr(jetArrpT, jetArreta, jetArrphi, treeArrpT, treeArreta, treeArrphi, evnum, jetnum):
81+
indexArr = np.ones(len(jetArrpT))*-999
82+
83+
for i in range(len(jetArrpT)):
84+
for j in range(len(treeArrpT)):
85+
if(np.abs(jetArrpT[i]-treeArrpT[j])<0.00001 and np.abs(jetArreta[i]-treeArreta[j])<0.00001
86+
and np.abs(np.cos(jetArrphi[i])-np.cos(treeArrphi[j]))<0.00001):
87+
if(indexArr[i]!=-999):
88+
print(f"Error: double matched at Event={evnum}, Jet={jetnum} i={i}")
89+
continue
90+
indexArr[i] = j
91+
if(indexArr[i]==-999.0):
92+
print(f"Error: Event={evnum}, Jet={jetnum} i={i}, pT = {jetArrpT[i]}, eta = {jetArreta[i]}, phi = {jetArrphi[i]}")
93+
94+
return indexArr
95+
96+
def pseudojet_from_pt_eta_phi_m(pt, eta, phi, mass):
97+
# Convert (pt, eta, phi, mass) to (px, py, pz, E) and create a PseudoJet.
98+
px = pt * np.cos(phi)
99+
py = pt * np.sin(phi)
100+
pz = pt * np.sinh(eta)
101+
energy = np.sqrt(px**2 + py**2 + pz**2 + mass**2)
102+
103+
return fastjet.PseudoJet(px, py, pz, energy)

‎RecoTracker/LSTCore/standalone/analysis/jets/reformat_jets.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
##########################################################
2+
#
3+
# The file trackingNtuple.root has the data I need.
4+
# I can section it off into jets.
5+
# I then put these jets into a new n-tuple,
6+
#
7+
##########################################################
8+
9+
import matplotlib.pyplot as plt
10+
import ROOT
11+
from ROOT import TFile
12+
from myjets import getLists, createJets, matchArr
13+
import numpy as np
14+
15+
# Load existing tree
16+
file = TFile("/data2/segmentlinking/CMSSW_12_2_0_pre2/trackingNtuple_ttbar_PU200.root")
17+
# file = TFile("trackingNtuple100.root")
18+
old_tree = file["trackingNtuple"]["tree"]
19+
20+
# Create a new ROOT file to store the new TTree
21+
new_file = ROOT.TFile("new_tree_temp.root", "RECREATE")
22+
23+
# Create a new subdirectory in the new file
24+
new_directory = new_file.mkdir("trackingNtuple")
25+
26+
# Change the current directory to the new subdirectory
27+
new_directory.cd()
28+
29+
# Create a new TTree with the same structure as the old one but empty
30+
new_tree = old_tree.CloneTree(0)
31+
32+
# Account for bug in 12_2_X branch
33+
new_tree.SetBranchStatus("ph2_bbxi", False)
34+
35+
# Create a variable to hold the new leaves' data (a list of floats)
36+
new_leaf_deltaEta = ROOT.std.vector('float')()
37+
new_leaf_deltaPhi = ROOT.std.vector('float')()
38+
new_leaf_deltaR = ROOT.std.vector('float')()
39+
new_leaf_jet_eta = ROOT.std.vector('float')()
40+
new_leaf_jet_phi = ROOT.std.vector('float')()
41+
new_leaf_jet_pt = ROOT.std.vector('float')()
42+
43+
44+
# Create a new branch in the tree
45+
new_tree.Branch("sim_deltaEta", new_leaf_deltaEta)
46+
new_tree.Branch("sim_deltaPhi", new_leaf_deltaPhi)
47+
new_tree.Branch("sim_deltaR", new_leaf_deltaR)
48+
new_tree.Branch("sim_jet_eta", new_leaf_jet_eta)
49+
new_tree.Branch("sim_jet_phi", new_leaf_jet_phi)
50+
new_tree.Branch("sim_jet_pt", new_leaf_jet_pt)
51+
52+
# Loop over entries in the old tree
53+
for ind in range(old_tree.GetEntries()):
54+
old_tree.GetEntry(ind)
55+
56+
# Clear the vector to start fresh for this entry
57+
new_leaf_deltaEta.clear()
58+
new_leaf_deltaPhi.clear()
59+
new_leaf_deltaR.clear()
60+
new_leaf_jet_eta.clear()
61+
new_leaf_jet_phi.clear()
62+
new_leaf_jet_pt.clear()
63+
64+
# Creates the lists that will fill the leaves
65+
pTList, etaList, phiList, massList = getLists(old_tree, hardSc=True, pTcut=True)
66+
cluster, jets = createJets(pTList, etaList, phiList, massList)
67+
68+
jetIndex = np.array([])
69+
eta_diffs = np.array([])
70+
phi_diffs = np.array([])
71+
deltaRs = np.array([])
72+
jet_eta = np.array([])
73+
jet_phi = np.array([])
74+
jet_pt = np.array([])
75+
76+
for j, jet in enumerate(jets):
77+
const = jet.constituents()
78+
c_len = len(const)
79+
c_pts = np.zeros(c_len)
80+
c_etas = np.zeros(c_len)
81+
c_phis = np.zeros(c_len)
82+
83+
for k in range(c_len):
84+
c_pts[k] = const[k].pt()
85+
c_etas[k] = const[k].eta()
86+
c_phis[k] = const[k].phi()
87+
88+
# Reorder particles within jet (jet clustering does not respect original index)
89+
jetIndex = np.append(jetIndex, matchArr(c_pts, c_etas, c_phis,
90+
np.array(old_tree.sim_pt), np.array(old_tree.sim_eta), np.array(old_tree.sim_phi),
91+
ind, j)) # order restored by matching pT
92+
jetIndex = jetIndex.astype(int)
93+
94+
# Compute the distance to jet
95+
etaval = c_etas-jet.eta()
96+
phival = c_phis-jet.phi()
97+
eta_diffs = np.append(eta_diffs, etaval)
98+
phi_diffs = np.append(phi_diffs, phival)
99+
100+
rval = np.sqrt(etaval**2 + np.arccos(np.cos(phival))**2)
101+
deltaRs = np.append(deltaRs, rval)
102+
103+
# Save values of closest jet
104+
jet_eta_val = np.ones(c_len)*jet.eta()
105+
jet_eta = np.append(jet_eta, jet_eta_val)
106+
jet_phi_val = np.ones(c_len)*jet.phi()
107+
jet_phi = np.append(jet_phi, jet_phi_val)
108+
jet_pt_val = np.ones(c_len)*jet.pt()
109+
jet_pt = np.append(jet_pt, jet_pt_val)
110+
111+
# Reordering branches appropriately
112+
length = len(np.array(old_tree.sim_pt))
113+
114+
re_eta_diffs = np.ones(length)*(-999)
115+
re_phi_diffs = np.ones(length)*(-999)
116+
re_deltaRs = np.ones(length)*(-999)
117+
re_jet_eta = np.ones(length)*(-999)
118+
re_jet_phi = np.ones(length)*(-999)
119+
re_jet_pt = np.ones(length)*(-999)
120+
121+
for i in range(len(jetIndex)):
122+
re_eta_diffs[jetIndex[i]] = eta_diffs[i]
123+
re_phi_diffs[jetIndex[i]] = phi_diffs[i]
124+
re_deltaRs[jetIndex[i]] = deltaRs[i]
125+
re_jet_eta[jetIndex[i]] = jet_eta[i]
126+
re_jet_phi[jetIndex[i]] = jet_phi[i]
127+
re_jet_pt[jetIndex[i]] = jet_pt[i]
128+
129+
# Add the list elements to the vector
130+
for value in re_eta_diffs:
131+
new_leaf_deltaEta.push_back(value)
132+
for value in re_phi_diffs:
133+
new_leaf_deltaPhi.push_back(value)
134+
for value in re_deltaRs:
135+
new_leaf_deltaR.push_back(value)
136+
for value in re_jet_eta:
137+
new_leaf_jet_eta.push_back(value)
138+
for value in re_jet_phi:
139+
new_leaf_jet_phi.push_back(value)
140+
for value in re_jet_pt:
141+
new_leaf_jet_pt.push_back(value)
142+
143+
# Fill the tree with the new values
144+
new_tree.Fill()
145+
146+
# break # Just for testing purposes
147+
148+
# new_tree.Scan("[email protected]():[email protected]()")
149+
150+
# Write the tree back to the file
151+
new_tree.Write()
152+
153+
# Debugging: print new_tree events
154+
# new_tree.GetEntry(0) # only look at first event
155+
# for i in range(3): # only look at first 3 tracks in event
156+
# print(f"Track {i}: sim_phi = {new_tree.sim_phi[i]}\t sim_eta = {new_tree.sim_eta[i]} \t sim_pt = {new_tree.sim_pt[i]} \t sim_deltaR = {new_tree.sim_deltaR[i]}")
157+
158+
new_file.Close()
159+
file.Close()
160+

‎RecoTracker/LSTCore/standalone/bin/lst.cc‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,8 @@ int main(int argc, char **argv) {
7070
"I,job_index",
7171
"job_index of split jobs (--nsplit_jobs must be set. index starts from 0. i.e. 0, 1, 2, 3, etc...)",
7272
cxxopts::value<int>())("3,tc_pls_triplets", "Allow triplet pLSs in TC collection")(
73-
"2,no_pls_dupclean", "Disable pLS duplicate cleaning (both steps)")("h,help", "Print help");
73+
"2,no_pls_dupclean", "Disable pLS duplicate cleaning (both steps)")("h,help", "Print help")(
74+
"J,jet_branches", "Accounts for specific jet branches in input root file for testing");
7475

7576
auto result = options.parse(argc, argv);
7677

@@ -91,6 +92,7 @@ int main(int argc, char **argv) {
9192

9293
// A default value one
9394
TString TrackingNtupleDir = gSystem->Getenv("TRACKINGNTUPLEDIR");
95+
9496
if (ana.input_raw_string.EqualTo("muonGun"))
9597
ana.input_file_list_tstring = TString::Format("%s/trackingNtuple_10mu_pt_0p5_2.root", TrackingNtupleDir.Data());
9698
else if (ana.input_raw_string.EqualTo("muonGun_highPt"))
@@ -261,6 +263,10 @@ int main(int argc, char **argv) {
261263
// --no_pls_dupclean
262264
ana.no_pls_dupclean = result["no_pls_dupclean"].as<bool>();
263265

266+
//_______________________________________________________________________________
267+
// --jet_branches
268+
ana.jet_branches = result["jet_branches"].as<bool>();
269+
264270
// Printing out the option settings overview
265271
std::cout << "=========================================================" << std::endl;
266272
std::cout << " Running for Acc = " << alpaka::getAccName<ALPAKA_ACCELERATOR_NAMESPACE::Acc3D>() << std::endl;

‎RecoTracker/LSTCore/standalone/code/core/AnalysisConfig.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,9 @@ class AnalysisConfig {
131131

132132
// Boolean to disable pLS duplicate cleaning
133133
bool no_pls_dupclean;
134+
135+
// Boolean to enable jet branches
136+
bool jet_branches;
134137
};
135138

136139
extern AnalysisConfig ana;

0 commit comments

Comments
 (0)