HGCalTriggerNtupleGen.cc

Go to the documentation of this file.

#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
#include "SimDataFormats/PileupSummaryInfo/interface/PileupSummaryInfo.h"
#include "DataFormats/GeometrySurface/interface/Plane.h"
#include "DataFormats/ForwardDetId/interface/HGCalDetId.h"
#include "DataFormats/HcalDetId/interface/HcalDetId.h"


#include "L1Trigger/L1THGCal/interface/HGCalTriggerNtupleBase.h"
#include "L1Trigger/L1THGCal/interface/HGCalTriggerTools.h"

#include "MagneticField/Engine/interface/MagneticField.h"
#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
#include "TrackPropagation/RungeKutta/interface/defaultRKPropagator.h"
#include "TrackPropagation/RungeKutta/interface/RKPropagatorInS.h"
#include "FastSimulation/Event/interface/FSimEvent.h"
#include "FastSimulation/Particle/interface/ParticleTable.h"
#include "FastSimulation/CaloGeometryTools/interface/Transform3DPJ.h"
#include "SimGeneral/HepPDTRecord/interface/PDTRecord.h"

#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/ESWatcher.h"
#include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
#include "L1Trigger/L1THGCal/interface/HGCalTriggerGeometryBase.h"

// NOTE: most of this code is borrowed by https://github.com/CMS-HGCAL/reco-ntuples
// kudos goes to the original authors. Ideally the 2 repos should be merged since they share part of the use case
#include <memory>

namespace HGCal_helpers {

  class Coordinates {
   public:
    Coordinates() : x(0), y(0), z(0), eta(0), phi(0) {}
    float x, y, z, eta, phi;
    inline math::XYZTLorentzVectorD toVector() { return math::XYZTLorentzVectorD(x, y, z, 0); }
  };


  class SimpleTrackPropagator {
   public:
    SimpleTrackPropagator(const MagneticField *f)
        : field_(f), prod_(field_, alongMomentum), absz_target_(0) {
      ROOT::Math::SMatrixIdentity id;
      AlgebraicSymMatrix55 C(id);
      const float uncert = 0.001;
      C *= uncert;
      err_ = CurvilinearTrajectoryError(C);
    }
    void setPropagationTargetZ(const float &z);

    bool propagate(const double px, const double py, const double pz, const double x, const double y,
                   const double z, const float charge, Coordinates &coords) const;

    bool propagate(const math::XYZTLorentzVectorD &momentum, const math::XYZTLorentzVectorD &position,
                   const float charge, Coordinates &coords) const;

   private:
    SimpleTrackPropagator() : field_(nullptr), prod_(field_, alongMomentum), absz_target_(0) {}
    const RKPropagatorInS &RKProp() const { return prod_.propagator; }
    Plane::PlanePointer targetPlaneForward_, targetPlaneBackward_;
    const MagneticField *field_;
    CurvilinearTrajectoryError err_;
    defaultRKPropagator::Product prod_;
    float absz_target_;
  };

  void SimpleTrackPropagator::setPropagationTargetZ(const float &z) {
    targetPlaneForward_ = Plane::build(Plane::PositionType(0, 0, std::abs(z)), Plane::RotationType());
    targetPlaneBackward_ =
        Plane::build(Plane::PositionType(0, 0, -std::abs(z)), Plane::RotationType());
    absz_target_ = std::abs(z);
  }
  bool SimpleTrackPropagator::propagate(const double px, const double py, const double pz,
                                        const double x, const double y, const double z,
                                        const float charge, Coordinates &output) const {
    output = Coordinates();

    typedef TrajectoryStateOnSurface TSOS;
    GlobalPoint startingPosition(x, y, z);
    GlobalVector startingMomentum(px, py, pz);
    Plane::PlanePointer startingPlane =
        Plane::build(Plane::PositionType(x, y, z), Plane::RotationType());
    TSOS startingStateP(
        GlobalTrajectoryParameters(startingPosition, startingMomentum, charge, field_), err_,
        *startingPlane);

    TSOS trackStateP;
    if (pz > 0) {
      trackStateP = RKProp().propagate(startingStateP, *targetPlaneForward_);
    } else {
      trackStateP = RKProp().propagate(startingStateP, *targetPlaneBackward_);
    }
    if (trackStateP.isValid()) {
      output.x = trackStateP.globalPosition().x();
      output.y = trackStateP.globalPosition().y();
      output.z = trackStateP.globalPosition().z();
      output.phi = trackStateP.globalPosition().phi();
      output.eta = trackStateP.globalPosition().eta();
      return true;
    }
    return false;
  }

