g4simulation/g4tracking/TruthClusterizerBase.cc

0001 #include "TruthClusterizerBase.h"
0002
0003 #include <g4tracking/TrkrTruthTrack.h>
0004 #include <g4tracking/TrkrTruthTrackContainer.h>
0005 #include <g4tracking/TrkrTruthTrackContainerv1.h>
0006
0007 #include <trackbase/TrkrClusterContainer.h>  // for TrkrClusterContainer
0008 #include <trackbase/TrkrClusterContainerv4.h>
0009 #include <trackbase/TrkrDefs.h>
0010 #include <trackbase/TrkrHit.h>
0011 #include <trackbase/TrkrHitSet.h>
0012 #include <trackbase/TrkrHitSetContainer.h>  // for TrkrHitSetContainer
0013 #include <trackbase/TrkrHitSetContainerv1.h>
0014 #include <trackbase/TrkrHitv2.h>  // for TrkrHit
0015
0016 #include <g4main/PHG4Hit.h>
0017 #include <g4main/PHG4TruthInfoContainer.h>
0018
0019 #include <phool/PHCompositeNode.h>
0020 #include <phool/PHIODataNode.h>
0021 #include <phool/PHNode.h>
0022 #include <phool/PHNodeIterator.h>
0023 #include <phool/PHObject.h>  // for PHObject
0024 #include <phool/getClass.h>
0025 #include <phool/phool.h>
0026
0027 #include <boost/format.hpp>
0028
0029 #include <cassert>
0030 #include <iostream>
0031 #include <string>   // for operator<<
0032 #include <utility>  // for pair
0033 #include <vector>
0034
0035 TruthClusterizerBase::TruthClusterizerBase()
0036   : m_hits{new TrkrHitSetContainerv1}
0037 {
0038 }
0039
0040 void TruthClusterizerBase::init_clusterizer_base(PHCompositeNode*& _topNode, int _verbosity)
0041 {
0042   m_topNode = _topNode;
0043   m_verbosity = _verbosity;
0044   PHNodeIterator iter(m_topNode);
0045   auto dstNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0046   assert(dstNode);
0047   PHNodeIterator dstiter(dstNode);
0048   auto DetNode = dynamic_cast<PHCompositeNode*>(dstiter.findFirst("PHCompositeNode", "TRKR"));
0049   if (!DetNode)
0050   {
0051     DetNode = new PHCompositeNode("TRKR");
0052     dstNode->addNode(DetNode);
0053   }
0054
0055   m_truthtracks = findNode::getClass<TrkrTruthTrackContainer>(m_topNode, "TRKR_TRUTHTRACKCONTAINER");
0056   if (!m_truthtracks)
0057   {
0058     m_truthtracks = new TrkrTruthTrackContainerv1();
0059     auto newNode = new PHIODataNode<PHObject>(m_truthtracks, "TRKR_TRUTHTRACKCONTAINER", "PHObject");
0060     DetNode->addNode(newNode);
0061   }
0062
0063   m_clusters = findNode::getClass<TrkrClusterContainer>(m_topNode, "TRKR_TRUTHCLUSTERCONTAINER");
0064   if (!m_clusters)
0065   {
0066     m_clusters = new TrkrClusterContainerv4;
0067     auto newNode = new PHIODataNode<PHObject>(m_clusters, "TRKR_TRUTHCLUSTERCONTAINER", "PHObject");
0068     DetNode->addNode(newNode);
0069   }
0070
0071   m_truthinfo = findNode::getClass<PHG4TruthInfoContainer>(m_topNode, "G4TruthInfo");
0072   if (!m_truthinfo)
0073   {
0074     std::cout << PHWHERE << " PHG4TruthInfoContainer node not found on node tree" << std::endl;
0075     assert(m_truthinfo);
0076   }
0077 }
0078
0079 TruthClusterizerBase::~TruthClusterizerBase()
0080 {
0081   delete m_hits;
0082 }
0083
0084 void TruthClusterizerBase::check_g4hit_status(PHG4Hit* hit)
0085 {
0086   int new_trkid = (hit == nullptr) ? -1 : hit->get_trkid();
0087   m_is_new_track = (new_trkid != m_trkid);
0088   if (m_verbosity > 5)
0089   {
0090     std::cout << PHWHERE << std::endl
0091               << " -> Checking status of PHG4Hit. Track id(" << new_trkid << ")" << std::endl;
0092   }
0093   if (!m_is_new_track)
0094   {
0095     return;
0096   }
0097
0098   m_trkid = new_trkid;  // although m_current_track won't catch up until update_track is called
0099   m_was_emb = m_is_emb;
0100   m_is_emb = m_truthinfo->isEmbeded(m_trkid);
0101 }
0102
0103 // to call if m_was_emb=true and after clustering
0104 void TruthClusterizerBase::transfer_clusters(TrkrClusterContainer* pass_clusters)
0105 {
0106   m_hits->Reset();  // clear out the old  hits
0107   for (auto hitsetkey : pass_clusters->getHitSetKeys())
0108   {
0109     m_hitsetkey_cnt.try_emplace(hitsetkey, 0);
0110     unsigned int& cnt = m_hitsetkey_cnt[hitsetkey];
0111     auto range = pass_clusters->getClusters(hitsetkey);
0112     for (auto cluster = range.first; cluster != range.second; ++cluster)
0113     {
0114       auto ckey = TrkrDefs::genClusKey(hitsetkey, cnt);
0115       m_clusters->addClusterSpecifyKey(ckey, cluster->second);
0116       m_current_track->addCluster(ckey);
0117       ++cnt;
0118     }
0119   }
0120   m_was_emb = false;
0121 }
0122
0123 // if m_is_new_track
0124 void TruthClusterizerBase::update_track()
0125 {
0126   m_current_track = m_is_emb ? m_truthtracks->getTruthTrack(m_trkid, m_truthinfo) : nullptr;
0127 }
0128
0129 void TruthClusterizerBase::addhitset(
0130     TrkrDefs::hitsetkey hitsetkey,
0131     TrkrDefs::hitkey hitkey,
0132     float neffelectrons)
0133 {
0134   if (!m_is_emb)
0135   {
0136     return;
0137   }
0138   TrkrHitSetContainer::Iterator hitsetit = m_hits->findOrAddHitSet(hitsetkey);
0139   // See if this hit already exists
0140   TrkrHit* hit = nullptr;
0141   hit = hitsetit->second->getHit(hitkey);
0142   if (!hit)
0143   {
0144     // create a new one
0145     hit = new TrkrHitv2();
0146     hitsetit->second->addHitSpecificKey(hitkey, hit);
0147   }
0148   // Either way, add the energy to it  -- adc values will be added at digitization
0149   hit->addEnergy(neffelectrons);
0150 }
0151
0152 void TruthClusterizerBase::print_clusters(int nclusprint)
0153 {
0154   std::cout << PHWHERE << ": content of clusters " << std::endl;
0155   auto& tmap = m_truthtracks->getMap();
0156   std::cout << " Number of tracks: " << tmap.size() << std::endl;
0157   for (auto& _pair : tmap)
0158   {
0159     auto& track = _pair.second;
0160
0161     std::cout << (boost::format("id(%2i) phi:eta:pt(%5.2f:%5.2f:%5.2f) nclusters(%d) ") % (int) track->getTrackid() % track->getPhi() % track->getPseudoRapidity() % track->getPt() % track->getClusters().size()).str();
0162     if (m_verbosity <= 10)
0163     {
0164       std::cout << std::endl;
0165     }
0166     else
0167     {
0168       int nclus = 0;
0169       for (auto cluskey : track->getClusters())
0170       {
0171         std::cout << " "
0172                   << ((int) TrkrDefs::getHitSetKeyFromClusKey(cluskey)) << ":index(" << ((int) TrkrDefs::getClusIndex(cluskey)) << ")";
0173         ++nclus;
0174         if (nclusprint > 0 && nclus >= nclusprint)
0175         {
0176           std::cout << " ... ";
0177           break;
0178         }
0179       }
0180     }
0181   }
0182   std::cout << PHWHERE << " ----- end of clusters " << std::endl;
0183 }