Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-02-14 14:33:59

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