g4simulation/g4eval/SvtxTruthRecoTableEval.cc

0001
0002 #include "SvtxTruthRecoTableEval.h"
0003 #include "SvtxEvalStack.h"
0004 #include "SvtxTrackEval.h"
0005 #include "SvtxTruthEval.h"
0006
0007 #include <fun4all/Fun4AllReturnCodes.h>
0008 #include <phool/PHCompositeNode.h>
0009 #include <phool/PHDataNode.h>
0010 #include <phool/PHNode.h>
0011 #include <phool/PHNodeIterator.h>
0012 #include <phool/PHObject.h>
0013 #include <phool/getClass.h>
0014 #include <phool/phool.h>
0015
0016 #include <g4main/PHG4Particle.h>
0017 #include <g4main/PHG4TruthInfoContainer.h>
0018
0019 #include <trackbase_historic/PHG4ParticleSvtxMap_v1.h>
0020 #include <trackbase_historic/SvtxPHG4ParticleMap_v1.h>
0021 #include <trackbase_historic/SvtxTrack.h>
0022 #include <trackbase_historic/SvtxTrackMap.h>
0023
0024 #include <CLHEP/Vector/ThreeVector.h>
0025
0026
0027 //____________________________________________________________________________..
0028 SvtxTruthRecoTableEval::SvtxTruthRecoTableEval(const std::string &name)
0029   : SubsysReco(name)
0030 {
0031 }
0032
0033 //____________________________________________________________________________..
0034 SvtxTruthRecoTableEval::~SvtxTruthRecoTableEval() = default;
0035
0036 //____________________________________________________________________________..
0037 int SvtxTruthRecoTableEval::Init(PHCompositeNode * /*unused*/)
0038 {
0039   return Fun4AllReturnCodes::EVENT_OK;
0040 }
0041
0042 //____________________________________________________________________________..
0043 int SvtxTruthRecoTableEval::InitRun(PHCompositeNode *topNode)
0044 {
0045   if (createNodes(topNode) != Fun4AllReturnCodes::EVENT_OK)
0046   {
0047     return Fun4AllReturnCodes::ABORTEVENT;
0048   }
0049
0050   return Fun4AllReturnCodes::EVENT_OK;
0051 }
0052
0053 //____________________________________________________________________________..
0054 int SvtxTruthRecoTableEval::process_event(PHCompositeNode *topNode)
0055 {
0056   if (!m_svtxevalstack)
0057   {
0058     m_svtxevalstack = std::make_unique<SvtxEvalStack>(topNode);
0059     m_svtxevalstack->set_strict(false);
0060     m_svtxevalstack->set_verbosity(Verbosity());
0061     m_svtxevalstack->set_use_initial_vertex(true);
0062     m_svtxevalstack->set_use_genfit_vertex(false);
0063     m_svtxevalstack->next_event(topNode);
0064   }
0065   else
0066   {
0067     m_svtxevalstack->next_event(topNode);
0068   }
0069
0070   if (Verbosity() > 1)
0071   {
0072     std::cout << "Fill truth map " << std::endl;
0073   }
0074   fillTruthMap(topNode);
0075
0076   if (Verbosity() > 1)
0077   {
0078     std::cout << "Fill reco map " << std::endl;
0079   }
0080   fillRecoMap(topNode);
0081
0082   return Fun4AllReturnCodes::EVENT_OK;
0083 }
0084
0085 //____________________________________________________________________________..
0086 int SvtxTruthRecoTableEval::ResetEvent(PHCompositeNode * /*unused*/)
0087 {
0088   if (Verbosity() > 0)
0089   {
0090     std::cout << "Truth track map " << std::endl;
0091     m_truthMap->identify();
0092     std::cout << std::endl
0093               << "Reco track map " << std::endl;
0094     m_recoMap->identify();
0095   }
0096   return Fun4AllReturnCodes::EVENT_OK;
0097 }
0098
0099 //____________________________________________________________________________..
0100 int SvtxTruthRecoTableEval::End(PHCompositeNode * /*unused*/)
0101 {
0102   return Fun4AllReturnCodes::EVENT_OK;
0103 }
0104
0105 void SvtxTruthRecoTableEval::fillTruthMap(PHCompositeNode *topNode)
0106 {
0107   PHG4TruthInfoContainer *truthinfo = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
0108   assert(truthinfo);
0109
0110   SvtxTrackEval *trackeval = m_svtxevalstack->get_track_eval();
0111   trackeval->set_verbosity(Verbosity());
0112   assert(trackeval);
0113
0114   PHG4TruthInfoContainer::ConstRange range = truthinfo->GetParticleRange();
0115   if (m_scanForPrimaries)
0116   {
0117     range = truthinfo->GetPrimaryParticleRange();
0118   }
0119
0120   for (auto iter = range.first; iter != range.second; ++iter)
0121   {
0122     PHG4Particle *g4particle = iter->second;
0123
0124     const double momentum = CLHEP::
0125                                 Hep3Vector(g4particle->get_px(), g4particle->get_py(), g4particle->get_pz())
0126                                     .mag();
0127
0128     // only record particle above minimal momentum requirement.
0129     if (momentum < m_minMomentumTruthMap)
0130     {
0131       continue;
0132     }
0133
0134     int gtrackID = g4particle->get_track_id();
0135     const std::set<SvtxTrack *> &alltracks = trackeval->all_tracks_from(g4particle);
0136
0137     // not to record zero associations
0138     if (alltracks.empty())
0139     {
0140       continue;
0141     }
0142
0143     PHG4ParticleSvtxMap::WeightedRecoTrackMap recomap;
0144
0145     for (auto *track : alltracks)
0146     {
0147       /// We fill the map with a key corresponding to the ncluster contribution.
0148       /// This weight could in principle be anything we choose
0149       float clusCont = trackeval->get_nclusters_contribution(track, g4particle);
0150
0151       auto iterator = recomap.find(clusCont);
0152       if (iterator == recomap.end())
0153       {
0154         std::set<unsigned int> dumset;
0155         dumset.insert(track->get_id());
0156         recomap.insert(std::make_pair(clusCont, dumset));
0157       }
0158       else
0159       {
0160         iterator->second.insert(track->get_id());
0161       }
0162     }
0163
0164     if (Verbosity() > 1)
0165     {
0166       std::cout << " Inserting gtrack id " << gtrackID << " with map size " << recomap.size() << std::endl;
0167     }
0168
0169     m_truthMap->insert(gtrackID, recomap);
0170   }
0171
0172   m_truthMap->setProcessed(true);
0173 }
0174
0175 void SvtxTruthRecoTableEval::fillRecoMap(PHCompositeNode *topNode)
0176 {
0177   SvtxTrackMap *trackMap = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
0178
0179   assert(trackMap);
0180
0181   SvtxTrackEval *trackeval = m_svtxevalstack->get_track_eval();
0182   assert(trackeval);
0183
0184   for (const auto &[key, track] : *trackMap)
0185   {
0186     const std::set<PHG4Particle *> &allparticles = trackeval->all_truth_particles(track);
0187     SvtxPHG4ParticleMap::WeightedTruthTrackMap truthmap;
0188     for (PHG4Particle *g4particle : allparticles)
0189     {
0190       float clusCont = trackeval->get_nclusters_contribution(track, g4particle);
0191       auto iterator = truthmap.find(clusCont);
0192       if (iterator == truthmap.end())
0193       {
0194         std::set<int> dumset;
0195         dumset.insert(g4particle->get_track_id());
0196         truthmap.insert(std::make_pair(clusCont, dumset));
0197       }
0198       else
0199       {
0200         iterator->second.insert(g4particle->get_track_id());
0201       }
0202     }
0203     if (Verbosity() > 1)
0204     {
0205       std::cout << " Inserting track id " << key << " with truth map size " << truthmap.size() << std::endl;
0206     }
0207     m_recoMap->insert(key, truthmap);
0208   }
0209
0210   m_recoMap->setProcessed(true);
0211 }
0212
0213 int SvtxTruthRecoTableEval::createNodes(PHCompositeNode *topNode)
0214 {
0215   PHNodeIterator iter(topNode);
0216
0217   PHCompositeNode *dstNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "DST"));
0218
0219   if (!dstNode)
0220   {
0221     std::cerr << "DST node is missing, quitting" << std::endl;
0222     throw std::runtime_error("Failed to find DST node in SvtxTruthRecoTableEval::createNodes");
0223   }
0224
0225   PHCompositeNode *svtxNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "SVTX"));
0226
0227   if (!svtxNode)
0228   {
0229     svtxNode = new PHCompositeNode("SVTX");
0230     dstNode->addNode(svtxNode);
0231   }
0232
0233   m_truthMap = findNode::getClass<PHG4ParticleSvtxMap>(topNode, "PHG4ParticleSvtxMap");
0234   if (!m_truthMap)
0235   {
0236     m_truthMap = new PHG4ParticleSvtxMap_v1;
0237     PHIODataNode<PHObject> *truthNode =
0238         new PHIODataNode<PHObject>(m_truthMap, "PHG4ParticleSvtxMap", "PHObject");
0239     svtxNode->addNode(truthNode);
0240   }
0241
0242   m_recoMap = findNode::getClass<SvtxPHG4ParticleMap>(topNode, "SvtxPHG4ParticleMap");
0243   if (!m_recoMap)
0244   {
0245     m_recoMap = new SvtxPHG4ParticleMap_v1;
0246     PHIODataNode<PHObject> *recoNode =
0247         new PHIODataNode<PHObject>(m_recoMap, "SvtxPHG4ParticleMap", "PHObject");
0248     svtxNode->addNode(recoNode);
0249   }
0250
0251   return Fun4AllReturnCodes::EVENT_OK;
0252 }