  bool SimpleTrackPropagator::propagate(const math::XYZTLorentzVectorD &momentum,
                                        const math::XYZTLorentzVectorD &position, const float charge,
                                        Coordinates &output) const {
    return propagate(momentum.px(), momentum.py(), momentum.pz(), position.x(), position.y(),
                     position.z(), charge, output);
  }

}  // HGCal_helpers


class HGCalTriggerNtupleGen : public HGCalTriggerNtupleBase
{

    public:
        HGCalTriggerNtupleGen(const edm::ParameterSet& );

        void initialize(TTree&, const edm::ParameterSet&, edm::ConsumesCollector&&) final;
        void fill(const edm::Event&, const edm::EventSetup& ) final;

        typedef ROOT::Math::Transform3DPJ Transform3D;
        typedef ROOT::Math::Transform3DPJ::Point Point;

        enum ReachHGCal {
          notReach = 0,
          outsideEESurface = 1,
          onEESurface = 2
        };

    private:
        void clear() final;

        edm::EDGetToken gen_token_;
        edm::EDGetToken gen_PU_token_;
        
        int gen_n_;
        int gen_PUNumInt_;
        float gen_TrueNumInt_;

        float vtx_x_;
        float vtx_y_;
        float vtx_z_;


        // GenParticles
        //
        std::vector<float> genpart_eta_;
        std::vector<float> genpart_phi_;
        std::vector<float> genpart_pt_;
        std::vector<float> genpart_energy_;
        std::vector<float> genpart_dvx_;
        std::vector<float> genpart_dvy_;
        std::vector<float> genpart_dvz_;
        std::vector<float> genpart_ovx_;
        std::vector<float> genpart_ovy_;
        std::vector<float> genpart_ovz_;
        std::vector<float> genpart_exx_;
        std::vector<float> genpart_exy_;
        std::vector<int> genpart_mother_;
        std::vector<float> genpart_exphi_;
        std::vector<float> genpart_exeta_;
        std::vector<float> genpart_fbrem_;
        std::vector<int> genpart_pid_;
        std::vector<int> genpart_gen_;
        std::vector<int> genpart_reachedEE_;
        std::vector<bool> genpart_fromBeamPipe_;
        std::vector<std::vector<float>> genpart_posx_;
        std::vector<std::vector<float>> genpart_posy_;
        std::vector<std::vector<float>> genpart_posz_;

        // reco::GenParticles
        //
        std::vector<float> gen_eta_;
        std::vector<float> gen_phi_;
        std::vector<float> gen_pt_;
        std::vector<float> gen_energy_;
        std::vector<int> gen_charge_;
        std::vector<int> gen_pdgid_;
        std::vector<int> gen_status_;
        std::vector<std::vector<int>> gen_daughters_;


        // -------convenient tool to deal with simulated tracks
        std::unique_ptr<FSimEvent> mySimEvent_;

        // and also the magnetic field
        const MagneticField *aField_;

        HGCalTriggerTools triggerTools_;

        // edm::EDGetTokenT<std::vector<reco::GenParticle> > genParticles_;
        edm::EDGetToken simTracks_token_;
        edm::EDGetToken simVertices_token_;
        edm::EDGetToken hepmcev_token_;

        edm::ESWatcher<PDTRecord> pdt_watcher_;
        edm::ESWatcher<IdealMagneticFieldRecord> magfield_watcher_;


};

DEFINE_EDM_PLUGIN(HGCalTriggerNtupleFactory,
        HGCalTriggerNtupleGen,
        "HGCalTriggerNtupleGen" );


HGCalTriggerNtupleGen::
HGCalTriggerNtupleGen(const edm::ParameterSet& conf):HGCalTriggerNtupleBase(conf)
{
}


