Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:20:45

0001 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0002 #include "SimDataFormats/PileupSummaryInfo/interface/PileupSummaryInfo.h"
0003 #include "DataFormats/GeometrySurface/interface/Plane.h"
0004 #include "DataFormats/HcalDetId/interface/HcalDetId.h"
0005 
0006 #include "L1Trigger/L1THGCalUtilities/interface/HGCalTriggerNtupleBase.h"
0007 #include "L1Trigger/L1THGCal/interface/HGCalTriggerTools.h"
0008 
0009 #include "MagneticField/Engine/interface/MagneticField.h"
0010 #include "TrackPropagation/RungeKutta/interface/defaultRKPropagator.h"
0011 #include "TrackPropagation/RungeKutta/interface/RKPropagatorInS.h"
0012 #include "FastSimulation/Event/interface/FSimEvent.h"
0013 
0014 #include "FWCore/Framework/interface/ESHandle.h"
0015 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0016 #include "L1Trigger/L1THGCal/interface/HGCalTriggerGeometryBase.h"
0017 
0018 // NOTE: most of this code is borrowed by https://github.com/CMS-HGCAL/reco-ntuples
0019 // kudos goes to the original authors. Ideally the 2 repos should be merged since they share part of the use case
0020 #include <memory>
0021 
0022 namespace HGCal_helpers {
0023 
0024   class Coordinates {
0025   public:
0026     Coordinates() : x(0), y(0), z(0), eta(0), phi(0) {}
0027     float x, y, z, eta, phi;
0028     inline math::XYZTLorentzVectorD toVector() { return math::XYZTLorentzVectorD(x, y, z, 0); }
0029   };
0030 
0031   class SimpleTrackPropagator {
0032   public:
0033     SimpleTrackPropagator(const MagneticField *f) : field_(f), prod_(field_, alongMomentum), absz_target_(0) {
0034       ROOT::Math::SMatrixIdentity id;
0035       AlgebraicSymMatrix55 C(id);
0036       const float uncert = 0.001;
0037       C *= uncert;
0038       err_ = CurvilinearTrajectoryError(C);
0039     }
0040     void setPropagationTargetZ(const float &z);
0041 
0042     bool propagate(const double px,
0043                    const double py,
0044                    const double pz,
0045                    const double x,
0046                    const double y,
0047                    const double z,
0048                    const float charge,
0049                    Coordinates &coords) const;
0050 
0051     bool propagate(const math::XYZTLorentzVectorD &momentum,
0052                    const math::XYZTLorentzVectorD &position,
0053                    const float charge,
0054                    Coordinates &coords) const;
0055 
0056   private:
0057     SimpleTrackPropagator() : field_(nullptr), prod_(field_, alongMomentum), absz_target_(0) {}
0058     const RKPropagatorInS &RKProp() const { return prod_.propagator; }
0059     Plane::PlanePointer targetPlaneForward_, targetPlaneBackward_;
0060     const MagneticField *field_;
0061     CurvilinearTrajectoryError err_;
0062     defaultRKPropagator::Product prod_;
0063     float absz_target_;
0064   };
0065 
0066   void SimpleTrackPropagator::setPropagationTargetZ(const float &z) {
0067     targetPlaneForward_ = Plane::build(Plane::PositionType(0, 0, std::abs(z)), Plane::RotationType());
0068     targetPlaneBackward_ = Plane::build(Plane::PositionType(0, 0, -std::abs(z)), Plane::RotationType());
0069     absz_target_ = std::abs(z);
0070   }
0071   bool SimpleTrackPropagator::propagate(const double px,
0072                                         const double py,
0073                                         const double pz,
0074                                         const double x,
0075                                         const double y,
0076                                         const double z,
0077                                         const float charge,
0078                                         Coordinates &output) const {
0079     output = Coordinates();
0080 
0081     typedef TrajectoryStateOnSurface TSOS;
0082     GlobalPoint startingPosition(x, y, z);
0083     GlobalVector startingMomentum(px, py, pz);
0084     Plane::PlanePointer startingPlane = Plane::build(Plane::PositionType(x, y, z), Plane::RotationType());
0085     TSOS startingStateP(
0086         GlobalTrajectoryParameters(startingPosition, startingMomentum, charge, field_), err_, *startingPlane);
0087 
0088     TSOS trackStateP;
0089     if (pz > 0) {
0090       trackStateP = RKProp().propagate(startingStateP, *targetPlaneForward_);
0091     } else {
0092       trackStateP = RKProp().propagate(startingStateP, *targetPlaneBackward_);
0093     }
0094     if (trackStateP.isValid()) {
0095       output.x = trackStateP.globalPosition().x();
0096       output.y = trackStateP.globalPosition().y();
0097       output.z = trackStateP.globalPosition().z();
0098       output.phi = trackStateP.globalPosition().phi();
0099       output.eta = trackStateP.globalPosition().eta();
0100       return true;
0101     }
0102     return false;
0103   }
0104 
0105   bool SimpleTrackPropagator::propagate(const math::XYZTLorentzVectorD &momentum,
0106                                         const math::XYZTLorentzVectorD &position,
0107                                         const float charge,
0108                                         Coordinates &output) const {
0109     return propagate(
0110         momentum.px(), momentum.py(), momentum.pz(), position.x(), position.y(), position.z(), charge, output);
0111   }
0112 
0113 }  // namespace HGCal_helpers
0114 
0115 class HGCalTriggerNtupleGen : public HGCalTriggerNtupleBase {
0116 public:
0117   HGCalTriggerNtupleGen(const edm::ParameterSet &);
0118 
0119   void initialize(TTree &, const edm::ParameterSet &, edm::ConsumesCollector &&) final;
0120   void fill(const edm::Event &, const HGCalTriggerNtupleEventSetup &) final;
0121 
0122   enum ReachHGCal { notReach = 0, outsideEESurface = 1, onEESurface = 2 };
0123 
0124 private:
0125   void clear() final;
0126 
0127   edm::EDGetToken gen_token_;
0128   edm::EDGetToken gen_PU_token_;
0129 
0130   int gen_n_;
0131   int gen_PUNumInt_;
0132   float gen_TrueNumInt_;
0133 
0134   float vtx_x_;
0135   float vtx_y_;
0136   float vtx_z_;
0137 
0138   ////////////////////
0139   // GenParticles
0140   //
0141   std::vector<float> genpart_eta_;
0142   std::vector<float> genpart_phi_;
0143   std::vector<float> genpart_pt_;
0144   std::vector<float> genpart_energy_;
0145   std::vector<float> genpart_dvx_;
0146   std::vector<float> genpart_dvy_;
0147   std::vector<float> genpart_dvz_;
0148   std::vector<float> genpart_ovx_;
0149   std::vector<float> genpart_ovy_;
0150   std::vector<float> genpart_ovz_;
0151   std::vector<float> genpart_exx_;
0152   std::vector<float> genpart_exy_;
0153   std::vector<int> genpart_mother_;
0154   std::vector<float> genpart_exphi_;
0155   std::vector<float> genpart_exeta_;
0156   std::vector<float> genpart_fbrem_;
0157   std::vector<int> genpart_pid_;
0158   std::vector<int> genpart_gen_;
0159   std::vector<int> genpart_reachedEE_;
0160   std::vector<bool> genpart_fromBeamPipe_;
0161   std::vector<std::vector<float>> genpart_posx_;
0162   std::vector<std::vector<float>> genpart_posy_;
0163   std::vector<std::vector<float>> genpart_posz_;
0164 
0165   ////////////////////
0166   // reco::GenParticles
0167   //
0168   std::vector<float> gen_eta_;
0169   std::vector<float> gen_phi_;
0170   std::vector<float> gen_pt_;
0171   std::vector<float> gen_energy_;
0172   std::vector<int> gen_charge_;
0173   std::vector<int> gen_pdgid_;
0174   std::vector<int> gen_status_;
0175   std::vector<std::vector<int>> gen_daughters_;
0176 
0177   // -------convenient tool to deal with simulated tracks
0178   std::unique_ptr<FSimEvent> mySimEvent_;
0179 
0180   // and also the magnetic field
0181   const MagneticField *aField_;
0182 
0183   HGCalTriggerTools triggerTools_;
0184 
0185   edm::EDGetToken simTracks_token_;
0186   edm::EDGetToken simVertices_token_;
0187   edm::EDGetToken hepmcev_token_;
0188 };
0189 
0190 DEFINE_EDM_PLUGIN(HGCalTriggerNtupleFactory, HGCalTriggerNtupleGen, "HGCalTriggerNtupleGen");
0191 
0192 HGCalTriggerNtupleGen::HGCalTriggerNtupleGen(const edm::ParameterSet &conf) : HGCalTriggerNtupleBase(conf) {
0193   accessEventSetup_ = false;
0194 }
0195 
0196 void HGCalTriggerNtupleGen::initialize(TTree &tree, const edm::ParameterSet &conf, edm::ConsumesCollector &&collector) {
0197   edm::ParameterSet particleFilter_(conf.getParameter<edm::ParameterSet>("particleFilter"));
0198   mySimEvent_ = std::make_unique<FSimEvent>(particleFilter_);
0199 
0200   gen_token_ = collector.consumes<reco::GenParticleCollection>(conf.getParameter<edm::InputTag>("GenParticles"));
0201   gen_PU_token_ = collector.consumes<std::vector<PileupSummaryInfo>>(conf.getParameter<edm::InputTag>("GenPU"));
0202   tree.Branch("gen_n", &gen_n_, "gen_n/I");
0203   tree.Branch("gen_PUNumInt", &gen_PUNumInt_, "gen_PUNumInt/I");
0204   tree.Branch("gen_TrueNumInt", &gen_TrueNumInt_, "gen_TrueNumInt/F");
0205 
0206   hepmcev_token_ = collector.consumes<edm::HepMCProduct>(conf.getParameter<edm::InputTag>("MCEvent"));
0207 
0208   simTracks_token_ = collector.consumes<std::vector<SimTrack>>(conf.getParameter<edm::InputTag>("SimTracks"));
0209   simVertices_token_ = collector.consumes<std::vector<SimVertex>>(conf.getParameter<edm::InputTag>("SimVertices"));
0210 
0211   tree.Branch("vtx_x", &vtx_x_);
0212   tree.Branch("vtx_y", &vtx_y_);
0213   tree.Branch("vtx_z", &vtx_z_);
0214 
0215   tree.Branch("gen_eta", &gen_eta_);
0216   tree.Branch("gen_phi", &gen_phi_);
0217   tree.Branch("gen_pt", &gen_pt_);
0218   tree.Branch("gen_energy", &gen_energy_);
0219   tree.Branch("gen_charge", &gen_charge_);
0220   tree.Branch("gen_pdgid", &gen_pdgid_);
0221   tree.Branch("gen_status", &gen_status_);
0222   tree.Branch("gen_daughters", &gen_daughters_);
0223 
0224   tree.Branch("genpart_eta", &genpart_eta_);
0225   tree.Branch("genpart_phi", &genpart_phi_);
0226   tree.Branch("genpart_pt", &genpart_pt_);
0227   tree.Branch("genpart_energy", &genpart_energy_);
0228   tree.Branch("genpart_dvx", &genpart_dvx_);
0229   tree.Branch("genpart_dvy", &genpart_dvy_);
0230   tree.Branch("genpart_dvz", &genpart_dvz_);
0231   tree.Branch("genpart_ovx", &genpart_ovx_);
0232   tree.Branch("genpart_ovy", &genpart_ovy_);
0233   tree.Branch("genpart_ovz", &genpart_ovz_);
0234   tree.Branch("genpart_mother", &genpart_mother_);
0235   tree.Branch("genpart_exphi", &genpart_exphi_);
0236   tree.Branch("genpart_exeta", &genpart_exeta_);
0237   tree.Branch("genpart_exx", &genpart_exx_);
0238   tree.Branch("genpart_exy", &genpart_exy_);
0239   tree.Branch("genpart_fbrem", &genpart_fbrem_);
0240   tree.Branch("genpart_pid", &genpart_pid_);
0241   tree.Branch("genpart_gen", &genpart_gen_);
0242   tree.Branch("genpart_reachedEE", &genpart_reachedEE_);
0243   tree.Branch("genpart_fromBeamPipe", &genpart_fromBeamPipe_);
0244   tree.Branch("genpart_posx", &genpart_posx_);
0245   tree.Branch("genpart_posy", &genpart_posy_);
0246   tree.Branch("genpart_posz", &genpart_posz_);
0247 }
0248 
0249 void HGCalTriggerNtupleGen::fill(const edm::Event &iEvent, const HGCalTriggerNtupleEventSetup &es) {
0250   clear();
0251 
0252   edm::Handle<std::vector<PileupSummaryInfo>> PupInfo_h;
0253   iEvent.getByToken(gen_PU_token_, PupInfo_h);
0254   const std::vector<PileupSummaryInfo> &PupInfo = *PupInfo_h;
0255 
0256   mySimEvent_->initializePdt(&(*es.pdt));
0257   aField_ = &(*es.magfield);
0258   triggerTools_.setGeometry(es.geometry.product());
0259 
0260   // This balck magic is needed to use the mySimEvent_
0261   edm::Handle<edm::HepMCProduct> hevH;
0262   edm::Handle<std::vector<SimTrack>> simTracksHandle;
0263   edm::Handle<std::vector<SimVertex>> simVerticesHandle;
0264 
0265   iEvent.getByToken(hepmcev_token_, hevH);
0266   iEvent.getByToken(simTracks_token_, simTracksHandle);
0267   iEvent.getByToken(simVertices_token_, simVerticesHandle);
0268   mySimEvent_->fill(*simTracksHandle, *simVerticesHandle);
0269 
0270   HepMC::GenVertex *primaryVertex = *(hevH)->GetEvent()->vertices_begin();
0271   const float mm2cm = 0.1;
0272   vtx_x_ = primaryVertex->position().x() * mm2cm;  // to put in official units
0273   vtx_y_ = primaryVertex->position().y() * mm2cm;
0274   vtx_z_ = primaryVertex->position().z() * mm2cm;
0275 
0276   HGCal_helpers::SimpleTrackPropagator toHGCalPropagator(aField_);
0277   toHGCalPropagator.setPropagationTargetZ(triggerTools_.getLayerZ(1));
0278   std::vector<FSimTrack *> allselectedgentracks;
0279   const float eeInnerRadius = 25.;
0280   const float eeOuterRadius = 160.;
0281   unsigned int npart = mySimEvent_->nTracks();
0282   for (unsigned int i = 0; i < npart; ++i) {
0283     std::vector<float> xp, yp, zp;
0284     FSimTrack &myTrack(mySimEvent_->track(i));
0285     math::XYZTLorentzVectorD vtx(0, 0, 0, 0);
0286 
0287     int reachedEE = ReachHGCal::notReach;  // compute the extrapolations for the particles reaching EE
0288                                            // and for the gen particles
0289     double fbrem = -1;
0290 
0291     if (std::abs(myTrack.vertex().position().z()) >= triggerTools_.getLayerZ(1))
0292       continue;
0293 
0294     const unsigned nlayers = triggerTools_.lastLayerBH();
0295     if (myTrack.noEndVertex())  // || myTrack.genpartIndex()>=0)
0296     {
0297       HGCal_helpers::Coordinates propcoords;
0298       bool reachesHGCal =
0299           toHGCalPropagator.propagate(myTrack.momentum(), myTrack.vertex().position(), myTrack.charge(), propcoords);
0300       vtx = propcoords.toVector();
0301 
0302       if (reachesHGCal && vtx.Rho() < eeOuterRadius && vtx.Rho() > eeInnerRadius) {
0303         reachedEE = ReachHGCal::onEESurface;
0304         double dpt = 0;
0305 
0306         for (int i = 0; i < myTrack.nDaughters(); ++i)
0307           dpt += myTrack.daughter(i).momentum().pt();
0308         if (abs(myTrack.type()) == 11)
0309           fbrem = dpt / myTrack.momentum().pt();
0310       } else if (reachesHGCal && vtx.Rho() > eeOuterRadius)
0311         reachedEE = ReachHGCal::outsideEESurface;
0312 
0313       HGCal_helpers::SimpleTrackPropagator indiv_particleProp(aField_);
0314       for (unsigned il = 1; il <= nlayers; ++il) {
0315         const float charge = myTrack.charge();
0316         indiv_particleProp.setPropagationTargetZ(triggerTools_.getLayerZ(il));
0317         HGCal_helpers::Coordinates propCoords;
0318         indiv_particleProp.propagate(myTrack.momentum(), myTrack.vertex().position(), charge, propCoords);
0319 
0320         xp.push_back(propCoords.x);
0321         yp.push_back(propCoords.y);
0322         zp.push_back(propCoords.z);
0323       }
0324     } else {
0325       vtx = myTrack.endVertex().position();
0326     }
0327     auto orig_vtx = myTrack.vertex().position();
0328 
0329     allselectedgentracks.push_back(&mySimEvent_->track(i));
0330     // fill branches
0331     genpart_eta_.push_back(myTrack.momentum().eta());
0332     genpart_phi_.push_back(myTrack.momentum().phi());
0333     genpart_pt_.push_back(myTrack.momentum().pt());
0334     genpart_energy_.push_back(myTrack.momentum().energy());
0335     genpart_dvx_.push_back(vtx.x());
0336     genpart_dvy_.push_back(vtx.y());
0337     genpart_dvz_.push_back(vtx.z());
0338 
0339     genpart_ovx_.push_back(orig_vtx.x());
0340     genpart_ovy_.push_back(orig_vtx.y());
0341     genpart_ovz_.push_back(orig_vtx.z());
0342 
0343     HGCal_helpers::Coordinates hitsHGCal;
0344     toHGCalPropagator.propagate(myTrack.momentum(), orig_vtx, myTrack.charge(), hitsHGCal);
0345 
0346     genpart_exphi_.push_back(hitsHGCal.phi);
0347     genpart_exeta_.push_back(hitsHGCal.eta);
0348     genpart_exx_.push_back(hitsHGCal.x);
0349     genpart_exy_.push_back(hitsHGCal.y);
0350 
0351     genpart_fbrem_.push_back(fbrem);
0352     genpart_pid_.push_back(myTrack.type());
0353     genpart_gen_.push_back(myTrack.genpartIndex());
0354     genpart_reachedEE_.push_back(reachedEE);
0355     genpart_fromBeamPipe_.push_back(true);
0356 
0357     genpart_posx_.push_back(xp);
0358     genpart_posy_.push_back(yp);
0359     genpart_posz_.push_back(zp);
0360   }
0361 
0362   edm::Handle<std::vector<reco::GenParticle>> genParticlesHandle;
0363   iEvent.getByToken(gen_token_, genParticlesHandle);
0364   gen_n_ = genParticlesHandle->size();
0365 
0366   for (const auto &particle : *genParticlesHandle) {
0367     gen_eta_.push_back(particle.eta());
0368     gen_phi_.push_back(particle.phi());
0369     gen_pt_.push_back(particle.pt());
0370     gen_energy_.push_back(particle.energy());
0371     gen_charge_.push_back(particle.charge());
0372     gen_pdgid_.push_back(particle.pdgId());
0373     gen_status_.push_back(particle.status());
0374     std::vector<int> daughters(particle.daughterRefVector().size(), 0);
0375     for (unsigned j = 0; j < particle.daughterRefVector().size(); ++j) {
0376       daughters[j] = static_cast<int>(particle.daughterRefVector().at(j).key());
0377     }
0378     gen_daughters_.push_back(daughters);
0379   }
0380 
0381   // associate gen particles to mothers
0382   genpart_mother_.resize(genpart_posz_.size(), -1);
0383   for (size_t i = 0; i < allselectedgentracks.size(); i++) {
0384     const auto tracki = allselectedgentracks.at(i);
0385 
0386     for (size_t j = i + 1; j < allselectedgentracks.size(); j++) {
0387       const auto trackj = allselectedgentracks.at(j);
0388 
0389       if (!tracki->noMother()) {
0390         if (&tracki->mother() == trackj)
0391           genpart_mother_.at(i) = j;
0392       }
0393       if (!trackj->noMother()) {
0394         if (&trackj->mother() == tracki)
0395           genpart_mother_.at(j) = i;
0396       }
0397     }
0398   }
0399 
0400   for (const auto &PVI : PupInfo) {
0401     if (PVI.getBunchCrossing() == 0) {
0402       gen_PUNumInt_ = PVI.getPU_NumInteractions();
0403       gen_TrueNumInt_ = PVI.getTrueNumInteractions();
0404     }
0405   }
0406 }
0407 
0408 void HGCalTriggerNtupleGen::clear() {
0409   gen_n_ = 0;
0410   gen_PUNumInt_ = 0;
0411   gen_TrueNumInt_ = 0.;
0412 
0413   vtx_x_ = 0;
0414   vtx_y_ = 0;
0415   vtx_z_ = 0;
0416 
0417   //
0418   genpart_eta_.clear();
0419   genpart_phi_.clear();
0420   genpart_pt_.clear();
0421   genpart_energy_.clear();
0422   genpart_dvx_.clear();
0423   genpart_dvy_.clear();
0424   genpart_dvz_.clear();
0425   genpart_ovx_.clear();
0426   genpart_ovy_.clear();
0427   genpart_ovz_.clear();
0428   genpart_exx_.clear();
0429   genpart_exy_.clear();
0430   genpart_mother_.clear();
0431   genpart_exphi_.clear();
0432   genpart_exeta_.clear();
0433   genpart_fbrem_.clear();
0434   genpart_pid_.clear();
0435   genpart_gen_.clear();
0436   genpart_reachedEE_.clear();
0437   genpart_fromBeamPipe_.clear();
0438   genpart_posx_.clear();
0439   genpart_posy_.clear();
0440   genpart_posz_.clear();
0441 
0442   ////////////////////
0443   // reco::GenParticles
0444   //
0445   gen_eta_.clear();
0446   gen_phi_.clear();
0447   gen_pt_.clear();
0448   gen_energy_.clear();
0449   gen_charge_.clear();
0450   gen_pdgid_.clear();
0451   gen_status_.clear();
0452   gen_daughters_.clear();
0453 }