python/utils/deltar.py

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