g4simulation/g4eval/DSTTrackInfoWriter.cc

0001 /*!
0002  * \file DSTTrackInfoWriter.cc
0003  * \author Alex Patton <aopatton@mit.edu>
0004  */
0005
0006 #include "DSTTrackInfoWriter.h"
0007
0008 #include <fun4all/Fun4AllReturnCodes.h>
0009 #include <g4main/PHG4Hit.h>
0010 #include <g4main/PHG4HitContainer.h>
0011 #include <g4main/PHG4Particle.h>
0012 #include <g4main/PHG4TruthInfoContainer.h>
0013 #include <micromegas/MicromegasDefs.h>
0014 #include <phool/PHCompositeNode.h>
0015 #include <phool/PHNodeIterator.h>
0016 #include <phool/getClass.h>
0017 #include <trackbase/InttDefs.h>
0018 #include <trackbase/MvtxDefs.h>
0019 #include <trackbase/TpcDefs.h>
0020 #include <trackbase/TrkrCluster.h>
0021 #include <trackbase/TrkrClusterv4.h>
0022 #include <trackbase/TrkrDefs.h>
0023 // include new cluster
0024 #include <trackbase/TrkrClusterContainer.h>
0025 #include <trackbase/TrkrClusterHitAssoc.h>
0026 #include <trackbase/TrkrHit.h>
0027 #include <trackbase/TrkrHitSet.h>
0028 #include <trackbase/TrkrHitSetContainer.h>
0029 #include <trackbase/TrkrHitTruthAssoc.h>
0030 #include <trackbase_historic/SvtxTrack.h>
0031 #include <trackbase_historic/SvtxTrackMap.h>
0032 #include <trackbase_historic/TrackSeed_v1.h>
0033 // include SvtxTrackMap_v1 to write data to it
0034 #include <trackbase_historic/SvtxTrackInfo_v1.h>
0035 #include <trackbase_historic/SvtxTrackMap_v1.h>
0036 #include <trackbase_historic/SvtxTrack_v4.h>
0037 #include <trackbase_historic/TrackInfoContainer_v1.h>
0038 #include <trackbase_historic/TrackStateInfo_v1.h>
0039 #include <trackbase_historic/ActsTransformations.h>
0040
0041 #include <TClonesArray.h>
0042 #include <TFile.h>
0043 #include <TLine.h>
0044 #include <TTree.h>
0045
0046 #include <algorithm>
0047 #include <bitset>
0048 #include <cassert>
0049 #include <iostream>
0050 #include <numeric>
0051
0052 //_____________________________________________________________________
0053
0054 //_____________________________________________________________________
0055 DSTTrackInfoWriter::DSTTrackInfoWriter(const std::string& name)
0056   : SubsysReco(name)
0057 {
0058 }
0059
0060 //_____________________________________________________________________
0061 int DSTTrackInfoWriter::InitRun(PHCompositeNode* topNode)
0062 {
0063   if (Verbosity() > 1)
0064   {
0065     std::cout << "Writer Init start" << std::endl;
0066   }
0067   // find DST node
0068   PHNodeIterator iter(topNode);
0069   auto dstNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0070   if (!dstNode)
0071   {
0072     std::cout << "DSTTrackInfoWriter::Init - DST Node missing" << std::endl;
0073     return Fun4AllReturnCodes::ABORTEVENT;
0074   }
0075
0076   // get EVAL node
0077   iter = PHNodeIterator(dstNode);
0078   auto evalNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "EVAL"));
0079   if (!evalNode)
0080   {
0081     // create
0082     std::cout << "DSTTrackInfoWriter::Init - EVAL node missing - creating" << std::endl;
0083     evalNode = new PHCompositeNode("EVAL");
0084     dstNode->addNode(evalNode);
0085   }
0086
0087   // // TClonesArary
0088   // m_container->arrClsDST = new TClonesArray("DSTContainerv3::ClusterStruct");
0089   // m_container->trkrClsDST = new TClonesArray("TrkrClusterv4");
0090
0091   auto newInfoNode = new PHIODataNode<PHObject>(new TrackInfoContainer_v1, "TrackInfoContainer", "PHObject");
0092   evalNode->addNode(newInfoNode);
0093
0094   // auto newTrackNode = new PHIODataNode<PHObject>(new SvtxTrackArray_v1, "TRACK_ARRAYV1", "PHObject");
0095   // evalNode->addNode(newTrackNode);
0096
0097   // TClonesArary
0098   // fcl = new TFile("dstcl.root", "recreate");
0099   // tcl = new TTree("tcl", "dst tree");
0100   // arrEvt = new TClonesArray("DSTContainerv3::EventStruct");
0101   // arrTrk = new TClonesArray("DSTContainerv3::TrackStruct");
0102   // arrCls = new TClonesArray("DSTContainerv3::ClusterStruct");
0103   // tcl->Branch("evt", &arrEvt);
0104   // tcl->Branch("trk", &arrTrk);
0105   // tcl->Branch("cls", &arrCls);
0106   if (Verbosity() > 1)
0107   {
0108     std::cout << "Writer Init end" << std::endl;
0109   }
0110   return Fun4AllReturnCodes::EVENT_OK;
0111 }
0112
0113 //_____________________________________________________________________
0114 int DSTTrackInfoWriter::process_event(PHCompositeNode* topNode)
0115 {
0116   // make topNode run in Init
0117   // Init(topNode);
0118   //  load nodes
0119   if (Verbosity() > 1)
0120   {
0121     std::cout << __FILE__ << "::" << __func__ << "::" << __LINE__ << std::endl;
0122     std::cout << "DSTTrackInfoWriter::process_event" << std::endl;
0123   }
0124   auto res = load_nodes(topNode);
0125   if (res != Fun4AllReturnCodes::EVENT_OK)
0126   {
0127     return res;
0128   }
0129   if (Verbosity() > 1)
0130   {
0131     std::cout << "Return codes  start" << Fun4AllReturnCodes::EVENT_OK << std::endl;
0132   }
0133   // cleanup output container
0134   if (m_track_info_container) { m_track_info_container->Reset();
0135 }
0136   if (Verbosity() > 1)
0137   {
0138     std::cout << "Evalutate track info" << std::endl;
0139   }
0140   evaluate_track_info();
0141   if (Verbosity() > 1)
0142   {
0143     std::cout << "exiting event"
0144               << "\n";
0145
0146     std::cout << "Return codes end" << Fun4AllReturnCodes::EVENT_OK << std::endl;
0147   }
0148   return Fun4AllReturnCodes::EVENT_OK;
0149 }
0150
0151 //_____________________________________________________________________
0152 int DSTTrackInfoWriter::load_nodes(PHCompositeNode* topNode)
0153 {
0154   // get necessary nodes
0155   m_track_map = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
0156
0157   m_track_info_container = findNode::getClass<TrackInfoContainer_v1>(topNode, "TrackInfoContainer");
0158
0159   return Fun4AllReturnCodes::EVENT_OK;
0160 }
0161
0162 //____________________________________________________________________
0163 void DSTTrackInfoWriter::evaluate_track_info()
0164 {
0165   if (!(m_track_info_container))
0166   {
0167     return;
0168   }
0169
0170   m_track_info_container->Reset();
0171   // get track into track info
0172   Int_t iTrk = 0;
0173   // long unsigned int iKey = 0;
0174   if (Verbosity() > 1)
0175   {
0176     std::cout << "Before loop"
0177               << "\n";
0178   }
0179   ActsTransformations transformer;
0180   for (const auto& trackpair : *m_track_map)
0181   {
0182     const auto track = trackpair.second;
0183     // this track will have a TPC and Silicon seed
0184
0185     uint64_t hitbitmap = 0;
0186
0187     SvtxTrackInfo_v1* trackInfo = new SvtxTrackInfo_v1();
0188     if (Verbosity() > 1)
0189     {
0190       std::cout << "Before seeds"
0191                 << "\n";
0192     }
0193     TrackSeed* TPCSeed = track->get_tpc_seed();
0194     TrackSeed* SiliconSeed = track->get_silicon_seed();
0195     if (TPCSeed)
0196     {
0197       for (auto key_iter = TPCSeed->begin_cluster_keys(); key_iter != TPCSeed->end_cluster_keys(); ++key_iter)
0198       {
0199         const auto& cluster_key = *key_iter;
0200
0201         // store information in track array
0202         // std::cout << "TPC clusterkey: " << cluster_key <<"\n";
0203         // std::cout << "TPC subsurfkey: " << cluster->getSubSurfKey() << std::endl;
0204         uint8_t layer = TrkrDefs::getLayer(cluster_key);
0205         if (Verbosity() > 1)
0206         {
0207           std::cout << "Layer is: " << unsigned(layer) << std::endl;
0208         }
0209         hitbitmap = hitbitmap + ((uint64_t) 1 << layer);
0210
0211         // TrkrDefs::
0212       }
0213     }
0214     if (Verbosity() > 1)
0215     {
0216       std::cout << "Before Silicon seeds"
0217                 << "\n";
0218     }
0219     if (!SiliconSeed)
0220     {
0221       if (Verbosity() > 1)
0222       {
0223         std::cout << "Silicon Seed does not exist" << std::endl;
0224       }
0225     }
0226
0227     if (SiliconSeed)
0228     {
0229       for (auto key_iter = SiliconSeed->begin_cluster_keys(); key_iter != SiliconSeed->end_cluster_keys(); ++key_iter)
0230       {
0231         const auto& cluster_key = *key_iter;
0232
0233         uint8_t layer = TrkrDefs::getLayer(cluster_key);
0234         if (Verbosity() > 1)
0235         {
0236           std::cout << "Layer is: " << unsigned(layer) << std::endl;
0237         }
0238         hitbitmap = hitbitmap + ((uint64_t) 1 << layer);
0239       }
0240     }
0241     if (Verbosity() > 1)
0242     {
0243       std::cout << "After Track seeds"
0244                 << "\n";
0245     }
0246     trackInfo->set_hitbitmap(hitbitmap);
0247     trackInfo->set_chisq(track->get_chisq());
0248     trackInfo->set_ndf(track->get_ndf());
0249     trackInfo->set_crossing(track->get_crossing());
0250
0251     trackInfo->set_x(track->get_x());
0252     trackInfo->set_y(track->get_y());
0253     trackInfo->set_z(track->get_z());
0254     trackInfo->set_phi(track->get_phi());
0255     trackInfo->set_theta(2.* std::atan(std::exp(-1*track->get_eta())));
0256     trackInfo->set_qOp(track->get_charge() / track->get_p());
0257     if (Verbosity() > 1)
0258     {
0259       std::cout << "track crossing: " << track->get_crossing() << std::endl;
0260
0261       std::cout << "track.get_z(): " << track->get_z() << std::endl;
0262       std::cout << "trackInfo.get_z(): " << trackInfo->get_z() << std::endl;
0263       std::cout << "Hitbitmap: " << trackInfo->get_hitbitmap() << std::endl;
0264       std::cout << "crossing: " << trackInfo->get_crossing() << std::endl;
0265       std::cout << "chi^2: " << trackInfo->get_chisq() << std::endl;
0266       std::cout << "ndf: " << unsigned(trackInfo->get_ndf()) << std::endl;
0267     }
0268     auto actsMatrix = transformer.rotateSvtxTrackCovToActs(track);
0269     //! only take the first 5 entries, which correspond to d0, z0, phi, theta, q/p
0270     for (int i = 0; i < 5; i++)
0271     {
0272       for (int j = i; j < 5; j++)
0273       {
0274         trackInfo->set_covariance(i, j, actsMatrix(i,j));
0275       }
0276     }
0277     if (Verbosity() > 1)
0278     {
0279       std::cout << "Right before adding track info" << iTrk << std::endl;
0280     }
0281     m_track_info_container->add_trackinfo(iTrk, trackInfo);
0282     if (Verbosity() > 1)
0283     {
0284       std::cout << "Right after adding track info" << std::endl;
0285     }
0286     delete trackInfo;
0287     ++iTrk;
0288   }
0289
0290   // add trackinfo to trackinfocontainer
0291 }
0292 //_____________________________________________________________________
0293
0294 //_____________________________________________________________________