-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathutils.py
100 lines (80 loc) · 3.97 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import math
import bisect
import ROOT
from ROOT import RooArgSet
muMass = 0.1057
def deltaPhi(phi1, phi2):
result = phi1 - phi2;
if result > math.pi: result -= 2*math.pi
if result <= -math.pi: result += 2*math.pi
return result
def deltaR(v1, v2):
return math.sqrt(deltaPhi(v1.phi(), v2.phi())**2 + (v1.eta()-v2.eta())**2)
def deltaR(phi1, eta1, phi2, eta2):
return math.sqrt(deltaPhi(phi1, phi2)**2 + (eta1-eta2)**2)
class Properties:
"""Stores the min and max mass and the binning in pt"""
def __init__(self, minMass, maxMass, ptBinsX, ptBinsY, triggerMatchDeltaR, NoBkgd, minDeltaR):
self.minMass = minMass
self.maxMass = maxMass
self.ptBinsX = ptBinsX
self.ptBinsY = ptBinsY
self.triggerMatchDeltaR = triggerMatchDeltaR
self.NoBkgd = NoBkgd
self.minDeltaR = minDeltaR
def computeCosineAndMass(mu1, mu2):
muon1 = ROOT.TLorentzVector()
muon2 = ROOT.TLorentzVector()
muon1.SetPtEtaPhiM(mu1.pt, mu1.eta, mu1.phi, muMass)
muon2.SetPtEtaPhiM(mu2.pt, mu2.eta, mu2.phi, muMass)
cosine = math.cos(muon1.Angle(muon2.Vect()))
mass = (muon1+muon2).M()
return (cosine, mass)
def passDileptonSelection(track1, track2, cosine):
if track1.charge == track2.charge: return False
if deltaR(track1.phi, track1.eta, track2.phi, track2.eta) < 0.2: return False
if cosine <= -0.79: return False
return True
def fillSingleCandidate(mass, p, track1, track2, histoMap, datasetMap):
cosineAndMass = computeCosineAndMass(track1, track2)
if not passDileptonSelection(track1, track2, cosineAndMass[0]): return False
if cosineAndMass[1] < p.minMass or cosineAndMass[1] > p.maxMass: return False
mass.setVal(cosineAndMass[1])
bins = find_bins(track1.pt, track2.pt, p.ptBinsX, p.ptBinsY)
if bins[0] == -1 or bins[1] == -1: return False
if track1.charge == track2.charge: return False
histoMap[bins].Fill(mass.getVal())
#print bins, cosineAndMass[1], track1.pt, track2.pt
datasetMap[bins].add(RooArgSet(mass))
return True
def fillCandidates_tnp(mass, properties, matchedTags, probes, histoMap, datasetMap):
for tag in matchedTags:
for probe in probes:
if deltaR(tag.phi,tag.eta, probe.phi, probe.eta) > properties.minDeltaR :
fillSingleCandidate(mass, properties, tag, probe, histoMap, datasetMap)
def find_bins(pt1, pt2, ptBinsX, ptBinsY):
return (bisect.bisect_right(ptBinsX, pt1)-1, bisect.bisect_right(ptBinsY, pt2)-1)
def find_position(bin1, bin2, ptBinsX):
return bin1+len(ptBinsX)*bin2
def find_position_NoOverflow(bin1, bin2, ptBinsX):
return bin1+(len(ptBinsX)-1)*bin2
def passSelectionGlobalMuon(track):
if ( track.pt > 26 and abs(track.eta) < 2) and track.isolation/track.pt < 0.1 \
and track.trackerLayersWithMeasurement >= 6 and track.dxy < 30. and track.dz < 30.:
return True
def passSelection(track):
if ( track.pt > 26 and abs(track.eta) < 2) and track.isolation/track.pt < 0.1 \
and track.trackerLayersWithMeasurement >= 6 and track.dxy < 30. and track.dz < 30 and track.trackQuality:
return True
def fillTriggerMatchedGlobalMuon(track, triggerObjects, matchedTracks, p):
for triggerMuon in triggerObjects:
if (deltaR(triggerMuon.phi, triggerMuon.eta, track.phi, track.eta) < p.triggerMatchDeltaR) and passSelectionGlobalMuon(track):
matchedTracks.append(track)
def fillTriggerMatchedTrack(track, triggerObjects, matchedTracks, p):
for triggerMuon in triggerObjects:
if (deltaR(triggerMuon.phi, triggerMuon.eta, track.phi, track.eta) < p.triggerMatchDeltaR) and passSelection(track):
matchedTracks.append(track)
def buildName(baseName, ptBin1, ptBin2, ptBinsX, ptBinsY):
return baseName+str(ptBinsX[ptBin1])+"_"+str(ptBinsY[ptBin2])
def buildNamePars(baseName, ptBin1, ptBin12, ptBin2, ptBin22, ptBinsX, ptBinsY):
return baseName+str(ptBinsX[ptBin1])+"_"+str(ptBinsX[ptBin12])+"_"+str(ptBinsY[ptBin2])+"_"+str(ptBinsY[ptBin22])