def test_cwn(self):
fi = TFile('tree2.root','RECREATE')
tr = Tree('test_tree', 'A test tree')
tr.var('nvals', the_type=int)
tr.vector('x', 'nvals', 20)
tr.fill('nvals', 10)
tr.vfill('x', range(10))
tr.tree.Fill()
tr.reset()
tr.fill('nvals', 5)
tr.vfill('x', range(5))
tr.tree.Fill()
fi.Write()
fi.Close()
class TreeProducer(Analyzer):
def beginLoop(self, setup):
super(TreeProducer, self).beginLoop(setup)
self.rootfile = TFile('/'.join([self.dirName, 'tree.root']),
'recreate')
self.tree = Tree('events', '')
self.tree.var('weight', float)
self.tree.var('Jet1_tau1', float)
self.tree.var('Jet1_tau2', float)
self.tree.var('Jet1_tau3', float)
self.tree.var('Jet2_tau1', float)
self.tree.var('Jet2_tau2', float)
self.tree.var('Jet2_tau3', float)
self.tree.var('Jet1_tau32', float)
self.tree.var('Jet1_tau31', float)
self.tree.var('Jet1_tau21', float)
self.tree.var('Jet2_tau32', float)
self.tree.var('Jet2_tau31', float)
self.tree.var('Jet2_tau21', float)
bookParticle(self.tree, 'Jet1')
bookParticle(self.tree, 'Jet2')
bookParticle(self.tree, 'softDroppedJet1')
bookParticle(self.tree, 'softDroppedJet2')
bookParticle(self.tree, 'trimmedJet1')
bookParticle(self.tree, 'trimmedJet2')
bookParticle(self.tree, 'prunedJet1')
bookParticle(self.tree, 'prunedJet2')
self.tree.var('RSGravitonReconstructedMass', float)
self.tree.var('RSGravitonReconstructedMass_trimmed', float)
self.tree.var('RSGravitonReconstructedMass_softDropped', float)
self.tree.var('RSGravitonReconstructedMass_pruned', float)
self.tree.var('Jet1_Flow15', float)
self.tree.var('Jet1_Flow25', float)
self.tree.var('Jet1_Flow35', float)
self.tree.var('Jet1_Flow45', float)
self.tree.var('Jet1_Flow55', float)
self.tree.var('Jet2_Flow15', float)
self.tree.var('Jet2_Flow25', float)
self.tree.var('Jet2_Flow35', float)
self.tree.var('Jet2_Flow45', float)
self.tree.var('Jet2_Flow55', float)
self.tree.vector('Jet1_jetConstituent_pt',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet1_jetConstituent_eta',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet1_jetConstituent_phi',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet1_jetConstituent_e', '{}'.format(maxConstituents),
maxConstituents, float)
self.tree.vector('Jet2_jetConstituent_pt',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet2_jetConstituent_eta',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet2_jetConstituent_phi',
'{}'.format(maxConstituents), maxConstituents, float)
self.tree.vector('Jet2_jetConstituent_e', '{}'.format(maxConstituents),
maxConstituents, float)
self.tree.var('rapiditySeparation', float)
self.tree.var('transverseMomentumAsymmetry', float)
def process(self, event):
self.tree.reset()
jets = getattr(event, self.cfg_ana.fatjets)
if (len(jets) > 1):
self.tree.fill('weight', event.weight)
self.tree.fill('rapiditySeparation',
abs(jets[0].eta() - jets[1].eta()))
self.tree.fill('transverseMomentumAsymmetry',
(jets[0].pt() - jets[1].pt()) /
(jets[0].pt() + jets[1].pt()))
self.tree.fill('Jet1_tau1', jets[0].tau1)
self.tree.fill('Jet1_tau2', jets[0].tau2)
self.tree.fill('Jet1_tau3', jets[0].tau3)
self.tree.fill('Jet2_tau1', jets[1].tau1)
self.tree.fill('Jet2_tau2', jets[1].tau2)
self.tree.fill('Jet2_tau3', jets[1].tau3)
if (jets[0].tau1 != 0.0):
self.tree.fill('Jet1_tau31', jets[0].tau3 / jets[0].tau1)
self.tree.fill('Jet1_tau21', jets[0].tau2 / jets[0].tau1)
else:
self.tree.fill('Jet1_tau31', -99)
self.tree.fill('Jet1_tau21', -99)
if (jets[0].tau2 != 0.0):
self.tree.fill('Jet1_tau32', jets[0].tau3 / jets[0].tau2)
else:
self.tree.fill('Jet1_tau32', -99)
if (jets[1].tau1 != 0.0):
self.tree.fill('Jet2_tau31', jets[1].tau3 / jets[1].tau1)
self.tree.fill('Jet2_tau21', jets[1].tau2 / jets[1].tau1)
else:
self.tree.fill('Jet2_tau31', -99)
self.tree.fill('Jet2_tau21', -99)
if (jets[1].tau2 != 0.0):
self.tree.fill('Jet2_tau32', jets[1].tau3 / jets[1].tau2)
else:
self.tree.fill('Jet2_tau32', -99)
fillParticle(self.tree, 'Jet1', jets[0])
fillParticle(self.tree, 'Jet2', jets[1])
fillParticle(self.tree, 'softDroppedJet1',
jets[0].subjetsSoftDrop[0])
fillParticle(self.tree, 'softDroppedJet2',
jets[1].subjetsSoftDrop[0])
fillParticle(self.tree, 'trimmedJet1', jets[0].subjetsTrimming[0])
fillParticle(self.tree, 'trimmedJet2', jets[1].subjetsTrimming[0])
fillParticle(self.tree, 'prunedJet1', jets[0].subjetsPruning[0])
fillParticle(self.tree, 'prunedJet2', jets[1].subjetsPruning[0])
jet1_ungroomed = ROOT.TLorentzVector()
jet2_ungroomed = ROOT.TLorentzVector()
jet1_ungroomed.SetPtEtaPhiE(jets[0].pt(), jets[0].eta(),
jets[0].phi(), jets[0].e())
jet2_ungroomed.SetPtEtaPhiE(jets[1].pt(), jets[1].eta(),
jets[1].phi(), jets[1].e())
self.tree.fill('RSGravitonReconstructedMass',
(jet1_ungroomed + jet2_ungroomed).M())
jet1_trimmed = ROOT.TLorentzVector()
jet2_trimmed = ROOT.TLorentzVector()
jet1_trimmed.SetPtEtaPhiE(jets[0].subjetsTrimming[0].pt(),
jets[0].subjetsTrimming[0].eta(),
jets[0].subjetsTrimming[0].phi(),
jets[0].subjetsTrimming[0].e())
jet2_trimmed.SetPtEtaPhiE(jets[1].subjetsTrimming[0].pt(),
jets[1].subjetsTrimming[0].eta(),
jets[1].subjetsTrimming[0].phi(),
jets[1].subjetsTrimming[0].e())
self.tree.fill('RSGravitonReconstructedMass_trimmed',
(jet1_trimmed + jet2_trimmed).M())
jet1_pruned = ROOT.TLorentzVector()
jet2_pruned = ROOT.TLorentzVector()
jet1_pruned.SetPtEtaPhiE(jets[0].subjetsPruning[0].pt(),
jets[0].subjetsPruning[0].eta(),
jets[0].subjetsPruning[0].phi(),
jets[0].subjetsPruning[0].e())
jet2_pruned.SetPtEtaPhiE(jets[1].subjetsPruning[0].pt(),
jets[1].subjetsPruning[0].eta(),
jets[1].subjetsPruning[0].phi(),
jets[1].subjetsPruning[0].e())
self.tree.fill('RSGravitonReconstructedMass_pruned',
(jet1_pruned + jet2_pruned).M())
jet1_softDropped = ROOT.TLorentzVector()
jet2_softDropped = ROOT.TLorentzVector()
jet1_softDropped.SetPtEtaPhiE(jets[0].subjetsSoftDrop[0].pt(),
jets[0].subjetsSoftDrop[0].eta(),
jets[0].subjetsSoftDrop[0].phi(),
jets[0].subjetsSoftDrop[0].e())
jet2_softDropped.SetPtEtaPhiE(jets[1].subjetsSoftDrop[0].pt(),
jets[1].subjetsSoftDrop[0].eta(),
jets[1].subjetsSoftDrop[0].phi(),
jets[1].subjetsSoftDrop[0].e())
self.tree.fill('RSGravitonReconstructedMass_softDropped',
(jet1_softDropped + jet2_softDropped).M())
#############################################################################
Jet1_jc_pt = []
Jet1_jc_eta = []
Jet1_jc_phi = []
Jet1_jc_e = []
Jet2_jc_pt = []
Jet2_jc_eta = []
Jet2_jc_phi = []
Jet2_jc_e = []
for constituent in jets[0].jetConstituents[1:]:
Jet1_jc_pt.append(constituent.pt())
Jet1_jc_eta.append(constituent.eta())
Jet1_jc_phi.append(constituent.phi())
Jet1_jc_e.append(constituent.e())
for constituent in jets[1].jetConstituents[1:]:
Jet2_jc_pt.append(constituent.pt())
Jet2_jc_eta.append(constituent.eta())
Jet2_jc_phi.append(constituent.phi())
Jet2_jc_e.append(constituent.e())
self.tree.vfill("Jet1_jetConstituent_pt", Jet1_jc_pt)
self.tree.vfill("Jet1_jetConstituent_eta", Jet1_jc_eta)
self.tree.vfill("Jet1_jetConstituent_phi", Jet1_jc_phi)
self.tree.vfill("Jet1_jetConstituent_e", Jet1_jc_e)
self.tree.vfill("Jet2_jetConstituent_pt", Jet2_jc_pt)
self.tree.vfill("Jet2_jetConstituent_eta", Jet2_jc_eta)
self.tree.vfill("Jet2_jetConstituent_phi", Jet2_jc_phi)
self.tree.vfill("Jet2_jetConstituent_e", Jet2_jc_e)
#############################################################################
#Flow n,5
#############################################################################
#REQUIRES THE FOLLOWING IN heppy/analyzers/fcc/Reader.py AFTER LINE 151:
# particle_relations = defaultdict(list)
# for tjet in store.get(self.cfg_ana.fatjets):
# for i in range(tjet.particles_size()):
# particle_relations[Jet(tjet)].append(Particle(tjet.particles(i)))
# for fatjet, particles in particle_relations.items():
# fatjets[fatjet].jetConstituents = particles
#############################################################################
R = 0.8
flow_Jet1 = [0] * 5
flow_Jet2 = [0] * 5
constituent_vector = ROOT.TLorentzVector()
for n in range(1, 5 + 1):
for constituent in jets[0].jetConstituents[1:]:
constituent_vector.SetPtEtaPhiE(constituent.pt(),
constituent.eta(),
constituent.phi(),
constituent.e())
dR = jet1_ungroomed.DeltaR(constituent_vector)
if ((dR >= (n - 1) / 5 * R) and (dR < n / 5 * R)):
flow_Jet1[n - 1] += abs(constituent.pt()) / abs(
jets[0].pt())
for constituent in jets[1].jetConstituents[1:]:
constituent_vector.SetPtEtaPhiE(constituent.pt(),
constituent.eta(),
constituent.phi(),
constituent.e())
dR = jet2_ungroomed.DeltaR(constituent_vector)
if ((dR >= (n - 1) / 5 * R) and (dR < n / 5 * R)):
flow_Jet2[n - 1] += abs(constituent.pt()) / abs(
jets[1].pt())
self.tree.fill('Jet1_Flow15', flow_Jet1[0])
self.tree.fill('Jet2_Flow15', flow_Jet2[0])
self.tree.fill('Jet1_Flow25', flow_Jet1[1])
self.tree.fill('Jet2_Flow25', flow_Jet2[1])
self.tree.fill('Jet1_Flow35', flow_Jet1[2])
self.tree.fill('Jet2_Flow35', flow_Jet2[2])
self.tree.fill('Jet1_Flow45', flow_Jet1[3])
self.tree.fill('Jet2_Flow45', flow_Jet2[3])
self.tree.fill('Jet1_Flow55', flow_Jet1[4])
self.tree.fill('Jet2_Flow55', flow_Jet2[4])
self.tree.tree.Fill()
def write(self, setup):
self.rootfile.Write()
self.rootfile.Close()
