packages/trackreco/DSTClusterPruning.cc

0001 /*!
0002  * \file DSTClusterPruning.cc
0003  * \author Alex Patton <aopatton@mit.edu>
0004  */
0005
0006 #include "DSTClusterPruning.h"
0007
0008 #include <fun4all/Fun4AllReturnCodes.h>
0009 #include <phool/PHCompositeNode.h>
0010 #include <phool/PHNodeIterator.h>
0011 #include <phool/getClass.h>
0012 #include <trackbase/TrkrCluster.h>
0013 #include <trackbase/TrkrClusterv4.h>
0014 #include <trackbase/TrkrClusterv5.h>
0015 // include new cluster
0016 #include <trackbase/TrkrClusterContainer.h>
0017 #include <trackbase/TrkrClusterHitAssoc.h>
0018 #include <trackbase/TrkrHit.h>
0019 #include <trackbase/TrkrHitSet.h>
0020 #include <trackbase/TrkrHitSetContainer.h>
0021 #include <trackbase/TrkrHitTruthAssoc.h>
0022 #include <trackbase_historic/SvtxTrack.h>
0023 #include <trackbase_historic/SvtxTrackMap.h>
0024 #include <trackbase_historic/TrackSeed.h>
0025
0026 #include <trackbase_historic/SvtxTrack.h>
0027 #include <trackbase_historic/SvtxTrackMap.h>
0028
0029 #include <algorithm>
0030 #include <bitset>
0031 #include <cassert>
0032 #include <iostream>
0033 #include <numeric>
0034 #include <utility>  //for pair
0035
0036 #include <TFile.h>
0037 #include <TLine.h>
0038 #include <TNtuple.h>
0039 #include <TTree.h>
0040 //_____________________________________________________________________
0041
0042 //_____________________________________________________________________
0043 DSTClusterPruning::DSTClusterPruning(const std::string& name)
0044   : SubsysReco(name)
0045 {
0046 }
0047
0048 //_____________________________________________________________________
0049 int DSTClusterPruning::InitRun(PHCompositeNode* topNode)
0050 {
0051   // find DST node
0052   PHNodeIterator iter(topNode);
0053   auto dstNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0054   if (!dstNode)
0055   {
0056     std::cout << "DSTClusterPruning::InitRun - DST Node missing" << std::endl;
0057     return Fun4AllReturnCodes::ABORTEVENT;
0058   }
0059
0060   // get EVAL node
0061   iter = PHNodeIterator(dstNode);
0062   auto evalNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "EVAL"));
0063   if (!evalNode)
0064   {
0065     // create
0066     std::cout << "DSTClusterPruning::InitRun - EVAL node missing - creating" << std::endl;
0067     evalNode = new PHCompositeNode("EVAL");
0068     dstNode->addNode(evalNode);
0069   }
0070
0071   auto svtxNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "SVTX"));
0072   if (!svtxNode)
0073   {
0074     // create
0075     std::cout << "DSTClusterPruning::InitRun - SVTX node missing - creating" << std::endl;
0076     svtxNode = new PHCompositeNode("SVTX");
0077     dstNode->addNode(svtxNode);
0078   }
0079
0080   auto trkrNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "TRKR"));
0081   if (!trkrNode)
0082   {
0083     // create
0084     std::cout << "DSTClusterPruning::InitRun - TRKR node missing - creating" << std::endl;
0085     trkrNode = new PHCompositeNode("TRKR");
0086     dstNode->addNode(trkrNode);
0087   }
0088
0089   // make new cluster container
0090   auto clsNode = findNode::getClass<TrkrClusterContainer>(trkrNode, "TRKR_CLUSTER_SEED");
0091   if (!clsNode)
0092   {
0093     auto newClusterNode = new PHIODataNode<PHObject>(new TrkrClusterContainerv4, "TRKR_CLUSTER_SEED", "PHObject");
0094     trkrNode->addNode(newClusterNode);
0095   }
0096
0097   auto res = load_nodes(topNode);
0098   if (res != Fun4AllReturnCodes::EVENT_OK)
0099   {
0100     return res;
0101   }
0102   if (Verbosity() > 0)
0103   {
0104     std::cout << "Pruner InitRun end" << std::endl;
0105   }
0106   return Fun4AllReturnCodes::EVENT_OK;
0107 }
0108
0109 //_____________________________________________________________________
0110 int DSTClusterPruning::process_event(PHCompositeNode* /*topNode*/)
0111 {
0112   // make topNode run in Init
0113   // Init(topNode);
0114   //  load nodes
0115   if (Verbosity() > 1)
0116   {
0117     std::cout << __FILE__ << "::" << __func__ << "::" << __LINE__ << std::endl;
0118     std::cout << "DSTClusterPruning::process_event" << std::endl;
0119   }
0120
0121   prune_clusters();
0122   if (Verbosity() > 3)
0123   {
0124     print_clusters();
0125   }
0126   return Fun4AllReturnCodes::EVENT_OK;
0127 }
0128
0129 //_____________________________________________________________________
0130 int DSTClusterPruning::load_nodes(PHCompositeNode* topNode)
0131 {
0132   // get necessary nodes
0133   m_track_map = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
0134
0135   m_cluster_map = findNode::getClass<TrkrClusterContainer>(topNode, "TRKR_CLUSTER");
0136
0137   // look for reduced cluster
0138   m_reduced_cluster_map = findNode::getClass<TrkrClusterContainer>(topNode, "TRKR_CLUSTER_SEED");
0139
0140   // m_reduced_track_map = findNode::getClass<SvtxTrackMap>(topNode, "ReducedTrackContainer");
0141   m_track_seed_container = findNode::getClass<TrackSeedContainer>(topNode, "SvtxTrackSeedContainer");
0142   m_tpc_track_seed_container = findNode::getClass<TrackSeedContainer>(topNode, "TpcTrackSeedContainer");
0143   m_silicon_track_seed_container = findNode::getClass<TrackSeedContainer>(topNode, "SiliconTrackSeedContainer");
0144
0145   return Fun4AllReturnCodes::EVENT_OK;
0146 }
0147
0148 //_____________________________________________________________________
0149 void DSTClusterPruning::prune_clusters()
0150 {
0151   // use this to create object that looks through both tracks and clusters and saves into new object
0152   // make sure clusters exist
0153   // make sure tracks exist
0154   if (Verbosity() > 1){
0155     std::cout << "In Prune clusters" << std::endl;
0156     if(!m_cluster_map){
0157       std::cout << "No cluster map" << std::endl;
0158     }
0159     if(!m_reduced_cluster_map){
0160       std::cout << "No reduced cluster map" << std::endl;
0161     }
0162     if(!m_track_seed_container){
0163       std::cout << "No track seed container" << std::endl;
0164     }
0165     if(!m_silicon_track_seed_container){
0166       std::cout << "No silicon track seed container" << std::endl;
0167     }
0168     if(!m_tpc_track_seed_container){
0169       std::cout << "No tpc track seed container" << std::endl;
0170     }
0171   }
0172
0173   if (!(m_cluster_map && m_reduced_cluster_map && m_track_seed_container && m_silicon_track_seed_container && m_tpc_track_seed_container))
0174   {
0175     if (Verbosity() > 1){
0176       std::cout << "Missing container" << std::endl;
0177     }
0178     return;
0179   }
0180   for (const auto& trackseed : *m_track_seed_container)
0181   {
0182     if (!trackseed)
0183     {
0184       std::cout << "No TrackSeed" << std::endl;
0185       continue;
0186     }
0187
0188     unsigned int tpcIndex = trackseed->get_tpc_seed_index();
0189     unsigned int siliconIndex = trackseed->get_silicon_seed_index();
0190
0191     TrackSeed* TPCSeed = nullptr;
0192     if (tpcIndex <= m_tpc_track_seed_container->end() - m_tpc_track_seed_container->begin())
0193     {
0194       TPCSeed = *(m_tpc_track_seed_container->begin() + tpcIndex);
0195     }
0196     TrackSeed* SiliconSeed = nullptr;
0197     if (siliconIndex <= m_silicon_track_seed_container->end() - m_silicon_track_seed_container->begin())
0198     {
0199       SiliconSeed = *(m_silicon_track_seed_container->begin() + siliconIndex);
0200     }
0201     if (TPCSeed)
0202     {
0203       if (Verbosity() > 1)
0204       {
0205         std::cout << "We are about to loop over cluster keys in TPC Seed" << std::endl;
0206         TPCSeed->identify();
0207       }
0208       for (auto key_iter = TPCSeed->begin_cluster_keys(); key_iter != TPCSeed->end_cluster_keys(); ++key_iter)
0209       {
0210         const auto& cluster_key = *key_iter;
0211         auto cluster = m_cluster_map->findCluster(cluster_key);
0212         if (!cluster)
0213         {
0214           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0215           continue;
0216         }
0217         if (!m_reduced_cluster_map->findCluster(cluster_key))
0218         {
0219           m_cluster = new TrkrClusterv5();
0220           m_cluster->CopyFrom(cluster);
0221           m_reduced_cluster_map->addClusterSpecifyKey(cluster_key, m_cluster);
0222         }
0223       }
0224     }
0225
0226     if (SiliconSeed)
0227     {
0228       if (Verbosity() > 1)
0229       {
0230         std::cout << "We are about to loop over cluster keys in Silicon Seed" << std::endl;
0231         SiliconSeed->identify();
0232       }
0233       for (auto key_iter = SiliconSeed->begin_cluster_keys(); key_iter != SiliconSeed->end_cluster_keys(); ++key_iter)
0234       {
0235         const auto& cluster_key = *key_iter;
0236         auto cluster = m_cluster_map->findCluster(cluster_key);
0237         if (!cluster)
0238         {
0239           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0240           continue;
0241         }
0242         if (!m_reduced_cluster_map->findCluster(cluster_key))
0243         {
0244           m_cluster = new TrkrClusterv5();
0245           m_cluster->CopyFrom(cluster);
0246           m_reduced_cluster_map->addClusterSpecifyKey(cluster_key, m_cluster);
0247         }
0248       }
0249     }
0250   }
0251 }
0252
0253 // fill original clusters
0254 //_____________________________________________________________________
0255 void DSTClusterPruning::fill_clusters()
0256 {
0257   // use this to create object that looks through both tracks and clusters and saves into new object
0258   // make sure clusters exist
0259   // if(!(m_cluster_map&&m_hitsetcontainer&&m_container)) return;
0260   // make sure tracks exist
0261   if (!(m_cluster_map && m_reduced_cluster_map && m_track_seed_container && m_silicon_track_seed_container && m_tpc_track_seed_container))
0262   {
0263     return;
0264   }
0265
0266   for (const auto& trackseed : *m_track_seed_container)
0267   {
0268     if (!trackseed)
0269     {
0270       continue;
0271     }
0272
0273     unsigned int tpcIndex = trackseed->get_tpc_seed_index();
0274     unsigned int siliconIndex = trackseed->get_silicon_seed_index();
0275
0276     TrackSeed* TPCSeed = nullptr;
0277     if (tpcIndex <= m_tpc_track_seed_container->end() - m_tpc_track_seed_container->begin())
0278     {
0279       TPCSeed = *(m_tpc_track_seed_container->begin() + tpcIndex);
0280     }
0281     TrackSeed* SiliconSeed = nullptr;
0282     if (siliconIndex <= m_silicon_track_seed_container->end() - m_silicon_track_seed_container->begin())
0283     {
0284       SiliconSeed = *(m_silicon_track_seed_container->begin() + siliconIndex);
0285     }
0286
0287     if (TPCSeed)
0288     {
0289       for (auto key_iter = TPCSeed->begin_cluster_keys(); key_iter != TPCSeed->end_cluster_keys(); ++key_iter)
0290       {
0291         const auto& cluster_key = *key_iter;
0292         auto cluster = m_reduced_cluster_map->findCluster(cluster_key);
0293         if (!cluster)
0294         {
0295           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0296           continue;
0297         }
0298         if (!m_cluster_map->findCluster(cluster_key))
0299         {
0300           m_cluster = new TrkrClusterv5();
0301           m_cluster->CopyFrom(cluster);
0302           m_cluster_map->addClusterSpecifyKey(cluster_key, m_cluster);
0303         }
0304       }
0305     }
0306
0307     if (SiliconSeed)
0308     {
0309       // std::cout << "We are about to loop over cluster keys in Silicon Seed" << std::endl;
0310       // SiliconSeed->identify();
0311       for (auto key_iter = SiliconSeed->begin_cluster_keys(); key_iter != SiliconSeed->end_cluster_keys(); ++key_iter)
0312       {
0313         const auto& cluster_key = *key_iter;
0314         auto cluster = m_reduced_cluster_map->findCluster(cluster_key);
0315         if (!cluster)
0316         {
0317           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0318           continue;
0319         }
0320         if (!m_cluster_map->findCluster(cluster_key))
0321         {
0322           m_cluster = new TrkrClusterv5();
0323           m_cluster->CopyFrom(cluster);
0324           m_cluster_map->addClusterSpecifyKey(cluster_key, m_cluster);
0325
0326           // m_reduced_cluster_map->addClusterSpecifyKey(cluster_key, cluster);
0327         }
0328       }
0329     }
0330   }
0331 }
0332
0333 // print clusters to test
0334 void DSTClusterPruning::print_clusters()
0335 {
0336   // make sure tracks exist
0337   if (!(m_track_map && m_cluster_map && m_reduced_cluster_map))
0338   {
0339     return;
0340   }
0341
0342   for (const auto& trackpair : *m_track_map)
0343   {
0344     std::cout << "start of loop" << "\n";
0345
0346     // unsigned int key = trackpair.first;
0347     const auto track = trackpair.second;
0348
0349     TrackSeed* TPCSeed = track->get_tpc_seed();
0350     TrackSeed* SiliconSeed = track->get_silicon_seed();
0351
0352     if (TPCSeed)
0353     {
0354       std::cout << "We are about to loop over cluster keys in TPC Seed" << std::endl;
0355       TPCSeed->identify();
0356       for (auto key_iter = TPCSeed->begin_cluster_keys(); key_iter != TPCSeed->end_cluster_keys(); ++key_iter)
0357       {
0358         const auto& cluster_key = *key_iter;
0359         auto cluster = m_cluster_map->findCluster(cluster_key);
0360         auto reducedcluster = m_reduced_cluster_map->findCluster(cluster_key);
0361         if (!cluster)
0362         {
0363           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0364           continue;
0365         }
0366         if (!reducedcluster)
0367         {
0368           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find reducedcluster for key " << cluster_key << std::endl;
0369           continue;
0370         }
0371         std::cout << "ClusterKey: " << cluster_key << std::endl;
0372         std::cout << "Cluster map Cluster Local X: " << cluster->getLocalX() << " Local Y: " << cluster->getLocalY() << std::endl;
0373         std::cout << "Reduced map Cluster Local X: " << cluster->getLocalX() << " Local Y: " << cluster->getLocalY() << std::endl;
0374       }
0375     }
0376
0377     if (SiliconSeed)
0378     {
0379       std::cout << "We are about to loop over cluster keys in Silicon Seed" << std::endl;
0380       SiliconSeed->identify();
0381       for (auto key_iter = SiliconSeed->begin_cluster_keys(); key_iter != SiliconSeed->end_cluster_keys(); ++key_iter)
0382       {
0383         const auto& cluster_key = *key_iter;
0384         auto cluster = m_cluster_map->findCluster(cluster_key);
0385
0386         auto reducedcluster = m_reduced_cluster_map->findCluster(cluster_key);
0387         if (!cluster)
0388         {
0389           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find cluster for key " << cluster_key << std::endl;
0390           continue;
0391         }
0392         if (!reducedcluster)
0393         {
0394           std::cout << "DSTClusterPruning::evaluate_tracks - unable to find reducedcluster for key " << cluster_key << std::endl;
0395           continue;
0396         }
0397         std::cout << "ClusterKey: " << cluster_key << std::endl;
0398         std::cout << "Cluster map Cluster Local X: " << cluster->getLocalX() << " Local Y: " << cluster->getLocalY() << std::endl;
0399         std::cout << "Reduced map Cluster Local X: " << cluster->getLocalX() << " Local Y: " << cluster->getLocalY() << std::endl;
0400       }
0401     }
0402
0403     std::cout << "end of loop" << "\n";
0404   }
0405 }