RecoTau/python/DeltaR.py

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