LHETablesProducer.cc

Go to the documentation of this file.

#include "FWCore/Framework/interface/global/EDProducer.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/Utilities/interface/transform.h"
#include "DataFormats/NanoAOD/interface/FlatTable.h"
#include "SimDataFormats/GeneratorProducts/interface/LHEEventProduct.h"
#include "TLorentzVector.h"

#include <vector>
#include <iostream>


class LHETablesProducer : public edm::global::EDProducer<> {
    public:
        LHETablesProducer( edm::ParameterSet const & params ) :
            lheTag_(edm::vector_transform(params.getParameter<std::vector<edm::InputTag>>("lheInfo"), [this](const edm::InputTag & tag) { return mayConsume<LHEEventProduct>(tag); })),
            precision_( params.getParameter<int>("precision") ),
            storeLHEParticles_( params.getParameter<bool>("storeLHEParticles") )
        {
            produces<nanoaod::FlatTable>("LHE");
            if (storeLHEParticles_)
              produces<nanoaod::FlatTable>("LHEPart");

        }

        ~LHETablesProducer() override {}

        void produce(edm::StreamID id, edm::Event& iEvent, const edm::EventSetup& iSetup) const override {
            auto lheTab  = std::make_unique<nanoaod::FlatTable>(1, "LHE", true);

            edm::Handle<LHEEventProduct> lheInfo;
            for (const auto & lheTag: lheTag_) {
              iEvent.getByToken(lheTag, lheInfo);
              if (lheInfo.isValid()) {
                break;
              }
            }
            if (lheInfo.isValid()) {
              auto lhePartTab = fillLHEObjectTable(*lheInfo, *lheTab);
              if (storeLHEParticles_) iEvent.put(std::move(lhePartTab), "LHEPart");
            } else {
              if (storeLHEParticles_) { // need to store a dummy table anyway to make the framework happy
                  auto lhePartTab = std::make_unique<nanoaod::FlatTable>(1, "LHEPart", true);
                  iEvent.put(std::move(lhePartTab), "LHEPart");
              }
            }
            iEvent.put(std::move(lheTab), "LHE");

        }

