NanoAOD/plugins/LHETablesProducer.cc

0001 #include "FWCore/Framework/interface/global/EDProducer.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0006 #include "FWCore/Utilities/interface/transform.h"
0007 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0008 #include "SimDataFormats/GeneratorProducts/interface/LHEEventProduct.h"
0009 #include "TLorentzVector.h"
0010
0011 #include <vector>
0012 #include <iostream>
0013
0014 class LHETablesProducer : public edm::global::EDProducer<> {
0015 public:
0016   LHETablesProducer(edm::ParameterSet const& params)
0017       : lheTag_(edm::vector_transform(params.getParameter<std::vector<edm::InputTag>>("lheInfo"),
0018                                       [this](const edm::InputTag& tag) { return mayConsume<LHEEventProduct>(tag); })),
0019         precision_(params.getParameter<int>("precision")),
0020         storeLHEParticles_(params.getParameter<bool>("storeLHEParticles")) {
0021     produces<nanoaod::FlatTable>("LHE");
0022     if (storeLHEParticles_)
0023       produces<nanoaod::FlatTable>("LHEPart");
0024   }
0025
0026   ~LHETablesProducer() override {}
0027
0028   void produce(edm::StreamID id, edm::Event& iEvent, const edm::EventSetup& iSetup) const override {
0029     auto lheTab = std::make_unique<nanoaod::FlatTable>(1, "LHE", true);
0030
0031     edm::Handle<LHEEventProduct> lheInfo;
0032     for (const auto& lheTag : lheTag_) {
0033       iEvent.getByToken(lheTag, lheInfo);
0034       if (lheInfo.isValid()) {
0035         break;
0036       }
0037     }
0038     if (lheInfo.isValid()) {
0039       auto lhePartTab = fillLHEObjectTable(*lheInfo, *lheTab);
0040       if (storeLHEParticles_)
0041         iEvent.put(std::move(lhePartTab), "LHEPart");
0042     } else {
0043       if (storeLHEParticles_) {  // need to store a dummy table anyway to make the framework happy
0044         auto lhePartTab = std::make_unique<nanoaod::FlatTable>(1, "LHEPart", true);
0045         iEvent.put(std::move(lhePartTab), "LHEPart");
0046       }
0047     }
0048     iEvent.put(std::move(lheTab), "LHE");
0049   }
0050
0051   std::unique_ptr<nanoaod::FlatTable> fillLHEObjectTable(const LHEEventProduct& lheProd,
0052                                                          nanoaod::FlatTable& out) const {
0053     double lheHT = 0, lheHTIncoming = 0;
0054     unsigned int lheNj = 0, lheNb = 0, lheNc = 0, lheNuds = 0, lheNglu = 0;
0055     double lheVpt = 0;
0056     double alphaS = 0;
0057
0058     const auto& hepeup = lheProd.hepeup();
0059     const auto& pup = hepeup.PUP;
0060     int lep = -1, lepBar = -1, nu = -1, nuBar = -1;
0061     std::vector<float> vals_pt;
0062     std::vector<float> vals_eta;
0063     std::vector<float> vals_phi;
0064     std::vector<float> vals_mass;
0065     std::vector<float> vals_pz;
0066     std::vector<int> vals_pid;
0067     std::vector<int> vals_status;
0068     std::vector<int> vals_spin;
0069     std::vector<int16_t> vals_firstMotherIdx;
0070     std::vector<int16_t> vals_lastMotherIdx;
0071     alphaS = hepeup.AQCDUP;
0072
0073     int nOutPart = 0;
0074     std::vector<int> newIdxs(pup.size());
0075
0076     for (unsigned int i = 0, n = pup.size(); i < n; ++i) {
0077       int status = hepeup.ISTUP[i];
0078       int idabs = std::abs(hepeup.IDUP[i]);
0079
0080       if (status == 1 || status == -1 || (status == 2 && (idabs >= 23 && idabs <= 25))) {
0081         newIdxs[i] = nOutPart;
0082
0083         nOutPart += 1;
0084         TLorentzVector p4(pup[i][0], pup[i][1], pup[i][2], pup[i][3]);  // x,y,z,t
0085         vals_pid.push_back(hepeup.IDUP[i]);
0086         vals_spin.push_back(hepeup.SPINUP[i]);
0087         vals_status.push_back(status);
0088         if (status == -1) {
0089           vals_pt.push_back(0);
0090           vals_eta.push_back(0);
0091           vals_phi.push_back(0);
0092           vals_mass.push_back(0);
0093           vals_pz.push_back(p4.Pz());
0094           vals_firstMotherIdx.push_back(-1);
0095           vals_lastMotherIdx.push_back(-1);
0096         } else {
0097           vals_pt.push_back(p4.Pt());
0098           vals_eta.push_back(p4.Eta());
0099           vals_phi.push_back(p4.Phi());
0100           vals_mass.push_back(p4.M());
0101           vals_pz.push_back(0);
0102           vals_firstMotherIdx.push_back(newIdxs[std::max(hepeup.MOTHUP[i].first - 1, 0)]);
0103           vals_lastMotherIdx.push_back(newIdxs[std::max(hepeup.MOTHUP[i].second - 1, 0)]);
0104         }
0105       } else {
0106         newIdxs[i] = -1;
0107       }
0108       if ((status == 1) && ((idabs == 21) || (idabs > 0 && idabs < 7))) {  //# gluons and quarks
0109         // object counters
0110         lheNj++;
0111         if (idabs == 5)
0112           lheNb++;
0113         else if (idabs == 4)
0114           lheNc++;
0115         else if (idabs <= 3)
0116           lheNuds++;
0117         else if (idabs == 21)
0118           lheNglu++;
0119         // HT
0120         double pt = std::hypot(pup[i][0], pup[i][1]);  // first entry is px, second py
0121         lheHT += pt;
0122         int mothIdx = std::max(
0123             hepeup.MOTHUP[i].first - 1,
0124             0);  //first and last mother as pair; first entry has index 1 in LHE; incoming particles return motherindex 0
0125         int mothIdxTwo = std::max(hepeup.MOTHUP[i].second - 1, 0);
0126         int mothStatus = hepeup.ISTUP[mothIdx];
0127         int mothStatusTwo = hepeup.ISTUP[mothIdxTwo];
0128         bool hasIncomingAsMother = mothStatus < 0 || mothStatusTwo < 0;
0129         if (hasIncomingAsMother)
0130           lheHTIncoming += pt;
0131       } else if (idabs == 12 || idabs == 14 || idabs == 16) {
0132         (hepeup.IDUP[i] > 0 ? nu : nuBar) = i;
0133       } else if (idabs == 11 || idabs == 13 || idabs == 15) {
0134         (hepeup.IDUP[i] > 0 ? lep : lepBar) = i;
0135       }
0136     }
0137
0138     std::pair<int, int> v(0, 0);
0139     if (lep != -1 && lepBar != -1)
0140       v = std::make_pair(lep, lepBar);
0141     else if (lep != -1 && nuBar != -1)
0142       v = std::make_pair(lep, nuBar);
0143     else if (nu != -1 && lepBar != -1)
0144       v = std::make_pair(nu, lepBar);
0145     else if (nu != -1 && nuBar != -1)
0146       v = std::make_pair(nu, nuBar);
0147     if (v.first != -1 && v.second != -1) {
0148       lheVpt = std::hypot(pup[v.first][0] + pup[v.second][0], pup[v.first][1] + pup[v.second][1]);
0149     }
0150
0151     out.addColumnValue<uint8_t>("Njets", lheNj, "Number of jets (partons) at LHE step");
0152     out.addColumnValue<uint8_t>("Nb", lheNb, "Number of b partons at LHE step");
0153     out.addColumnValue<uint8_t>("Nc", lheNc, "Number of c partons at LHE step");
0154     out.addColumnValue<uint8_t>("Nuds", lheNuds, "Number of u,d,s partons at LHE step");
0155     out.addColumnValue<uint8_t>("Nglu", lheNglu, "Number of gluon partons at LHE step");
0156     out.addColumnValue<float>("HT", lheHT, "HT, scalar sum of parton pTs at LHE step");
0157     out.addColumnValue<float>(
0158         "HTIncoming", lheHTIncoming, "HT, scalar sum of parton pTs at LHE step, restricted to partons");
0159     out.addColumnValue<float>("Vpt", lheVpt, "pT of the W or Z boson at LHE step");
0160     out.addColumnValue<uint8_t>("NpNLO", lheProd.npNLO(), "number of partons at NLO");
0161     out.addColumnValue<uint8_t>("NpLO", lheProd.npLO(), "number of partons at LO");
0162     out.addColumnValue<float>("AlphaS", alphaS, "Per-event alphaS");
0163
0164     auto outPart = std::make_unique<nanoaod::FlatTable>(vals_pt.size(), "LHEPart", false);
0165     outPart->addColumn<float>("pt", vals_pt, "Pt of LHE particles", this->precision_);
0166     outPart->addColumn<float>("eta", vals_eta, "Pseodorapidity of LHE particles", this->precision_);
0167     outPart->addColumn<float>("phi", vals_phi, "Phi of LHE particles", this->precision_);
0168     outPart->addColumn<float>("mass", vals_mass, "Mass of LHE particles", this->precision_);
0169     outPart->addColumn<float>("incomingpz", vals_pz, "Pz of incoming LHE particles", this->precision_);
0170     outPart->addColumn<int>("pdgId", vals_pid, "PDG ID of LHE particles");
0171     outPart->addColumn<int>("status", vals_status, "LHE particle status; -1:incoming, 1:outgoing");
0172     outPart->addColumn<int>("spin", vals_spin, "Spin of LHE particles");
0173     outPart->addColumn<int16_t>(
0174         "firstMotherIdx", vals_firstMotherIdx, "Index of this particle's first mother in the LHEPart collection");
0175     outPart->addColumn<int16_t>(
0176         "lastMotherIdx", vals_lastMotherIdx, "Index of this particle's last mother in the LHEPart collection");
0177
0178     return outPart;
0179   }
0180
0181   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0182     edm::ParameterSetDescription desc;
0183     desc.add<std::vector<edm::InputTag>>("lheInfo", std::vector<edm::InputTag>{{"externalLHEProducer"}, {"source"}})
0184         ->setComment("tag(s) for the LHE information (LHEEventProduct)");
0185     desc.add<int>("precision", -1)->setComment("precision on the 4-momenta of the LHE particles");
0186     desc.add<bool>("storeLHEParticles", false)
0187         ->setComment("Whether we want to store the 4-momenta of the status 1 particles at LHE level");
0188     descriptions.add("lheInfoTable", desc);
0189   }
0190
0191 protected:
0192   const std::vector<edm::EDGetTokenT<LHEEventProduct>> lheTag_;
0193   const unsigned int precision_;
0194   const bool storeLHEParticles_;
0195 };
0196
0197 #include "FWCore/Framework/interface/MakerMacros.h"
0198 DEFINE_FWK_MODULE(LHETablesProducer);