python/analyzers/Matcher.py

0001 from heppy.framework.analyzer import Analyzer
0002 from heppy.utils.deltar import matchObjectCollection, deltaR
0003
0004 import collections
0005
0006 class Matcher(Analyzer):
0007     '''Particle matcher.
0008
0009     Works with any kind of object with a p4 function.
0010
0011     Simple example configuration:
0012
0013     from heppy_fcc.analyzers.Matcher import Matcher
0014     papas_jet_match = cfg.Analyzer(
0015       Matcher,
0016       instance_label = 'papas',
0017       match_particles = 'gen_jets',
0018       particles = 'papas_jets'
0019     )
0020
0021     particles: Name of the collection containing the particles to be matched.
0022     match_particles: Name of the collection containing the particles where a match
0023                is to be found.
0024
0025     In this particular case, each jet in "papas_jets" will end up with a new
0026     attribute called "match". This attribute can be either the closest gen jet in the
0027     "gen_jets" collection in case a gen_jet is found within delta R = 0.3,
0028     or None in case a match cannot be found in this cone.
0029
0030     More complex example configuration:
0031
0032     papas_particle_match_g2r = cfg.Analyzer(
0033       Matcher,
0034       instance_label = 'papas_g2r',
0035       particles = 'gen_particles_stable',
0036       match_particles = [
0037         ('papas_rec_particles', None),
0038         ('papas_rec_particles', 211),
0039         ('papas_rec_particles', 130),
0040         ('papas_rec_particles', 22)
0041       ]
0042       )
0043
0044     In this case, each gen particle in gen_particles_stable will end up with the following
0045     new attributes:
0046       - "match"    : closest reconstructed particle in "papas_rec_particles", if any.
0047       - "match_211": closest reconstructed particle of pdgId 211 in "papas_rec_particles",
0048                      if any.
0049       - etc.
0050
0051     '''
0052
0053
0054     def beginLoop(self, setup):
0055         super(Matcher, self).beginLoop(setup)
0056         self.match_collections = []
0057         if isinstance( self.cfg_ana.match_particles, str):
0058             self.match_collections.append( (self.cfg_ana.match_particles, None) )
0059         else:
0060             self.match_collections = self.cfg_ana.match_particles
0061
0062     def process(self, event):
0063         particles = getattr(event, self.cfg_ana.particles)
0064         # match_particles = getattr(event, self.cfg_ana.match_particles)
0065         for collname, pdgid in self.match_collections:
0066             match_ptcs = getattr(event, collname)
0067             match_ptcs_filtered = match_ptcs
0068             if pdgid is not None:
0069                 match_ptcs_filtered = [ptc for ptc in match_ptcs
0070                                        if ptc.pdgid()==pdgid]
0071             pairs = matchObjectCollection(particles, match_ptcs_filtered,
0072                                           0.3**2)
0073             for ptc in particles:
0074                 matchname = 'match'
0075                 if pdgid:
0076                     matchname = 'match_{pdgid}'.format(pdgid=pdgid)
0077                 match = pairs[ptc]
0078                 setattr(ptc, matchname, match)
0079                 if match:
0080                     drname = 'dr'
0081                     if pdgid:
0082                         drname = 'dr_{pdgid}'.format(pdgid=pdgid)
0083                     dr = deltaR(ptc.theta(), ptc.phi(),
0084                                 match.theta(), match.phi())
0085                     setattr(ptc, drname, dr)
0086                     # print dr, ptc, match