void
HGCalTriggerNtupleGen::
initialize(TTree& tree, const edm::ParameterSet& conf, edm::ConsumesCollector&& collector)
{

    edm::ParameterSet particleFilter_(conf.getParameter<edm::ParameterSet>("particleFilter"));
    mySimEvent_ = std::make_unique<FSimEvent>(particleFilter_);

    gen_token_ = collector.consumes<reco::GenParticleCollection>(conf.getParameter<edm::InputTag>("GenParticles"));
    gen_PU_token_ = collector.consumes<std::vector<PileupSummaryInfo>>(conf.getParameter<edm::InputTag>("GenPU"));
    tree.Branch("gen_n", &gen_n_, "gen_n/I");
    tree.Branch("gen_PUNumInt", &gen_PUNumInt_ ,"gen_PUNumInt/I");
    tree.Branch("gen_TrueNumInt", &gen_TrueNumInt_ ,"gen_TrueNumInt/F");

    hepmcev_token_ = collector.consumes<edm::HepMCProduct>(conf.getParameter<edm::InputTag>("MCEvent"));

    simTracks_token_ = collector.consumes<std::vector<SimTrack>>(conf.getParameter<edm::InputTag>("SimTracks"));
    simVertices_token_ = collector.consumes<std::vector<SimVertex>>(conf.getParameter<edm::InputTag>("SimVertices"));

    tree.Branch("vtx_x", &vtx_x_);
    tree.Branch("vtx_y", &vtx_y_);
    tree.Branch("vtx_z", &vtx_z_);


    tree.Branch("gen_eta", &gen_eta_);
    tree.Branch("gen_phi", &gen_phi_);
    tree.Branch("gen_pt", &gen_pt_);
    tree.Branch("gen_energy", &gen_energy_);
    tree.Branch("gen_charge", &gen_charge_);
    tree.Branch("gen_pdgid", &gen_pdgid_);
    tree.Branch("gen_status", &gen_status_);
    tree.Branch("gen_daughters", &gen_daughters_);

    tree.Branch("genpart_eta", &genpart_eta_);
    tree.Branch("genpart_phi", &genpart_phi_);
    tree.Branch("genpart_pt", &genpart_pt_);
    tree.Branch("genpart_energy", &genpart_energy_);
    tree.Branch("genpart_dvx", &genpart_dvx_);
    tree.Branch("genpart_dvy", &genpart_dvy_);
    tree.Branch("genpart_dvz", &genpart_dvz_);
    tree.Branch("genpart_ovx", &genpart_ovx_);
    tree.Branch("genpart_ovy", &genpart_ovy_);
    tree.Branch("genpart_ovz", &genpart_ovz_);
    tree.Branch("genpart_mother", &genpart_mother_);
    tree.Branch("genpart_exphi", &genpart_exphi_);
    tree.Branch("genpart_exeta", &genpart_exeta_);
    tree.Branch("genpart_exx", &genpart_exx_);
    tree.Branch("genpart_exy", &genpart_exy_);
    tree.Branch("genpart_fbrem", &genpart_fbrem_);
    tree.Branch("genpart_pid", &genpart_pid_);
    tree.Branch("genpart_gen", &genpart_gen_);
    tree.Branch("genpart_reachedEE", &genpart_reachedEE_);
    tree.Branch("genpart_fromBeamPipe", &genpart_fromBeamPipe_);
    tree.Branch("genpart_posx", &genpart_posx_);
    tree.Branch("genpart_posy", &genpart_posy_);
    tree.Branch("genpart_posz", &genpart_posz_);


}

