NuclearSeedGenerator/plugins/NuclearTrackCorrector.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    NuclearTrackCorrector
0004 // Class:      NuclearTrackCorrector
0005 //
0006 /**\class NuclearTrackCorrector NuclearTrackCorrector.cc RecoTracker/NuclearSeedGenerator/plugin/NuclearTrackCorrector.cc
0007
0008  Description: <one line class summary>
0009
0010  Implementation:
0011      <Notes on implementation>
0012 */
0013 //
0014 // Original Author:  Loic QUERTENMONT, Vincent ROBERFROID
0015 //         Created:  Tue Sep 18 14:22:48 CEST 2007
0016 //
0017 //
0018
0019 #include "NuclearTrackCorrector.h"
0020 #include "FWCore/Framework/interface/EventSetup.h"
0021 #include "Geometry/Records/interface/TrackerDigiGeometryRecord.h"
0022 #include "RecoTracker/CkfPattern/interface/TransientInitialStateEstimator.h"
0023 #include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"
0024 #include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
0025
0026 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0027
0028 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0029
0030 using namespace edm;
0031 using namespace std;
0032 using namespace reco;
0033
0034 NuclearTrackCorrector::NuclearTrackCorrector(const edm::ParameterSet& iConfig) : theInitialState(nullptr) {
0035   str_Input_Trajectory = iConfig.getParameter<std::string>("InputTrajectory");
0036   str_Input_NuclearInteraction = iConfig.getParameter<std::string>("InputNuclearInteraction");
0037   verbosity = iConfig.getParameter<int>("Verbosity");
0038   KeepOnlyCorrectedTracks = iConfig.getParameter<bool>("KeepOnlyCorrectedTracks");
0039
0040   theAlgo = new TrackProducerAlgorithm<reco::Track>(iConfig);
0041
0042   produces<TrajectoryCollection>();
0043   produces<TrajectoryToTrajectoryMap>();
0044
0045   produces<reco::TrackExtraCollection>();
0046   produces<reco::TrackCollection>();
0047   produces<TrackToTrajectoryMap>();
0048
0049   produces<TrackToTrackMap>();
0050
0051   theGToken = esConsumes();
0052   theMFToken = esConsumes();
0053   std::string fitterName = iConfig.getParameter<std::string>("Fitter");
0054   theFitterToken = esConsumes(edm::ESInputTag("", fitterName));
0055   std::string propagatorName = iConfig.getParameter<std::string>("Propagator");
0056   thePropagatorToken = esConsumes(edm::ESInputTag("", propagatorName));
0057 }
0058
0059 NuclearTrackCorrector::~NuclearTrackCorrector() {}
0060
0061 void NuclearTrackCorrector::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0062   // Create Output Collections
0063   // --------------------------------------------------------------------------------------------------
0064   auto Output_traj = std::make_unique<TrajectoryCollection>();
0065   auto Output_trajmap = std::make_unique<TrajectoryToTrajectoryMap>();
0066
0067   auto Output_trackextra = std::make_unique<reco::TrackExtraCollection>();
0068   auto Output_track = std::make_unique<reco::TrackCollection>();
0069   auto Output_trackmap = std::make_unique<TrackToTrajectoryMap>();
0070
0071   // Load Reccord
0072   // --------------------------------------------------------------------------------------------------
0073   theFitter = iSetup.getHandle(theFitterToken);
0074
0075   thePropagator = iSetup.getHandle(thePropagatorToken);
0076   theG = iSetup.getHandle(theGToken);
0077
0078   reco::TrackExtraRefProd rTrackExtras = iEvent.getRefBeforePut<reco::TrackExtraCollection>();
0079
0080   theMF = iSetup.getHandle(theMFToken);
0081
0082   // Load Inputs
0083   // --------------------------------------------------------------------------------------------------
0084   edm::Handle<TrajectoryCollection> temp_m_TrajectoryCollection;
0085   iEvent.getByLabel(str_Input_Trajectory, temp_m_TrajectoryCollection);
0086   const TrajectoryCollection m_TrajectoryCollection = *(temp_m_TrajectoryCollection.product());
0087
0088   edm::Handle<NuclearInteractionCollection> temp_m_NuclearInteractionCollection;
0089   iEvent.getByLabel(str_Input_NuclearInteraction, temp_m_NuclearInteractionCollection);
0090   const NuclearInteractionCollection m_NuclearInteractionCollection = *(temp_m_NuclearInteractionCollection.product());
0091
0092   edm::Handle<TrajTrackAssociationCollection> h_TrajToTrackCollection;
0093   iEvent.getByLabel(str_Input_Trajectory, h_TrajToTrackCollection);
0094   m_TrajToTrackCollection = h_TrajToTrackCollection.product();
0095
0096   // Correct the trajectories (Remove trajectory's hits that are located after the nuclear interacion)
0097   // --------------------------------------------------------------------------------------------------
0098   if (verbosity >= 1) {
0099     LogDebug("NuclearTrackCorrector") << "Number of trajectories                    = " << m_TrajectoryCollection.size()
0100                                       << std::endl
0101                                       << "Number of nuclear interactions            = "
0102                                       << m_NuclearInteractionCollection.size();
0103   }
0104
0105   std::map<reco::TrackRef, TrajectoryRef> m_TrackToTrajMap;
0106   swap_map(temp_m_TrajectoryCollection, m_TrackToTrajMap);
0107
0108   for (unsigned int i = 0; i < m_NuclearInteractionCollection.size(); i++) {
0109     const reco::NuclearInteraction& ni = m_NuclearInteractionCollection[i];
0110     if (ni.likelihood() < 0.4)
0111       continue;
0112
0113     reco::TrackRef primTrackRef = ni.primaryTrack().castTo<reco::TrackRef>();
0114
0115     TrajectoryRef trajRef = m_TrackToTrajMap[primTrackRef];
0116
0117     Trajectory newTraj;
0118     if (newTrajNeeded(newTraj, trajRef, ni)) {
0119       AlgoProductCollection algoResults;
0120       bool isOK = getTrackFromTrajectory(newTraj, trajRef, algoResults);
0121
0122       if (isOK) {
0123         pair<unsigned int, unsigned int> tempory_pair;
0124         tempory_pair.first = Output_track->size();
0125         tempory_pair.second = i;
0126         Indice_Map.push_back(tempory_pair);
0127
0128         reco::TrackExtraRef teref = reco::TrackExtraRef(rTrackExtras, i);
0129         reco::TrackExtra newTrackExtra = getNewTrackExtra(algoResults);
0130         (algoResults[0].track)->setExtra(teref);
0131
0132         Output_track->push_back(*algoResults[0].track);
0133         Output_trackextra->push_back(newTrackExtra);
0134         Output_traj->push_back(newTraj);
0135       }
0136     } else {
0137       if (!KeepOnlyCorrectedTracks) {
0138         Output_track->push_back(*primTrackRef);
0139         Output_trackextra->push_back(*primTrackRef->extra());
0140         Output_traj->push_back(*trajRef);
0141       }
0142     }
0143   }
0144   const edm::OrphanHandle<TrajectoryCollection> Handle_traj = iEvent.put(std::move(Output_traj));
0145   const edm::OrphanHandle<reco::TrackCollection> Handle_tracks = iEvent.put(std::move(Output_track));
0146   iEvent.put(std::move(Output_trackextra));
0147
0148   // Make Maps between elements
0149   // --------------------------------------------------------------------------------------------------
0150   if (Handle_tracks->size() != Handle_traj->size()) {
0151     printf("ERROR Handle_tracks->size() != Handle_traj->size() \n");
0152     return;
0153   }
0154
0155   auto Output_tracktrackmap = std::make_unique<TrackToTrackMap>(Handle_tracks, m_TrajToTrackCollection->refProd().val);
0156
0157   for (unsigned int i = 0; i < Indice_Map.size(); i++) {
0158     TrajectoryRef InTrajRef(temp_m_TrajectoryCollection, Indice_Map[i].second);
0159     TrajectoryRef OutTrajRef(Handle_traj, Indice_Map[i].first);
0160     reco::TrackRef TrackRef(Handle_tracks, Indice_Map[i].first);
0161
0162     Output_trajmap->insert(OutTrajRef, InTrajRef);
0163     Output_trackmap->insert(TrackRef, InTrajRef);
0164
0165     try {
0166       reco::TrackRef PrimaryTrackRef = m_TrajToTrackCollection->operator[](InTrajRef);
0167       Output_tracktrackmap->insert(TrackRef, PrimaryTrackRef);
0168     } catch (edm::Exception const&) {
0169     }
0170   }
0171   iEvent.put(std::move(Output_trajmap));
0172   iEvent.put(std::move(Output_trackmap));
0173   iEvent.put(std::move(Output_tracktrackmap));
0174
0175   if (verbosity >= 3)
0176     printf("-----------------------\n");
0177 }
0178
0179 //----------------------------------------------------------------------------------------
0180 bool NuclearTrackCorrector::newTrajNeeded(Trajectory& newtrajectory,
0181                                           const TrajectoryRef& trajRef,
0182                                           const NuclearInteraction& ni) {
0183   bool needNewTraj = false;
0184   reco::Vertex::Point vtx_pos = ni.vertex().position();
0185   double vtx_pos_mag = sqrt(vtx_pos.X() * vtx_pos.X() + vtx_pos.Y() * vtx_pos.Y() + vtx_pos.Z() * vtx_pos.Z());
0186   if (verbosity >= 2)
0187     printf("Nuclear Interaction pos = %f\n", vtx_pos_mag);
0188
0189   newtrajectory = Trajectory(trajRef->seed(), alongMomentum);
0190
0191   // Look all the Hits of the trajectory and keep only Hits before seeds
0192   Trajectory::DataContainer Measurements = trajRef->measurements();
0193   if (verbosity >= 2)
0194     LogDebug("NuclearTrackCorrector") << "Size of Measurements  = " << Measurements.size();
0195
0196   for (unsigned int m = Measurements.size() - 1; m != (unsigned int)-1; m--) {
0197     if (!Measurements[m].recHit()->isValid())
0198       continue;
0199     GlobalPoint hit_pos = theG->idToDet(Measurements[m].recHit()->geographicalId())
0200                               ->surface()
0201                               .toGlobal(Measurements[m].recHit()->localPosition());
0202
0203     if (verbosity >= 2)
0204       printf("Hit pos = %f", hit_pos.mag());
0205
0206     if (hit_pos.mag() > vtx_pos_mag) {
0207       if (verbosity >= 2)
0208         printf(" X ");
0209       needNewTraj = true;
0210     } else {
0211       newtrajectory.push(Measurements[m]);
0212     }
0213     if (verbosity >= 2)
0214       printf("\n");
0215   }
0216
0217   return needNewTraj;
0218 }
0219
0220 //----------------------------------------------------------------------------------------
0221 bool NuclearTrackCorrector::getTrackFromTrajectory(const Trajectory& newTraj,
0222                                                    const TrajectoryRef& initialTrajRef,
0223                                                    AlgoProductCollection& algoResults) {
0224   const Trajectory* it = &newTraj;
0225
0226   TransientTrackingRecHit::RecHitContainer hits;
0227   it->validRecHits(hits);
0228
0229   float ndof = 0;
0230   for (unsigned int h = 0; h < hits.size(); h++) {
0231     if (hits[h]->isValid()) {
0232       ndof = ndof + hits[h]->dimension() * hits[h]->weight();
0233     } else {
0234       LogDebug("NuclearSeedGenerator") << " HIT IS INVALID ???";
0235     }
0236   }
0237
0238   ndof = ndof - 5;
0239   reco::TrackRef theT = m_TrajToTrackCollection->operator[](initialTrajRef);
0240   LogDebug("NuclearSeedGenerator") << " TrackCorrector - number of valid hits" << hits.size() << "\n"
0241                                    << "                - number of hits from Track " << theT->recHitsSize() << "\n"
0242                                    << "                - number of valid hits from initial track "
0243                                    << theT->numberOfValidHits();
0244
0245   if (hits.size() > 1) {
0246     TrajectoryStateOnSurface theInitialStateForRefitting =
0247         getInitialState(&(*theT), hits, theG.product(), theMF.product());
0248
0249     reco::BeamSpot bs;
0250     return theAlgo->buildTrack(theFitter.product(),
0251                                thePropagator.product(),
0252                                algoResults,
0253                                hits,
0254                                theInitialStateForRefitting,
0255                                it->seed(),
0256                                ndof,
0257                                bs,
0258                                theT->seedRef());
0259   }
0260
0261   return false;
0262 }
0263 //----------------------------------------------------------------------------------------
0264 reco::TrackExtra NuclearTrackCorrector::getNewTrackExtra(const AlgoProductCollection& algoResults) {
0265   Trajectory* theTraj = algoResults[0].trajectory;
0266   PropagationDirection seedDir = algoResults[0].pDir;
0267
0268   TrajectoryStateOnSurface outertsos;
0269   TrajectoryStateOnSurface innertsos;
0270   unsigned int innerId, outerId;
0271   if (theTraj->direction() == alongMomentum) {
0272     outertsos = theTraj->lastMeasurement().updatedState();
0273     innertsos = theTraj->firstMeasurement().updatedState();
0274     outerId = theTraj->lastMeasurement().recHit()->geographicalId().rawId();
0275     innerId = theTraj->firstMeasurement().recHit()->geographicalId().rawId();
0276   } else {
0277     outertsos = theTraj->firstMeasurement().updatedState();
0278     innertsos = theTraj->lastMeasurement().updatedState();
0279     outerId = theTraj->firstMeasurement().recHit()->geographicalId().rawId();
0280     innerId = theTraj->lastMeasurement().recHit()->geographicalId().rawId();
0281   }
0282
0283   GlobalPoint v = outertsos.globalParameters().position();
0284   GlobalVector p = outertsos.globalParameters().momentum();
0285   math::XYZVector outmom(p.x(), p.y(), p.z());
0286   math::XYZPoint outpos(v.x(), v.y(), v.z());
0287   v = innertsos.globalParameters().position();
0288   p = innertsos.globalParameters().momentum();
0289   math::XYZVector inmom(p.x(), p.y(), p.z());
0290   math::XYZPoint inpos(v.x(), v.y(), v.z());
0291
0292   return reco::TrackExtra(outpos,
0293                           outmom,
0294                           true,
0295                           inpos,
0296                           inmom,
0297                           true,
0298                           outertsos.curvilinearError(),
0299                           outerId,
0300                           innertsos.curvilinearError(),
0301                           innerId,
0302                           seedDir);
0303 }
0304 //----------------------------------------------------------------------------------------
0305 TrajectoryStateOnSurface NuclearTrackCorrector::getInitialState(const reco::Track* theT,
0306                                                                 TransientTrackingRecHit::RecHitContainer& hits,
0307                                                                 const TrackingGeometry* theG,
0308                                                                 const MagneticField* theMF) {
0309   TrajectoryStateOnSurface theInitialStateForRefitting;
0310   //the starting state is the state closest to the first hit along seedDirection.
0311
0312   //avoiding to use transientTrack, it should be faster;
0313   TrajectoryStateOnSurface innerStateFromTrack = trajectoryStateTransform::innerStateOnSurface(*theT, *theG, theMF);
0314   TrajectoryStateOnSurface outerStateFromTrack = trajectoryStateTransform::outerStateOnSurface(*theT, *theG, theMF);
0315   TrajectoryStateOnSurface initialStateFromTrack =
0316       ((innerStateFromTrack.globalPosition() - hits.front()->globalPosition()).mag2() <
0317        (outerStateFromTrack.globalPosition() - hits.front()->globalPosition()).mag2())
0318           ? innerStateFromTrack
0319           : outerStateFromTrack;
0320
0321   // error is rescaled, but correlation are kept.
0322   initialStateFromTrack.rescaleError(100);
0323   theInitialStateForRefitting = TrajectoryStateOnSurface(initialStateFromTrack.localParameters(),
0324                                                          initialStateFromTrack.localError(),
0325                                                          initialStateFromTrack.surface(),
0326                                                          theMF);
0327   return theInitialStateForRefitting;
0328 }
0329 //----------------------------------------------------------------------------------------
0330 void NuclearTrackCorrector::swap_map(const edm::Handle<TrajectoryCollection>& trajColl,
0331                                      std::map<reco::TrackRef, edm::Ref<TrajectoryCollection> >& result) {
0332   for (unsigned int i = 0; i < trajColl->size(); i++) {
0333     TrajectoryRef InTrajRef(trajColl, i);
0334     reco::TrackRef PrimaryTrackRef = m_TrajToTrackCollection->operator[](InTrajRef);
0335     result[PrimaryTrackRef] = InTrajRef;
0336   }
0337 }
0338
0339 DEFINE_FWK_MODULE(NuclearTrackCorrector);