RecoTau/python/DeltaR.py

0001 #ROOTTOOLS
0002 import math
0003 import copy
0004
0005 def deltaR2( e1, p1, e2, p2):
0006     de = e1 - e2
0007     dp = deltaPhi(p1, p2)
0008     return de*de + dp*dp
0009
0010
0011 def deltaR( *args ):
0012     return math.sqrt( deltaR2(*args) )
0013
0014
0015 def deltaPhi( p1, p2):
0016     '''Computes delta phi, handling periodic limit conditions.'''
0017     res = p1 - p2
0018     while res > math.pi:
0019         res -= 2*math.pi
0020     while res < -math.pi:
0021         res += 2*math.pi
0022     return res
0023
0024 def matchObjectCollection3 ( objects, matchCollection, deltaRMax = 0.3, filter = lambda x,y : True ):
0025     '''Univoque association of an element from matchCollection to an element of objects.
0026     Reco and Gen objects get the "matched" attribute, true if they are part of a matched tuple.
0027     By default, the matching is true only if delta R is smaller than 0.3.
0028     '''
0029     #
0030     pairs = {}
0031     if len(objects)==0:
0032         return pairs
0033     if len(matchCollection)==0:
0034         return dict( zip(objects, [None]*len(objects)) )
0035     # build all possible combinations
0036
0037     allPairs = sorted([(deltaR2 (object.eta(), object.phi(), match.eta(), match.phi()), (object, match)) for object in objects for match in matchCollection if list(filter(object,match)) ])
0038     #
0039     # to flag already matched objects
0040     # FIXME this variable remains appended to the object, I do not like it
0041
0042     for object in objects:
0043         object.matched = False
0044     for match in matchCollection:
0045         match.matched = False
0046     #
0047
0048     deltaR2Max = deltaRMax * deltaRMax
0049     for dR2, (object, match) in allPairs:
0050         if dR2 > deltaR2Max:
0051             break
0052         if dR2 < deltaR2Max and object.matched == False and match.matched == False:
0053             object.matched = True
0054             match.matched = True
0055             pairs[object] = match
0056     #
0057
0058     for object in objects:
0059         if object.matched == False:
0060             pairs[object] = None
0061     #
0062
0063     return pairs
0064     # by now, the matched attribute remains in the objects, for future usage
0065     # one could remove it with delattr (object, attrname)
0066
0067
0068 def cleanObjectCollection2( objects, masks, deltaRMin ):
0069     '''Masks objects using a deltaR cut, another algorithm (same results).'''
0070     if len(objects)==0:
0071         return objects
0072     deltaR2Min = deltaRMin*deltaRMin
0073     cleanObjects = copy.copy( objects )
0074     for mask in masks:
0075         tooClose = []
0076         for idx, object in enumerate(cleanObjects):
0077             dR2 = deltaR2( object.eta(), object.phi(),
0078                            mask.eta(), mask.phi() )
0079             if dR2 < deltaR2Min:
0080                 tooClose.append( idx )
0081         nRemoved = 0
0082         for idx in tooClose:
0083             # yes, everytime an object is removed, the list of objects is updated
0084             # so need to update the index accordingly.
0085             # example: to remove : ele 1 and 2
0086             #  first, ele 1 is removed
0087             #  -> ele 2 is now at index 1
0088             # one should again remove the element at index 1
0089             idx -= nRemoved
0090             del cleanObjects[idx]
0091             nRemoved += 1
0092     return cleanObjects
0093
0094
0095 def cleanObjectCollection( objects, masks, deltaRMin ):
0096     '''Masks objects using a deltaR cut.'''
0097     if len(objects)==0 or len(masks)==0:
0098         return objects, []
0099     deltaR2Min = deltaRMin*deltaRMin
0100     cleanObjects = []
0101     dirtyObjects = []
0102     for object in objects:
0103         ok = True
0104         for mask in masks:
0105             dR2 = deltaR2( object.eta(), object.phi(),
0106                            mask.eta(), mask.phi() )
0107             if dR2 < deltaR2Min:
0108                 ok = False
0109         if ok:
0110             cleanObjects.append( object )
0111         else:
0112             dirtyObjects.append( object )
0113     return cleanObjects, dirtyObjects
0114
0115
0116
0117 def bestMatch( object, matchCollection):
0118     '''Return the best match to object in matchCollection, which is the closest object in delta R'''
0119     deltaR2Min = float('+inf')
0120     bm = None
0121     for match in matchCollection:
0122         dR2 = deltaR2( object.eta(), object.phi(),
0123                        match.eta(), match.phi() )
0124         if dR2 < deltaR2Min:
0125             deltaR2Min = dR2
0126             bm = match
0127     return bm, deltaR2Min
0128
0129
0130 def matchObjectCollection( objects, matchCollection, deltaR2Max):
0131     pairs = {}
0132     if len(objects)==0:
0133         return pairs
0134     if len(matchCollection)==0:
0135         return dict( zip(objects, [None]*len(objects)) )
0136     for object in objects:
0137         bm, dr2 = bestMatch( object, matchCollection )
0138         if dr2<deltaR2Max:
0139             pairs[object] = bm
0140         else:
0141             pairs[object] = None
0142     return pairs
0143
0144
0145 def matchObjectCollection2 ( objects, matchCollection, deltaRMax = 0.3 ):
0146     '''Univoque association of an element from matchCollection to an element of objects.
0147     Reco and Gen objects get the "matched" attribute, true if they are part of a matched tuple.
0148     By default, the matching is true only if delta R is smaller than 0.3.
0149     '''
0150
0151     pairs = {}
0152     if len(objects)==0:
0153         return pairs
0154     if len(matchCollection)==0:
0155         return dict( zip(objects, [None]*len(objects)) )
0156     # build all possible combinations
0157     allPairs = sorted([(deltaR2 (object.eta(), object.phi(), match.eta(), match.phi()), (object, match)) for object in objects for match in matchCollection])
0158
0159     # to flag already matched objects
0160     # FIXME this variable remains appended to the object, I do not like it
0161     for object in objects:
0162         object.matched = False
0163     for match in matchCollection:
0164         match.matched = False
0165
0166     deltaR2Max = deltaRMax * deltaRMax
0167     for dR2, (object, match) in allPairs:
0168         if dR2 > deltaR2Max:
0169             break
0170         if dR2 < deltaR2Max and object.matched == False and match.matched == False:
0171             object.matched = True
0172             match.matched = True
0173             pairs[object] = match
0174
0175     for object in objects:
0176         if object.matched == False:
0177             pairs[object] = None
0178
0179     return pairs
0180     # by now, the matched attribute remains in the objects, for future usage
0181     # one could remove it with delattr (object, attrname)
0182
0183
0184
0185 if __name__ == '__main__':
0186
0187     import sys
0188     args = sys.argv[1:]
0189     fargs = map( float, args )
0190
0191     print('dR2 = ', deltaR2( *fargs ))
0192     print('dR = ', deltaR( *fargs ))
0193
0194
0195