void
HGCalTriggerNtupleGen::
fill(const edm::Event& iEvent, const edm::EventSetup& es)
{
    clear();

    edm::Handle<std::vector< PileupSummaryInfo > >  PupInfo_h;
    iEvent.getByToken(gen_PU_token_, PupInfo_h);
    const std::vector< PileupSummaryInfo >& PupInfo = *PupInfo_h;

    if(pdt_watcher_.check(es))
    {
      edm::ESHandle<HepPDT::ParticleDataTable> pdt;
      es.get<PDTRecord>().get(pdt);
      mySimEvent_->initializePdt(&(*pdt));
    }

    if(magfield_watcher_.check(es))
    {
      edm::ESHandle<MagneticField> magfield;
      es.get<IdealMagneticFieldRecord>().get(magfield);
      aField_ = &(*magfield);
    }

    triggerTools_.eventSetup(es);

    // This balck magic is needed to use the mySimEvent_
    ParticleTable::Sentry ptable(mySimEvent_->theTable());
    edm::Handle<edm::HepMCProduct> hevH;
    edm::Handle<std::vector<SimTrack>> simTracksHandle;
    edm::Handle<std::vector<SimVertex>> simVerticesHandle;

    iEvent.getByToken(hepmcev_token_, hevH);
    iEvent.getByToken(simTracks_token_, simTracksHandle);
    iEvent.getByToken(simVertices_token_, simVerticesHandle);
    mySimEvent_->fill(*simTracksHandle, *simVerticesHandle);

    HepMC::GenVertex *primaryVertex = *(hevH)->GetEvent()->vertices_begin();
    const float mm2cm = 0.1;
    vtx_x_ = primaryVertex->position().x() * mm2cm;  // to put in official units
    vtx_y_ = primaryVertex->position().y() * mm2cm;
    vtx_z_ = primaryVertex->position().z() * mm2cm;
    Point sim_pv(vtx_x_, vtx_y_, vtx_z_);


    HGCal_helpers::SimpleTrackPropagator toHGCalPropagator(aField_);
    toHGCalPropagator.setPropagationTargetZ(triggerTools_.getLayerZ(1));
    std::vector<FSimTrack *> allselectedgentracks;
    const float eeInnerRadius = 25.;
    const float eeOuterRadius = 160.;
    unsigned int npart = mySimEvent_->nTracks();
    for (unsigned int i = 0; i < npart; ++i) {
      std::vector<float> xp, yp, zp;
      FSimTrack &myTrack(mySimEvent_->track(i));
      math::XYZTLorentzVectorD vtx(0, 0, 0, 0);

      int reachedEE = ReachHGCal::notReach;  // compute the extrapolations for the particles reaching EE
                          // and for the gen particles
      double fbrem = -1;

      if (std::abs(myTrack.vertex().position().z()) >= triggerTools_.getLayerZ(1)) continue;

      const unsigned nlayers = triggerTools_.lastLayerBH();
      if (myTrack.noEndVertex())  // || myTrack.genpartIndex()>=0)
      {
        HGCal_helpers::Coordinates propcoords;
        bool reachesHGCal = toHGCalPropagator.propagate(
            myTrack.momentum(), myTrack.vertex().position(), myTrack.charge(), propcoords);
        vtx = propcoords.toVector();

        if (reachesHGCal && vtx.Rho() < eeOuterRadius && vtx.Rho() > eeInnerRadius) {
          reachedEE = ReachHGCal::onEESurface;
          double dpt = 0;

          for (int i = 0; i < myTrack.nDaughters(); ++i) dpt += myTrack.daughter(i).momentum().pt();
          if (abs(myTrack.type()) == 11) fbrem = dpt / myTrack.momentum().pt();
        } else if (reachesHGCal && vtx.Rho() > eeOuterRadius)
          reachedEE = ReachHGCal::outsideEESurface;

        HGCal_helpers::SimpleTrackPropagator indiv_particleProp(aField_);
        for (unsigned il = 1; il <= nlayers; ++il) {
          const float charge = myTrack.charge();
          indiv_particleProp.setPropagationTargetZ(triggerTools_.getLayerZ(il));
          HGCal_helpers::Coordinates propCoords;
          indiv_particleProp.propagate(myTrack.momentum(), myTrack.vertex().position(), charge,
                                       propCoords);

          xp.push_back(propCoords.x);
          yp.push_back(propCoords.y);
          zp.push_back(propCoords.z);
        }
      } else {
        vtx = myTrack.endVertex().position();
      }
      auto orig_vtx = myTrack.vertex().position();

      allselectedgentracks.push_back(&mySimEvent_->track(i));
      // fill branches
      genpart_eta_.push_back(myTrack.momentum().eta());
      genpart_phi_.push_back(myTrack.momentum().phi());
      genpart_pt_.push_back(myTrack.momentum().pt());
      genpart_energy_.push_back(myTrack.momentum().energy());
      genpart_dvx_.push_back(vtx.x());
      genpart_dvy_.push_back(vtx.y());
      genpart_dvz_.push_back(vtx.z());

      genpart_ovx_.push_back(orig_vtx.x());
      genpart_ovy_.push_back(orig_vtx.y());
      genpart_ovz_.push_back(orig_vtx.z());

      HGCal_helpers::Coordinates hitsHGCal;
      toHGCalPropagator.propagate(myTrack.momentum(), orig_vtx, myTrack.charge(), hitsHGCal);

      genpart_exphi_.push_back(hitsHGCal.phi);
      genpart_exeta_.push_back(hitsHGCal.eta);
      genpart_exx_.push_back(hitsHGCal.x);
      genpart_exy_.push_back(hitsHGCal.y);

      genpart_fbrem_.push_back(fbrem);
      genpart_pid_.push_back(myTrack.type());
      genpart_gen_.push_back(myTrack.genpartIndex());
      genpart_reachedEE_.push_back(reachedEE);
      genpart_fromBeamPipe_.push_back(true);

      genpart_posx_.push_back(xp);
      genpart_posy_.push_back(yp);
      genpart_posz_.push_back(zp);
    }


    edm::Handle<std::vector<reco::GenParticle>> genParticlesHandle;
    iEvent.getByToken(gen_token_, genParticlesHandle);
    gen_n_ = genParticlesHandle->size();

    for (const auto& particle : *genParticlesHandle) {
      gen_eta_.push_back(particle.eta());
      gen_phi_.push_back(particle.phi());
      gen_pt_.push_back(particle.pt());
      gen_energy_.push_back(particle.energy());
      gen_charge_.push_back(particle.charge());
      gen_pdgid_.push_back(particle.pdgId());
      gen_status_.push_back(particle.status());
      std::vector<int> daughters(particle.daughterRefVector().size(), 0);
      for (unsigned j = 0; j < particle.daughterRefVector().size(); ++j) {
        daughters[j] = static_cast<int>(particle.daughterRefVector().at(j).key());
      }
      gen_daughters_.push_back(daughters);
    }


    // associate gen particles to mothers
    genpart_mother_.resize(genpart_posz_.size(), -1);
    for (size_t i = 0; i < allselectedgentracks.size(); i++) {
      const auto tracki = allselectedgentracks.at(i);

      for (size_t j = i + 1; j < allselectedgentracks.size(); j++) {
        const auto trackj = allselectedgentracks.at(j);

        if (!tracki->noMother()) {
          if (&tracki->mother() == trackj) genpart_mother_.at(i) = j;
        }
        if (!trackj->noMother()) {
          if (&trackj->mother() == tracki) genpart_mother_.at(j) = i;
        }
      }
    }

    for(const auto& PVI : PupInfo)
    {
        if(PVI.getBunchCrossing() == 0)
        { 
            gen_PUNumInt_ = PVI.getPU_NumInteractions();
            gen_TrueNumInt_ = PVI.getTrueNumInteractions();
        }
    }

}


