packages/trackreco/PHTruthTrackSeeding.cc

0001 #include "PHTruthTrackSeeding.h"
0002
0003 #include <trackbase_historic/TrackSeed.h>
0004 #include <trackbase_historic/TrackSeedContainer.h>
0005 #include <trackbase_historic/TrackSeed_FastSim_v2.h>
0006
0007 #include <trackbase/TrkrCluster.h>
0008 #include <trackbase/TrkrClusterContainer.h>
0009
0010 #include <globalvertex/SvtxVertex.h>
0011 #include <globalvertex/SvtxVertexMap.h>
0012
0013 #include <trackbase/TrkrHitTruthAssoc.h>
0014 #include <trackbase_historic/ActsTransformations.h>
0015
0016 #include <fun4all/Fun4AllReturnCodes.h>
0017
0018 #include <g4eval/SvtxClusterEval.h>
0019 #include <g4eval/SvtxEvalStack.h>
0020 #include <g4eval/SvtxEvaluator.h>
0021 #include <g4eval/SvtxTrackEval.h>
0022
0023 #include <g4main/PHG4Hit.h>  // for PHG4Hit
0024 #include <g4main/PHG4HitContainer.h>
0025 #include <g4main/PHG4HitDefs.h>   // for keytype
0026 #include <g4main/PHG4Particle.h>  // for PHG4Particle
0027 #include <g4main/PHG4TruthInfoContainer.h>
0028 #include <g4main/PHG4VtxPoint.h>
0029
0030 #include <phool/PHCompositeNode.h>
0031 #include <phool/PHIODataNode.h>
0032 #include <phool/PHNode.h>
0033 #include <phool/PHNodeIterator.h>
0034 #include <phool/PHObject.h>
0035 #include <phool/PHRandomSeed.h>
0036 #include <phool/getClass.h>
0037 #include <phool/phool.h>
0038
0039 #include <trackbase/TrkrCluster.h>
0040 #include <trackbase/TrkrClusterContainer.h>
0041 #include <trackbase/TrkrClusterCrossingAssoc.h>
0042 #include <trackbase/TrkrHitTruthAssoc.h>
0043
0044 #include <trackbase_historic/SvtxTrackSeed_v1.h>
0045 #include <trackbase_historic/TrackSeedContainer_v1.h>
0046 #include <trackbase_historic/TrackSeed_v2.h>
0047 #include <trackbase_historic/TrackSeedHelper.h>
0048
0049 #include <TDatabasePDG.h>
0050 #include <TParticlePDG.h>
0051
0052 #include <gsl/gsl_randist.h>
0053 #include <gsl/gsl_rng.h>  // for gsl_rng_alloc
0054
0055 #include <cassert>
0056 #include <cmath>
0057 #include <cstdlib>   // for exit
0058 #include <iostream>  // for operator<<, std::endl
0059 #include <map>       // for multimap, map<>::c...
0060 #include <memory>
0061 #include <set>
0062 #include <utility>  // for pair
0063
0064 class PHCompositeNode;
0065
0066 namespace
0067 {
0068   template <class T>
0069   inline constexpr T square(const T& x)
0070   {
0071     return x * x;
0072   }
0073 }  // namespace
0074
0075 PHTruthTrackSeeding::PHTruthTrackSeeding(const std::string& name)
0076   : PHTrackSeeding(name)
0077   , _clustereval(nullptr)
0078 {
0079   // initialize random generator
0080   const uint seed = PHRandomSeed();
0081   m_rng.reset(gsl_rng_alloc(gsl_rng_mt19937));
0082   gsl_rng_set(m_rng.get(), seed);
0083 }
0084
0085 int PHTruthTrackSeeding::Setup(PHCompositeNode* topNode)
0086 {
0087   std::cout << "Enter PHTruthTrackSeeding:: Setup" << std::endl;
0088
0089   int ret = PHTrackSeeding::Setup(topNode);
0090   if (ret != Fun4AllReturnCodes::EVENT_OK)
0091   {
0092     return ret;
0093   }
0094
0095   ret = GetNodes(topNode);
0096   if (ret != Fun4AllReturnCodes::EVENT_OK)
0097   {
0098     return ret;
0099   }
0100
0101   ret = CreateNodes(topNode);
0102   if (ret != Fun4AllReturnCodes::EVENT_OK)
0103   {
0104     return ret;
0105   }
0106
0107   _clustereval = new SvtxClusterEval(topNode);
0108   _clustereval->do_caching(true);
0109   return Fun4AllReturnCodes::EVENT_OK;
0110 }
0111
0112 int PHTruthTrackSeeding::Process(PHCompositeNode* topNode)
0113 {
0114   _clustereval->next_event(topNode);
0115   _track_map_combined->Clear();
0116
0117   using TrkClustersMap = std::map<int, std::set<TrkrCluster*> >;
0118   TrkClustersMap m_trackID_clusters;
0119
0120   std::vector<TrkrDefs::cluskey> ClusterKeyListSilicon;
0121   std::vector<TrkrDefs::cluskey> ClusterKeyListTpc;
0122
0123   PHG4TruthInfoContainer::ConstRange range = m_g4truth_container->GetPrimaryParticleRange();
0124   for (PHG4TruthInfoContainer::ConstIterator iter = range.first;
0125        iter != range.second;
0126        ++iter)
0127   {
0128     ClusterKeyListSilicon.clear();
0129     ClusterKeyListTpc.clear();
0130     PHG4Particle* g4particle = iter->second;
0131
0132     if (g4particle == nullptr)
0133     {
0134       std::cout << __PRETTY_FUNCTION__ << " - validity check failed: missing truth particle" << std::endl;
0135       exit(1);
0136     }
0137
0138     const float gtrackID = g4particle->get_track_id();
0139
0140     // momentum cut-off
0141     if (_min_momentum > 0)
0142     {
0143       const double monentum2 =
0144           g4particle->get_px() * g4particle->get_px() +
0145           g4particle->get_py() * g4particle->get_py() +
0146           g4particle->get_pz() * g4particle->get_pz();
0147
0148       if (monentum2 < _min_momentum * _min_momentum)
0149       {
0150         if (Verbosity() >= 3)
0151         {
0152           std::cout << __PRETTY_FUNCTION__ << " ignore low momentum particle " << gtrackID << std::endl;
0153           g4particle->identify();
0154         }
0155         continue;
0156       }
0157     }
0158
0159     for (unsigned int layer = _min_layer; layer < _max_layer; layer++)
0160     {
0161       TrkrDefs::cluskey cluskey = _clustereval->best_cluster_by_nhit(gtrackID, layer);
0162       if (cluskey != 0)
0163       {
0164         unsigned int trkrid = TrkrDefs::getTrkrId(cluskey);
0165         if (trkrid == TrkrDefs::mvtxId || trkrid == TrkrDefs::inttId)
0166         {
0167           ClusterKeyListSilicon.push_back(cluskey);
0168         }
0169         else
0170         {
0171           ClusterKeyListTpc.push_back(cluskey);
0172         }
0173       }
0174     }
0175
0176     unsigned int nsi = ClusterKeyListSilicon.size();
0177     unsigned int ntpc = ClusterKeyListTpc.size();
0178
0179     if (nsi + ntpc < _min_clusters_per_track)
0180     {
0181       continue;
0182     }
0183
0184     if (nsi > 0)
0185     {
0186       buildTrackSeed(ClusterKeyListSilicon, g4particle, _track_map_silicon);
0187     }
0188     if (ntpc > 0)
0189     {
0190       buildTrackSeed(ClusterKeyListTpc, g4particle, _track_map);
0191     }
0192
0193     if (nsi > 0 && ntpc > 0)
0194     {
0195       // Create the SvtxTrackSeed for the full track
0196       auto track = std::make_unique<SvtxTrackSeed_v1>();
0197       /// The ids will by definition be the last entry in the container because
0198       /// the seeds were just added
0199       track->set_tpc_seed_index(_track_map->size() - 1);
0200       track->set_silicon_seed_index(_track_map_silicon->size() - 1);
0201
0202       _track_map_combined->insert(track.get());
0203     }
0204   }
0205
0206   if (Verbosity() > 0)
0207   {
0208     // loop over the assembled tracks
0209     for (unsigned int phtrk_iter = 0;
0210          phtrk_iter < _track_map_combined->size();
0211          ++phtrk_iter)
0212     {
0213       auto seed = _track_map_combined->get(phtrk_iter);
0214       if (!seed)
0215       {
0216         continue;
0217       }
0218
0219       auto tpc_index = seed->get_tpc_seed_index();
0220       auto silicon_index = seed->get_silicon_seed_index();
0221
0222       std::cout << "SvtxSeedTrack: " << phtrk_iter
0223                 << " tpc_index " << tpc_index
0224                 << " silicon_index " << silicon_index
0225                 << std::endl;
0226
0227       std::cout << " ----------  silicon tracklet " << silicon_index << std::endl;
0228       auto silicon_tracklet = _track_map_silicon->get(silicon_index);
0229       if (!silicon_tracklet)
0230       {
0231         continue;
0232       }
0233       silicon_tracklet->identify();
0234
0235       std::cout << " ---------- tpc tracklet " << tpc_index << std::endl;
0236       auto tpc_tracklet = _track_map->get(tpc_index);
0237       if (!tpc_tracklet)
0238       {
0239         continue;
0240       }
0241       tpc_tracklet->identify();
0242     }
0243   }
0244
0245   //==================================
0246
0247   return Fun4AllReturnCodes::EVENT_OK;
0248 }
0249
0250 void PHTruthTrackSeeding::buildTrackSeed(const std::vector<TrkrDefs::cluskey>& clusters, PHG4Particle* g4particle, TrackSeedContainer* container)
0251 {
0252   // This method is called separately for silicon and tpc seeds
0253
0254   auto track = std::make_unique<TrackSeed_FastSim_v2>();
0255   bool silicon = false;
0256   bool tpc = false;
0257   for (const auto& cluskey : clusters)
0258   {
0259     if (TrkrDefs::getTrkrId(cluskey) == TrkrDefs::TrkrId::mvtxId ||
0260         TrkrDefs::getTrkrId(cluskey) == TrkrDefs::TrkrId::inttId)
0261     {
0262       silicon = true;
0263     }
0264     if (TrkrDefs::getTrkrId(cluskey) == TrkrDefs::TrkrId::tpcId)
0265     {
0266       tpc = true;
0267     }
0268     track->insert_cluster_key(cluskey);
0269   }
0270
0271   const auto particle = TDatabasePDG::Instance()->GetParticle(g4particle->get_pid());
0272   int charge = 1;
0273   if (particle)
0274   {
0275     if (particle->Charge() < 0)
0276     {
0277       charge = -1;
0278     }
0279   }
0280
0281   auto random1 = gsl_ran_flat(m_rng.get(), 0.95, 1.05);
0282   float px = g4particle->get_px() * random1;
0283   float py = g4particle->get_py() * random1;
0284   float pz = g4particle->get_pz() * random1;
0285   const auto g4vertex = m_g4truth_container->GetVtx(g4particle->get_vtx_id());
0286   auto random2 = gsl_ran_flat(m_rng.get(), -0.02, 0.02);
0287   float x = g4vertex->get_x() + random2;
0288   float y = g4vertex->get_y() + random2;
0289   float z = g4vertex->get_z() + random2;
0290
0291   float pt = std::sqrt(px * px + py * py);
0292   float phi = std::atan2(py, px);
0293   float R = 100 * pt / (0.3 * 1.4);
0294   float theta = std::atan2(pt, pz);
0295   if (theta < 0)
0296   {
0297     theta += M_PI;
0298   }
0299   if (theta > M_PI)
0300   {
0301     theta -= M_PI;
0302   }
0303
0304   float eta = -log(tan(theta / 2.));
0305
0306   // We have two equations, phi = atan2(-(X0-x),y-Y0) and
0307   // R^2 = (x-X0)^2 + (y-Y0)^2. Solve for X0 and Y0 knowing R and phi
0308   const float tanphisq = square(std::tan(phi));
0309   const float a = tanphisq + 1;
0310   const float b = -2 * y * (tanphisq + 1);
0311   const float c = (tanphisq + 1) * square(y) - square(R);
0312
0313   const float Y0_1 = (-b + std::sqrt(square(b) - 4 * a * c)) / (2. * a);
0314   const float Y0_2 = (-b - std::sqrt(square(b) - 4 * a * c)) / (2. * a);
0315   const float X0_1 = sqrt(pow(R, 2) - pow(Y0_1 - y, 2)) + x;
0316   const float X0_2 = -sqrt(pow(R, 2) - pow(Y0_2 - y, 2)) + x;
0317   track->set_X0(X0_1);
0318   track->set_Y0(Y0_1);
0319   track->set_qOverR(charge / R);
0320   track->set_slope(1. / std::tan(theta));
0321   track->set_Z0(z);
0322
0323   if (tpc)
0324   {
0325     // if this is the TPC, the track Z0 depends on the crossing
0326     // we must determine Z0 from the clusters instead of the truth for the TPC
0327     // this method calculates the Z0 and z slope from a fit to the clusters and overwrites them
0328     const unsigned int start_layer = 7;
0329     const unsigned int end_layer = 55;
0330
0331     // create local position array
0332     std::map<TrkrDefs::cluskey, Acts::Vector3> positions;
0333     std::transform(track->begin_cluster_keys(), track->end_cluster_keys(), std::inserter(positions, positions.end()),
0334                    [this](const auto& key)
0335                    { return std::make_pair(key, tgeometry->getGlobalPosition(key, m_clusterMap->findCluster(key))); });
0336
0337     TrackSeedHelper::lineFit(track.get(), positions, start_layer, end_layer);
0338   }
0339
0340   // Need to find the right one for the bend angle
0341   // TODO: this should account from distortion corrections, as done, e.g. in PHSimpleKFProp
0342   auto newphi = TrackSeedHelper::get_phi_fastsim(track.get());
0343
0344   // We have to pick the right one based on the bend angle, so iterate
0345   // through until you find the closest phi match
0346   if (std::fabs(newphi - phi) > 0.03)
0347   {
0348     track->set_X0(X0_2);
0349     newphi = TrackSeedHelper::get_phi_fastsim( track.get() );
0350
0351     if (std::fabs(newphi - phi) > 0.03)
0352     {
0353       track->set_Y0(Y0_2);
0354       newphi = TrackSeedHelper::get_phi_fastsim( track.get() );
0355
0356       if (std::fabs(newphi - phi) > 0.03)
0357       {
0358         track->set_X0(X0_1);
0359         newphi = TrackSeedHelper::get_phi_fastsim( track.get() );
0360       }
0361     }
0362   }
0363
0364   // make phi persistent
0365   track->set_phi(newphi);
0366
0367   if (Verbosity() > 2)
0368   {
0369     std::cout << "Charge is " << charge << std::endl;
0370     std::cout << "truth/reco px " << px << ", " << track->get_px() << std::endl;
0371     std::cout << "truth/reco py " << py << ", " << track->get_py() << std::endl;
0372     std::cout << "truth/reco pz " << pz << ", " << track->get_pz() << std::endl;
0373     std::cout << "truth/reco pt " << pt << ", " << track->get_pt() << std::endl;
0374     std::cout << "truth/reco phi " << phi << ", " << track->get_phi() << std::endl;
0375     std::cout << "truth/reco eta " << eta << ", " << track->get_eta() << std::endl;
0376     std::cout << "truth/reco x " << x << ", " << TrackSeedHelper::get_x(track.get()) << std::endl;
0377     std::cout << "truth/reco y " << y << ", " << TrackSeedHelper::get_y(track.get()) << std::endl;
0378     std::cout << "truth/reco z " << z << ", " << TrackSeedHelper::get_z(track.get()) << std::endl;
0379   }
0380
0381   // set intt crossing
0382   if (silicon)
0383   {
0384     // silicon tracklet
0385     /* inspired from PHtruthSiliconAssociation */
0386     const auto intt_crossings = getInttCrossings(track.get());
0387     if (intt_crossings.empty())
0388     {
0389       if (Verbosity() > 1)
0390       {
0391         std::cout << "PHTruthTrackSeeding::Process - Silicon track " << container->size() - 1 << " has no INTT clusters" << std::endl;
0392       }
0393       return;
0394     }
0395     else if (intt_crossings.size() > 1)
0396     {
0397       if (Verbosity() > 1)
0398       {
0399         std::cout << "PHTruthTrackSeeding::Process - INTT crossings not all the same for track " << container->size() - 1 << " crossing_keep - dropping this match " << std::endl;
0400       }
0401     }
0402     else
0403     {
0404       const auto& crossing = *intt_crossings.begin();
0405       track->set_crossing(crossing);
0406       if (Verbosity() > 1)
0407       {
0408         std::cout << "PHTruthTrackSeeding::Process - Combined track " << container->size() - 1 << " bunch crossing " << crossing << std::endl;
0409       }
0410     }
0411   }  // end if _min_layer
0412   else
0413   {
0414     // no INTT layers, crossing is unknown
0415     track->set_crossing(SHRT_MAX);
0416   }
0417
0418   container->insert(track.get());
0419 }
0420
0421 int PHTruthTrackSeeding::CreateNodes(PHCompositeNode* topNode)
0422 {
0423   // create nodes...
0424   PHNodeIterator iter(topNode);
0425
0426   PHCompositeNode* dstNode = static_cast<PHCompositeNode*>(iter.findFirst(
0427       "PHCompositeNode", "DST"));
0428   if (!dstNode)
0429   {
0430     std::cerr << PHWHERE << "DST Node missing, doing nothing." << std::endl;
0431     return Fun4AllReturnCodes::ABORTEVENT;
0432   }
0433   PHNodeIterator iter_dst(dstNode);
0434
0435   // Create the SVTX node
0436   PHCompositeNode* tb_node =
0437       dynamic_cast<PHCompositeNode*>(iter_dst.findFirst("PHCompositeNode",
0438                                                         "SVTX"));
0439   if (!tb_node)
0440   {
0441     tb_node = new PHCompositeNode("SVTX");
0442     dstNode->addNode(tb_node);
0443     if (Verbosity() > 0)
0444     {
0445       std::cout << PHWHERE << "SVTX node added" << std::endl;
0446     }
0447   }
0448
0449   _track_map = findNode::getClass<TrackSeedContainer>(topNode, "TpcTrackSeedContainer");
0450   if (!_track_map)
0451   {
0452     _track_map = new TrackSeedContainer_v1;
0453     PHIODataNode<PHObject>* tracks_node =
0454         new PHIODataNode<PHObject>(_track_map, "TpcTrackSeedContainer", "PHObject");
0455     tb_node->addNode(tracks_node);
0456   }
0457
0458   _track_map_silicon = findNode::getClass<TrackSeedContainer>(topNode, "SiliconTrackSeedContainer");
0459   if (!_track_map_silicon)
0460   {
0461     _track_map_silicon = new TrackSeedContainer_v1;
0462     PHIODataNode<PHObject>* tracks_node =
0463         new PHIODataNode<PHObject>(_track_map_silicon, "SiliconTrackSeedContainer", "PHObject");
0464     tb_node->addNode(tracks_node);
0465   }
0466
0467   _track_map_combined = findNode::getClass<TrackSeedContainer>(topNode, "SvtxTrackSeedContainer");
0468   if (!_track_map_combined)
0469   {
0470     _track_map_combined = new TrackSeedContainer_v1;
0471     PHIODataNode<PHObject>* node2 = new PHIODataNode<PHObject>(_track_map_combined, "SvtxTrackSeedContainer", "PHObject");
0472     tb_node->addNode(node2);
0473   }
0474
0475   return Fun4AllReturnCodes::EVENT_OK;
0476 }
0477 int PHTruthTrackSeeding::GetNodes(PHCompositeNode* topNode)
0478 {
0479   tgeometry = findNode::getClass<ActsGeometry>(topNode, "ActsGeometry");
0480   if (!tgeometry)
0481   {
0482     std::cerr << PHWHERE << "Error, can' find needed Acts nodes " << std::endl;
0483     return Fun4AllReturnCodes::ABORTEVENT;
0484   }
0485
0486   m_clusterMap = findNode::getClass<TrkrClusterContainer>(topNode, "TRKR_CLUSTER");
0487
0488   if (!m_clusterMap)
0489   {
0490     std::cerr << PHWHERE << "Error: Can't find node TRKR_CLUSTER" << std::endl;
0491     return Fun4AllReturnCodes::ABORTEVENT;
0492   }
0493
0494   m_cluster_crossing_map = findNode::getClass<TrkrClusterCrossingAssoc>(topNode, "TRKR_CLUSTERCROSSINGASSOC");
0495   if (!m_cluster_crossing_map)
0496   {
0497     std::cerr << PHWHERE << " ERROR: Can't find TRKR_CLUSTERCROSSINGASSOC " << std::endl;
0498     return Fun4AllReturnCodes::ABORTEVENT;
0499   }
0500
0501   m_g4truth_container = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
0502   if (!m_g4truth_container)
0503   {
0504     std::cerr << PHWHERE << " ERROR: Can't find node G4TruthInfo" << std::endl;
0505     return Fun4AllReturnCodes::ABORTEVENT;
0506   }
0507
0508   hittruthassoc = findNode::getClass<TrkrHitTruthAssoc>(topNode, "TRKR_HITTRUTHASSOC");
0509   if (!hittruthassoc)
0510   {
0511     std::cout << PHWHERE << "Failed to find TRKR_HITTRUTHASSOC node, quit!" << std::endl;
0512     exit(1);
0513   }
0514
0515   using nodePair = std::pair<std::string, PHG4HitContainer*&>;
0516   std::initializer_list<nodePair> nodes =
0517       {
0518           {"G4HIT_TPC", phg4hits_tpc},
0519           {"G4HIT_INTT", phg4hits_intt},
0520           {"G4HIT_MVTX", phg4hits_mvtx},
0521           {"G4HIT_MICROMEGAS", phg4hits_micromegas}};
0522
0523   for (auto&& node : nodes)
0524   {
0525     if (!(node.second = findNode::getClass<PHG4HitContainer>(topNode, node.first)))
0526     {
0527       std::cerr << PHWHERE << " PHG4HitContainer " << node.first << " not found" << std::endl;
0528     }
0529   }
0530
0531   return Fun4AllReturnCodes::EVENT_OK;
0532 }
0533
0534 int PHTruthTrackSeeding::End()
0535 {
0536   return 0;
0537 }
0538
0539 //_____________________________________________________________________________________________
0540 std::set<short int> PHTruthTrackSeeding::getInttCrossings(TrackSeed* si_track) const
0541 {
0542   if (Verbosity())
0543   {
0544     std::cout << "PHTruthTrackSeeding::getInttCrossings - entering " << std::endl;
0545   }
0546
0547   std::set<short int> intt_crossings;
0548
0549   // If the Si track contains an INTT hit, use it to get the bunch crossing offset
0550   // loop over associated clusters to get keys for silicon cluster
0551
0552   for (auto iter = si_track->begin_cluster_keys();
0553        iter != si_track->end_cluster_keys(); ++iter)
0554   {
0555     const TrkrDefs::cluskey& cluster_key = *iter;
0556     const unsigned int trkrid = TrkrDefs::getTrkrId(cluster_key);
0557     if (Verbosity() > 0)
0558     {
0559       std::cout << "    trkrid " << trkrid << " cluster_key " << cluster_key << std::endl;
0560     }
0561     if (trkrid == TrkrDefs::inttId)
0562     {
0563       // get layer from cluster key
0564       const unsigned int layer = TrkrDefs::getLayer(cluster_key);
0565
0566       // get the bunch crossings for all hits in this cluster
0567       const auto crossings = m_cluster_crossing_map->getCrossings(cluster_key);
0568       for (auto iter_A = crossings.first; iter_A != crossings.second; ++iter_A)
0569       {
0570         const auto& [key, crossing] = *iter_A;
0571         if (Verbosity())
0572         {
0573           std::cout << "    PHTruthTrackSeeding::getInttCrossings - si Track cluster " << key << " layer " << layer << " crossing " << crossing << std::endl;
0574         }
0575         intt_crossings.insert(crossing);
0576       }
0577     }
0578   }
0579
0580   return intt_crossings;
0581 }