analysis/jets/reformat_jets.py

0001 ##########################################################
0002 #
0003 # The file trackingNtuple.root has the data I need.
0004 # I can section it off into jets.
0005 # I then put these jets into a new n-tuple,
0006 #
0007 ##########################################################
0008
0009 import matplotlib.pyplot as plt
0010 import ROOT
0011 from ROOT import TFile
0012 from myjets import getLists, createJets, matchArr
0013 import numpy as np
0014
0015 # Load existing tree
0016 file =  TFile("/data2/segmentlinking/CMSSW_12_2_0_pre2/trackingNtuple_ttbar_PU200.root")
0017 # file = TFile("trackingNtuple100.root")
0018 old_tree = file["trackingNtuple"]["tree"]
0019
0020 # Create a new ROOT file to store the new TTree
0021 new_file = ROOT.TFile("new_tree_temp.root", "RECREATE")
0022
0023 # Create a new subdirectory in the new file
0024 new_directory = new_file.mkdir("trackingNtuple")
0025
0026 # Change the current directory to the new subdirectory
0027 new_directory.cd()
0028
0029 # Create a new TTree with the same structure as the old one but empty
0030 new_tree = old_tree.CloneTree(0)
0031
0032 # Account for bug in 12_2_X branch
0033 new_tree.SetBranchStatus("ph2_bbxi", False)
0034
0035 # Create a variable to hold the new leaves' data (a list of floats)
0036 new_leaf_deltaEta = ROOT.std.vector('float')()
0037 new_leaf_deltaPhi = ROOT.std.vector('float')()
0038 new_leaf_deltaR = ROOT.std.vector('float')()
0039 new_leaf_jet_eta = ROOT.std.vector('float')()
0040 new_leaf_jet_phi = ROOT.std.vector('float')()
0041 new_leaf_jet_pt = ROOT.std.vector('float')()
0042
0043
0044 # Create a new branch in the tree
0045 new_tree.Branch("sim_deltaEta", new_leaf_deltaEta)
0046 new_tree.Branch("sim_deltaPhi", new_leaf_deltaPhi)
0047 new_tree.Branch("sim_deltaR", new_leaf_deltaR)
0048 new_tree.Branch("sim_jet_eta", new_leaf_jet_eta)
0049 new_tree.Branch("sim_jet_phi", new_leaf_jet_phi)
0050 new_tree.Branch("sim_jet_pt", new_leaf_jet_pt)
0051
0052 # Loop over entries in the old tree
0053 for ind in range(old_tree.GetEntries()):
0054     old_tree.GetEntry(ind)
0055
0056     # Clear the vector to start fresh for this entry
0057     new_leaf_deltaEta.clear()
0058     new_leaf_deltaPhi.clear()
0059     new_leaf_deltaR.clear()
0060     new_leaf_jet_eta.clear()
0061     new_leaf_jet_phi.clear()
0062     new_leaf_jet_pt.clear()
0063
0064     # Creates the lists that will fill the leaves
0065     pTList, etaList, phiList, massList = getLists(old_tree, hardSc=True, pTcut=True)
0066     cluster, jets = createJets(pTList, etaList, phiList, massList)
0067
0068     jetIndex = np.array([])
0069     eta_diffs = np.array([])
0070     phi_diffs = np.array([])
0071     deltaRs = np.array([])
0072     jet_eta = np.array([])
0073     jet_phi = np.array([])
0074     jet_pt = np.array([])
0075
0076     for j, jet in enumerate(jets):
0077         const = jet.constituents()
0078         c_len = len(const)
0079         c_pts = np.zeros(c_len)
0080         c_etas = np.zeros(c_len)
0081         c_phis = np.zeros(c_len)
0082
0083         for k in range(c_len):
0084             c_pts[k] = const[k].pt()
0085             c_etas[k] = const[k].eta()
0086             c_phis[k] = const[k].phi()
0087
0088         # Reorder particles within jet (jet clustering does not respect original index)
0089         jetIndex = np.append(jetIndex, matchArr(c_pts, c_etas, c_phis,
0090                                                 np.array(old_tree.sim_pt), np.array(old_tree.sim_eta), np.array(old_tree.sim_phi),
0091                                                 ind, j)) # order restored by matching pT
0092         jetIndex = jetIndex.astype(int)
0093
0094         # Compute the distance to jet
0095         etaval = c_etas-jet.eta()
0096         phival = c_phis-jet.phi()
0097         eta_diffs = np.append(eta_diffs, etaval)
0098         phi_diffs = np.append(phi_diffs, phival)
0099
0100         rval = np.sqrt(etaval**2 + np.arccos(np.cos(phival))**2)
0101         deltaRs = np.append(deltaRs, rval)
0102
0103         # Save values of closest jet
0104         jet_eta_val = np.ones(c_len)*jet.eta()
0105         jet_eta = np.append(jet_eta, jet_eta_val)
0106         jet_phi_val = np.ones(c_len)*jet.phi()
0107         jet_phi = np.append(jet_phi, jet_phi_val)
0108         jet_pt_val = np.ones(c_len)*jet.pt()
0109         jet_pt = np.append(jet_pt, jet_pt_val)
0110
0111     # Reordering branches appropriately
0112     length = len(np.array(old_tree.sim_pt))
0113
0114     re_eta_diffs = np.ones(length)*(-999)
0115     re_phi_diffs = np.ones(length)*(-999)
0116     re_deltaRs = np.ones(length)*(-999)
0117     re_jet_eta = np.ones(length)*(-999)
0118     re_jet_phi = np.ones(length)*(-999)
0119     re_jet_pt = np.ones(length)*(-999)
0120
0121     for i in range(len(jetIndex)):
0122         re_eta_diffs[jetIndex[i]] = eta_diffs[i]
0123         re_phi_diffs[jetIndex[i]] = phi_diffs[i]
0124         re_deltaRs[jetIndex[i]] = deltaRs[i]
0125         re_jet_eta[jetIndex[i]] = jet_eta[i]
0126         re_jet_phi[jetIndex[i]] = jet_phi[i]
0127         re_jet_pt[jetIndex[i]] = jet_pt[i]
0128
0129     # Add the list elements to the vector
0130     for value in re_eta_diffs:
0131         new_leaf_deltaEta.push_back(value)
0132     for value in re_phi_diffs:
0133         new_leaf_deltaPhi.push_back(value)
0134     for value in re_deltaRs:
0135         new_leaf_deltaR.push_back(value)
0136     for value in re_jet_eta:
0137         new_leaf_jet_eta.push_back(value)
0138     for value in re_jet_phi:
0139         new_leaf_jet_phi.push_back(value)
0140     for value in re_jet_pt:
0141         new_leaf_jet_pt.push_back(value)
0142
0143     # Fill the tree with the new values
0144     new_tree.Fill()
0145
0146     # break # Just for testing purposes
0147
0148 # new_tree.Scan("sim_pt@.size():sim_jet_pt@.size()")
0149
0150 # Write the tree back to the file
0151 new_tree.Write()
0152
0153 # Debugging: print new_tree events
0154 # new_tree.GetEntry(0) # only look at first event
0155 # for i in range(3): # only look at first 3 tracks in event
0156 #     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]}")
0157
0158 new_file.Close()
0159 file.Close()
0160