void
HGCalTriggerNtupleGen::
clear()
{

    gen_n_ = 0;
    gen_PUNumInt_ = 0;
    gen_TrueNumInt_ = 0.;

    vtx_x_ = 0;
    vtx_y_ = 0;
    vtx_z_ = 0;


    //
    genpart_eta_.clear();
    genpart_phi_.clear();
    genpart_pt_.clear();
    genpart_energy_.clear();
    genpart_dvx_.clear();
    genpart_dvy_.clear();
    genpart_dvz_.clear();
    genpart_ovx_.clear();
    genpart_ovy_.clear();
    genpart_ovz_.clear();
    genpart_exx_.clear();
    genpart_exy_.clear();
    genpart_mother_.clear();
    genpart_exphi_.clear();
    genpart_exeta_.clear();
    genpart_fbrem_.clear();
    genpart_pid_.clear();
    genpart_gen_.clear();
    genpart_reachedEE_.clear();
    genpart_fromBeamPipe_.clear();
    genpart_posx_.clear();
    genpart_posy_.clear();
    genpart_posz_.clear();

    // reco::GenParticles
    //
    gen_eta_.clear();
    gen_phi_.clear();
    gen_pt_.clear();
    gen_energy_.clear();
    gen_charge_.clear();
    gen_pdgid_.clear();
    gen_status_.clear();
    gen_daughters_.clear();
}