Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2025-03-26 01:51:27

0001 import FWCore.ParameterSet.Config as cms
0002 
0003 from PhysicsTools.NanoAOD.common_cff import *
0004 from PhysicsTools.NanoAOD.simpleCandidateFlatTableProducer_cfi import simpleCandidateFlatTableProducer
0005 from PhysicsTools.NanoAOD.jetsAK8_cff import fatJetTable, subJetTable
0006 from PhysicsTools.NanoAOD.jetsAK4_Puppi_cff import jetPuppiTable
0007 from PhysicsTools.NanoAOD.genparticles_cff import *
0008 # when storing the flat table, always do "T"able in the naming convention
0009 
0010 def addGenCands(process, allPF = False, addAK4=False, addAK8=False):
0011     process.btvGenTask = cms.Task()
0012     process.schedule.associate(process.btvGenTask)
0013 
0014     process.finalGenParticles.select +=[
0015             "keep (4 <= abs(pdgId) <= 5) && statusFlags().isLastCopy()", # BTV: keep b/c quarks in their last copy
0016             "keep (abs(pdgId) == 310 || abs(pdgId) == 3122) && statusFlags().isLastCopy()", # BTV: keep K0s and Lambdas in their last copy
0017             "++keep (400 < abs(pdgId) < 600) || (4000 < abs(pdgId) < 6000)", # keep all B and C hadron + their decay products
0018         ]
0019 
0020     process.btvGenTable =  cms.EDProducer(
0021         "SimpleGenParticleFlatTableProducer",
0022         src = cms.InputTag("finalGenParticles"),
0023         name= cms.string("GenPart"),
0024         doc = cms.string("interesting gen particles "),
0025         singleton=cms.bool(False),
0026         variables=
0027         cms.PSet(
0028             genParticleTable.variables,
0029         vx = Var("vx", "float", doc="x coordinate of vertex position"),
0030         vy = Var("vy", "float", doc="y coordinate of vertex position"),
0031         vz = Var("vz", "float", doc="z coordinate of vertex position"),
0032         genPartIdxMother2 = Var("?numberOfMothers>1?motherRef(1).key():-1", "int", doc="index of the second mother particle, if valid")
0033         )
0034     )
0035     process.btvGenTask.add(process.btvGenTable)
0036     process.genParticleTablesTask.replace(process.genParticleTable,process.btvGenTable)
0037 
0038     if addAK4:
0039         process.btvMCTable = cms.EDProducer("BTVMCFlavourTableProducer",name=jetPuppiTable.name,src=cms.InputTag("linkedObjects","jets"),genparticles=cms.InputTag("prunedGenParticles"))
0040         process.btvGenTask.add(process.btvMCTable)
0041 
0042         process.btvAK4JetExtTable = cms.EDProducer(
0043             "SimplePATJetFlatTableProducer",
0044             src=jetPuppiTable.src,
0045             cut=jetPuppiTable.cut,
0046             name=jetPuppiTable.name,
0047             doc=jetPuppiTable.doc,
0048             singleton=cms.bool(False),  # the number of entries is variable
0049             extension=cms.bool(True),  # this is the extension table for Jets
0050             variables=cms.PSet(
0051             nBHadrons=Var("jetFlavourInfo().getbHadrons().size()",
0052                         int, doc="number of b-hadrons"),
0053             nCHadrons=Var("jetFlavourInfo().getcHadrons().size()",
0054                         int, doc="number of c-hadrons"),
0055             )
0056         )
0057         process.btvGenTask.add(process.btvAK4JetExtTable)
0058 
0059     if addAK8:
0060         process.btvSubJetMCExtTable = cms.EDProducer(
0061             "SimplePATJetFlatTableProducer",
0062             src = subJetTable.src,
0063             cut = subJetTable.cut,
0064             name = subJetTable.name,
0065             doc=subJetTable.doc,
0066             singleton = cms.bool(False),
0067             extension = cms.bool(True),
0068             variables = cms.PSet(
0069                 subGenJetAK8Idx = Var("?genJetFwdRef().backRef().isNonnull()?genJetFwdRef().backRef().key():-1",
0070                 int, doc="index of matched gen Sub jet"),
0071             )
0072         )
0073         process.btvGenTask.add(process.btvSubJetMCExtTable)
0074 
0075     if addAK4:
0076         process.genJetsAK4Constituents = cms.EDProducer("GenJetPackedConstituentPtrSelector",
0077             src = cms.InputTag("slimmedGenJets"),
0078             cut = cms.string("pt > 20")
0079         )
0080         process.btvGenTask.add(process.genJetsAK4Constituents)
0081     if addAK8:
0082         process.genJetsAK8Constituents = cms.EDProducer("GenJetPackedConstituentPtrSelector",
0083             src = cms.InputTag("slimmedGenJetsAK8"),
0084             cut = cms.string("pt > 100.")
0085         )
0086         process.btvGenTask.add(process.genJetsAK8Constituents)
0087 
0088     if allPF:
0089         genCandInput = cms.InputTag("packedGenParticles")
0090     elif not addAK8:
0091         process.genJetsConstituentsTable = cms.EDProducer("PackedGenParticlePtrMerger", src = cms.VInputTag(cms.InputTag("genJetsAK4Constituents", "constituents")), skipNulls = cms.bool(True), warnOnSkip = cms.bool(True))
0092         genCandInput = cms.InputTag("genJetsConstituentsTable")
0093     elif not addAK4:
0094         process.genJetsConstituentsTable = cms.EDProducer("PackedGenParticlePtrMerger", src = cms.VInputTag(cms.InputTag("genJetsAK8Constituents", "constituents")), skipNulls = cms.bool(True), warnOnSkip = cms.bool(True))
0095         genCandInput = cms.InputTag("genJetsConstituentsTable")
0096     else:
0097         process.genJetsConstituentsTable = cms.EDProducer("PackedGenParticlePtrMerger", src = cms.VInputTag(cms.InputTag("genJetsAK4Constituents", "constituents"), cms.InputTag("genJetsAK8Constituents", "constituents")), skipNulls = cms.bool(True), warnOnSkip = cms.bool(True))
0098         genCandInput = cms.InputTag("genJetsConstituentsTable")
0099     if not allPF:
0100         process.btvGenTask.add(process.genJetsConstituentsTable)
0101 
0102     GenCandVars = CandVars.clone()
0103     GenCandVars.pdgId.doc = cms.string("PDG id")
0104 
0105     process.genCandsTable = cms.EDProducer("SimpleCandidateFlatTableProducer",
0106         src = genCandInput,
0107         cut = cms.string(""), #we should not filter after pruning
0108         name= cms.string("GenCands"),
0109         doc = cms.string("Final-state gen particles"),
0110         singleton = cms.bool(False), # the number of entries is variable
0111         extension = cms.bool(False), # this is the main table for the GEN constituents
0112         variables = cms.PSet(GenCandVars)
0113     )
0114     process.btvGenTask.add(process.genCandsTable)
0115 
0116     kwargs = { }
0117     import os
0118     sv_sort = os.getenv('CMSSW_NANOAOD_SV_SORT')
0119     if sv_sort is not None: kwargs['sv_sort'] = cms.untracked.string(sv_sort)
0120     pf_sort = os.getenv('CMSSW_NANOAOD_PF_SORT')
0121     if pf_sort is not None: kwargs['pf_sort'] = cms.untracked.string(pf_sort)
0122 
0123     if addAK4:
0124         process.genAK4ConstituentsTable = cms.EDProducer("GenJetConstituentTableProducer",
0125             candidates = genCandInput,
0126             jets = cms.InputTag("genJetsAK4Constituents"), # Note: The name has "Constituents" in it, but these are the jets
0127             name = cms.string("GenJetCands"),
0128             nameSV = cms.string("GenJetSVs"),
0129             idx_name = cms.string("genCandsIdx"),
0130             idx_nameSV = cms.string("sVIdx"),
0131             readBtag = cms.bool(False),
0132             **kwargs,
0133         )
0134         process.btvGenTask.add(process.genAK4ConstituentsTable)
0135     if addAK8:
0136         process.genAK8ConstituentsTable = cms.EDProducer("GenJetConstituentTableProducer",
0137             candidates = genCandInput,
0138             jets = cms.InputTag("genJetsAK8Constituents"), # Note: The name has "Constituents" in it, but these are the jets
0139             name = cms.string("GenFatJetCands"),
0140             nameSV = cms.string("GenFatJetSVs"),
0141             idx_name = cms.string("genCandsIdx"),
0142             idx_nameSV = cms.string("sVIdx"),
0143             readBtag = cms.bool(False),
0144             **kwargs,
0145         )
0146         process.btvGenTask.add(process.genAK8ConstituentsTable)
0147 
0148     process.genCandMotherTable = cms.EDProducer("GenCandMotherTableProducer",
0149         src     = genCandInput,
0150         mcMap   = cms.InputTag("finalGenParticles"),
0151         objName = cms.string("GenCands"),
0152         branchName = cms.string("genPart"),
0153     )
0154     process.btvGenTask.add(process.genCandMotherTable)
0155 
0156     if allPF:
0157         pfCandInput = cms.InputTag("packedPFCandidates")
0158     else:
0159         pfCandInput = cms.InputTag("finalJetsConstituentsTable")
0160 
0161     process.genCandMCMatchTable = cms.EDProducer("PackedCandMCMatchTableProducer",
0162         src = pfCandInput,
0163         genparticles = genCandInput,
0164         objName = cms.string("PFCands"),
0165         branchName = cms.string("genCand"),
0166         docString = cms.string("MC matching to status==1 genCands"),
0167     )
0168     process.btvGenTask.add(process.genCandMCMatchTable)
0169 
0170     return process