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