analyzers/core/PileUpAnalyzer.py

0001 import os
0002 from PhysicsTools.Heppy.analyzers.core.VertexHistograms import VertexHistograms
0003 from PhysicsTools.Heppy.analyzers.core.Analyzer import Analyzer
0004 from PhysicsTools.Heppy.analyzers.core.AutoHandle import AutoHandle
0005 from PhysicsTools.HeppyCore.statistics.average import Average
0006 from PhysicsTools.Heppy.physicsutils.PileUpSummaryInfo import PileUpSummaryInfo
0007 import PhysicsTools.HeppyCore.framework.config as cfg
0008
0009 from ROOT import TFile, TH1F
0010
0011 class PileUpAnalyzer( Analyzer ):
0012     '''Computes pile-up weights for MC from the pile up histograms for MC and data.
0013     These histograms should be set on the components as
0014     puFileData, puFileMC attributes, as is done here:
0015
0016     http://cmssw.cvs.cern.ch/cgi-bin/cmssw.cgi/UserCode/CMG/CMGTools/H2TauTau/Colin/test_tauMu_2012_cfg.py?view=markup
0017
0018     THESE HISTOGRAMS MUST BE CONSISTENT, SEE
0019     https://twiki.cern.ch/twiki/bin/view/CMS/CMGToolsPileUpReweighting#Generating_pile_up_distributions
0020
0021     If the component is not MC, or if the puFileData and puFileMC are not
0022     set for the component, the reweighting is not done.
0023
0024     The analyzer sets event.vertexWeight.
0025     This weight is multiplied to the global event weight, event.eventWeight.
0026     When using this analyzer, make sure that the VertexAnalyzer is disabled,
0027     as you would be reweighting the MC PU distribution twice!
0028
0029     Additionally, this analyzer writes in the output an histogram containing the unweighting MC
0030     pile-up distribution, to be used in input of the weighting for a later pass.
0031
0032     Example of use:
0033
0034     puAna = cfg.Analyzer(
0035       "PileUpAnalyzer",
0036       # build unweighted pu distribution using number of pile up interactions if False
0037       # otherwise, use fill the distribution using number of true interactions
0038       true = True
0039       )
0040     '''
0041
0042     def __init__(self, cfg_ana, cfg_comp, looperName):
0043         super(PileUpAnalyzer, self).__init__(cfg_ana, cfg_comp, looperName)
0044
0045         self.doHists=True
0046
0047         if (hasattr(self.cfg_ana,'makeHists')) and (not self.cfg_ana.makeHists):
0048             self.doHists=False
0049
0050         self.allVertices = self.cfg_ana.allVertices if (hasattr(self.cfg_ana,'allVertices')) else "_AUTO_"
0051
0052         if self.cfg_comp.isMC and self.doHists:
0053             self.rawmcpileup = VertexHistograms('/'.join([self.dirName,
0054                                                           'rawMCPU.root']))
0055         self.enable = True
0056         ## if component is embed return (has no trigger obj)
0057         if self.cfg_comp.isEmbed :
0058           self.cfg_comp.puFileMC   = None
0059           self.cfg_comp.puFileData = None
0060
0061         if self.cfg_comp.isMC or self.cfg_comp.isEmbed:
0062             if not hasattr(self.cfg_comp,"puFileMC") or (self.cfg_comp.puFileMC is None and self.cfg_comp.puFileData is None):
0063                 self.enable = False
0064             else:
0065                 assert( os.path.isfile(os.path.expandvars(self.cfg_comp.puFileMC)) )
0066                 assert( os.path.isfile(os.path.expandvars(self.cfg_comp.puFileData)) )
0067
0068                 self.mcfile = TFile( self.cfg_comp.puFileMC )
0069                 self.mchist = self.mcfile.Get('pileup')
0070                 self.mchist.Scale( 1 / self.mchist.Integral() )
0071
0072                 self.datafile = TFile( self.cfg_comp.puFileData )
0073                 self.datahist = self.datafile.Get('pileup')
0074                 self.datahist.Scale( 1 / self.datahist.Integral() )
0075                 # import pdb; pdb.set_trace()
0076                 if self.mchist.GetNbinsX() != self.datahist.GetNbinsX():
0077                     raise ValueError('data and mc histograms must have the same number of bins')
0078                 if self.mchist.GetXaxis().GetXmin() != self.datahist.GetXaxis().GetXmin():
0079                     raise ValueError('data and mc histograms must have the same xmin')
0080                 if self.mchist.GetXaxis().GetXmax() != self.datahist.GetXaxis().GetXmax():
0081                     raise ValueError('data and mc histograms must have the same xmax')
0082
0083     def declareHandles(self):
0084         super(PileUpAnalyzer, self).declareHandles()
0085         self.mchandles['pusi'] =  AutoHandle(
0086             'slimmedAddPileupInfo',
0087             'std::vector<PileupSummaryInfo>',
0088             fallbackLabel="addPileupInfo"
0089             )
0090
0091         if self.allVertices == '_AUTO_':
0092             self.handles['vertices'] =  AutoHandle( "offlineSlimmedPrimaryVertices", 'std::vector<reco::Vertex>', fallbackLabel="offlinePrimaryVertices" )
0093         else:
0094             self.handles['vertices'] =  AutoHandle( self.allVertices, 'std::vector<reco::Vertex>' )
0095
0096     def beginLoop(self, setup):
0097         super(PileUpAnalyzer,self).beginLoop(setup)
0098         self.averages.add('puWeight', Average('puWeight') )
0099
0100
0101     def process(self, event):
0102         self.readCollections( event.input )
0103         ## if component is embed return (has no trigger obj)
0104         if self.cfg_comp.isEmbed :
0105           return True
0106
0107         event.puWeight = 1
0108         event.nPU = None
0109         event.pileUpVertex_z = []
0110         event.pileUpVertex_ptHat = []
0111         if self.cfg_comp.isMC:
0112             event.pileUpInfo = map( PileUpSummaryInfo,
0113                                     self.mchandles['pusi'].product() )
0114             for puInfo in event.pileUpInfo:
0115                 if puInfo.getBunchCrossing()==0:
0116                     # import pdb; pdb.set_trace()
0117                     if self.cfg_ana.true is False:
0118                         event.nPU = puInfo.nPU()
0119                     else:
0120                         event.nPU = puInfo.nTrueInteractions()
0121
0122                     if self.doHists:
0123                         self.rawmcpileup.hist.Fill( event.nPU )
0124
0125                     ##get z position of on-time pile-up sorted by pt-hat
0126                     ptHat_zPositions = zip(puInfo.getPU_pT_hats(),puInfo.getPU_zpositions())
0127                     ptHat_zPositions.sort(reverse=True)
0128                     for ptHat_zPosition in ptHat_zPositions:
0129                         event.pileUpVertex_z.append(ptHat_zPosition[1])
0130                         event.pileUpVertex_ptHat.append(ptHat_zPosition[0])
0131
0132             if event.nPU is None:
0133                 raise ValueError('nPU cannot be None! means that no pu info has been found for bunch crossing 0.')
0134         elif self.cfg_comp.isEmbed:
0135             vertices = self.handles['vertices'].product()
0136             event.nPU = len(vertices)
0137         else:
0138             return True
0139
0140         if self.enable:
0141             bin = self.datahist.FindBin(event.nPU)
0142             if bin<1 or bin>self.datahist.GetNbinsX():
0143                 event.puWeight = 0
0144             else:
0145                 data = self.datahist.GetBinContent(bin)
0146                 mc = self.mchist.GetBinContent(bin)
0147                 #Protect 0 division!!!!
0148                 if mc !=0.0:
0149                     event.puWeight = data/mc
0150                 else:
0151                     event.puWeight = 1
0152
0153         event.eventWeight *= event.puWeight
0154         self.averages['puWeight'].add( event.puWeight )
0155         return True
0156
0157     def write(self, setup):
0158         super(PileUpAnalyzer, self).write(setup)
0159         if self.cfg_comp.isMC and self.doHists:
0160             self.rawmcpileup.write()
0161
0162
0163 setattr(PileUpAnalyzer,"defaultConfig", cfg.Analyzer(
0164     class_object = PileUpAnalyzer,
0165     true = True,  # use number of true interactions for reweighting
0166     makeHists=False
0167 )
0168 )
0169