        std::unique_ptr<nanoaod::FlatTable> fillLHEObjectTable(const LHEEventProduct & lheProd, nanoaod::FlatTable & out) const {
            double lheHT = 0, lheHTIncoming = 0;
            unsigned int lheNj = 0, lheNb = 0, lheNc = 0, lheNuds = 0, lheNglu = 0;
            double lheVpt = 0;

            const auto & hepeup = lheProd.hepeup();
            const auto & pup = hepeup.PUP;
            int lep = -1, lepBar = -1, nu = -1, nuBar = -1;
            std::vector<float> vals_pt;
            std::vector<float> vals_eta;
            std::vector<float> vals_phi;
            std::vector<float> vals_mass;
            std::vector<int>   vals_pid;
            for (unsigned int i = 0, n = pup.size(); i  < n; ++i) {
                int status = hepeup.ISTUP[i];
                int idabs = std::abs(hepeup.IDUP[i]);
                if (status == 1){  
                  TLorentzVector p4(pup[i][0], pup[i][1], pup[i][2], pup[i][3]); // x,y,z,t
                  vals_pt.push_back(p4.Pt());
                  vals_eta.push_back(p4.Eta());
                  vals_phi.push_back(p4.Phi());
                  vals_mass.push_back(p4.M());
                  vals_pid.push_back(hepeup.IDUP[i]);
                }  
                if ( (status == 1) && ( ( idabs == 21 ) || (idabs > 0 && idabs < 7) ) ) { //# gluons and quarks
                    // object counters
                    lheNj++;
                    if (idabs==5) lheNb++;
                    else if (idabs==4) lheNc++;
                    else if (idabs <= 3 ) lheNuds++;
                    else if (idabs == 21) lheNglu++;
                    // HT
                    double pt = std::hypot( pup[i][0], pup[i][1] ); // first entry is px, second py
                    lheHT += pt;
                    int mothIdx = std::max(hepeup.MOTHUP[i].first-1,0); //first and last mother as pair; first entry has index 1 in LHE; incoming particles return motherindex 0
                    int mothIdxTwo = std::max(hepeup.MOTHUP[i].second-1,0);
                    int mothStatus  = hepeup.ISTUP[mothIdx];
                    int mothStatusTwo  = hepeup.ISTUP[mothIdxTwo];
                    bool hasIncomingAsMother = mothStatus<0 || mothStatusTwo<0;
                    if (hasIncomingAsMother) lheHTIncoming += pt;
                } else if (idabs == 12 || idabs == 14 || idabs == 16) {
                    (hepeup.IDUP[i] > 0 ? nu : nuBar) = i;
                } else if (idabs == 11 || idabs == 13 || idabs == 15) {
                    (hepeup.IDUP[i] > 0 ? lep : lepBar) = i;
                }
            }
            std::pair<int,int> v(0,0);
            if      (lep != -1 && lepBar != -1) v = std::make_pair(lep,lepBar);
            else if (lep != -1 &&  nuBar != -1) v = std::make_pair(lep, nuBar);
            else if (nu  != -1 && lepBar != -1) v = std::make_pair(nu ,lepBar);
            else if (nu  != -1 &&  nuBar != -1) v = std::make_pair(nu , nuBar);
            if (v.first != -1 && v.second != -1) {
                lheVpt = std::hypot( pup[v.first][0] + pup[v.second][0], pup[v.first][1] + pup[v.second][1] ); 
            }

            out.addColumnValue<uint8_t>("Njets", lheNj, "Number of jets (partons) at LHE step", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<uint8_t>("Nb",    lheNb, "Number of b partons at LHE step", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<uint8_t>("Nc",    lheNc, "Number of c partons at LHE step", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<uint8_t>("Nuds", lheNuds, "Number of u,d,s partons at LHE step", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<uint8_t>("Nglu", lheNglu, "Number of gluon partons at LHE step", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<float>("HT", lheHT, "HT, scalar sum of parton pTs at LHE step", nanoaod::FlatTable::FloatColumn);
            out.addColumnValue<float>("HTIncoming", lheHTIncoming, "HT, scalar sum of parton pTs at LHE step, restricted to partons", nanoaod::FlatTable::FloatColumn);
            out.addColumnValue<float>("Vpt", lheVpt, "pT of the W or Z boson at LHE step", nanoaod::FlatTable::FloatColumn);
            out.addColumnValue<uint8_t>("NpNLO", lheProd.npNLO(), "number of partons at NLO", nanoaod::FlatTable::UInt8Column);
            out.addColumnValue<uint8_t>("NpLO", lheProd.npLO(), "number of partons at LO", nanoaod::FlatTable::UInt8Column);
 
  
            auto outPart  = std::make_unique<nanoaod::FlatTable>(vals_pt.size(), "LHEPart", false);
            outPart->addColumn<float>("pt",     vals_pt,     "Pt of LHE particles", nanoaod::FlatTable::FloatColumn, this->precision_);
            outPart->addColumn<float>("eta",    vals_eta,    "Pseodorapidity of LHE particles", nanoaod::FlatTable::FloatColumn, this->precision_);
            outPart->addColumn<float>("phi",    vals_phi,    "Phi of LHE particles", nanoaod::FlatTable::FloatColumn, this->precision_);
            outPart->addColumn<float>("mass",   vals_mass,   "Mass of LHE particles", nanoaod::FlatTable::FloatColumn, this->precision_);
            outPart->addColumn<int>  ("pdgId",  vals_pid,    "PDG ID of LHE particles", nanoaod::FlatTable::IntColumn);

            return outPart;
        }

        static void fillDescriptions(edm::ConfigurationDescriptions & descriptions) {
            edm::ParameterSetDescription desc;
            desc.add<std::vector<edm::InputTag>>("lheInfo", std::vector<edm::InputTag>{{"externalLHEProducer"},{"source"}})->setComment("tag(s) for the LHE information (LHEEventProduct)");
            desc.add<int>("precision", -1)->setComment("precision on the 4-momenta of the LHE particles");
            desc.add<bool>("storeLHEParticles", false)->setComment("Whether we want to store the 4-momenta of the status 1 particles at LHE level");
            descriptions.add("lheInfoTable", desc);
        }

    protected:
        const std::vector<edm::EDGetTokenT<LHEEventProduct>> lheTag_;
        const unsigned int precision_;
        const bool storeLHEParticles_;
};

#include "FWCore/Framework/interface/MakerMacros.h"
DEFINE_FWK_MODULE(LHETablesProducer);