LWTrackNtuplizer/module/lwtrackntuplizer.cc

0001 #include "lwtrackntuplizer.h"
0002
0003 #include <fun4all/Fun4AllReturnCodes.h>
0004
0005 #include <phool/PHCompositeNode.h>
0006 #include <phool/getClass.h>
0007 #include <phool/phool.h>
0008
0009 #include <trackbase/ActsGeometry.h>
0010 #include <trackbase/TrkrClusterContainer.h>
0011 #include <trackbase/TrkrDefs.h>
0012 #include <trackbase_historic/SvtxTrack.h>
0013 #include <trackbase_historic/SvtxTrackMap.h>
0014 #include <trackbase_historic/TrackAnalysisUtils.h>
0015 #include <trackbase_historic/TrackSeed.h>
0016
0017 #include <globalvertex/SvtxVertex.h>
0018 #include <globalvertex/SvtxVertexMap.h>
0019
0020 #include <g4detectors/PHG4TpcGeom.h>
0021 #include <g4detectors/PHG4TpcGeomContainer.h>
0022
0023 #include <TFile.h>
0024 #include <TTree.h>
0025 #include <TVector3.h>
0026
0027 #include <cmath>
0028 #include <climits>
0029 #include <iostream>
0030 #include <limits>
0031
0032 namespace
0033 {
0034 template <class Container>
0035 void Clean(Container& c)
0036 {
0037   Container().swap(c);
0038 }
0039
0040 float nan_value()
0041 {
0042   return std::numeric_limits<float>::quiet_NaN();
0043 }
0044 }  // namespace
0045
0046 lwtrackntuplizer::lwtrackntuplizer(const std::string& name)
0047   : SubsysReco(name)
0048 {
0049 }
0050
0051 lwtrackntuplizer::~lwtrackntuplizer() = default;
0052
0053 int lwtrackntuplizer::Init(PHCompositeNode* /*topNode*/)
0054 {
0055   m_outFile = new TFile(m_outFileName.c_str(), "RECREATE");
0056   if (!m_outFile || m_outFile->IsZombie())
0057   {
0058     std::cout << PHWHERE << " failed to create output file " << m_outFileName << std::endl;
0059     return Fun4AllReturnCodes::ABORTRUN;
0060   }
0061
0062   m_outFile->SetCompressionLevel(7);
0063   SetupTree();
0064   Cleanup();
0065
0066   return Fun4AllReturnCodes::EVENT_OK;
0067 }
0068
0069 int lwtrackntuplizer::InitRun(PHCompositeNode* topNode)
0070 {
0071   const int ret = GetNodes(topNode);
0072   if (ret != Fun4AllReturnCodes::EVENT_OK)
0073   {
0074     return Fun4AllReturnCodes::ABORTRUN;
0075   }
0076
0077   return Fun4AllReturnCodes::EVENT_OK;
0078 }
0079
0080 int lwtrackntuplizer::process_event(PHCompositeNode* topNode)
0081 {
0082   const int ret = GetNodes(topNode);
0083   if (ret != Fun4AllReturnCodes::EVENT_OK)
0084   {
0085     return ret;
0086   }
0087
0088   Cleanup();
0089
0090   if (m_trackMap)
0091   {
0092     for (const auto& entry : *m_trackMap)
0093     {
0094       const auto* track = entry.second;
0095       if (!track)
0096       {
0097         continue;
0098       }
0099       FillTrack(track);
0100     }
0101   }
0102
0103   m_nTracks = m_trackID.size();
0104   if (m_outTree)
0105   {
0106     m_outTree->Fill();
0107   }
0108
0109   ++m_event;
0110   return Fun4AllReturnCodes::EVENT_OK;
0111 }
0112
0113 int lwtrackntuplizer::ResetEvent(PHCompositeNode* /*topNode*/)
0114 {
0115   Cleanup();
0116   return Fun4AllReturnCodes::EVENT_OK;
0117 }
0118
0119 int lwtrackntuplizer::EndRun(const int /*runnumber*/)
0120 {
0121   return Fun4AllReturnCodes::EVENT_OK;
0122 }
0123
0124 int lwtrackntuplizer::End(PHCompositeNode* /*topNode*/)
0125 {
0126   if (m_outFile)
0127   {
0128     m_outFile->cd();
0129     if (m_outTree)
0130     {
0131       m_outTree->Write("", TObject::kOverwrite);
0132     }
0133     m_outFile->Close();
0134     delete m_outFile;
0135     m_outFile = nullptr;
0136     m_outTree = nullptr;
0137   }
0138
0139   return Fun4AllReturnCodes::EVENT_OK;
0140 }
0141
0142 int lwtrackntuplizer::Reset(PHCompositeNode* /*topNode*/)
0143 {
0144   Cleanup();
0145   return Fun4AllReturnCodes::EVENT_OK;
0146 }
0147
0148 void lwtrackntuplizer::Print(const std::string& what) const
0149 {
0150   std::cout << "lwtrackntuplizer::Print - " << what
0151             << ", output=" << m_outFileName
0152             << ", tree=" << m_treeName
0153             << ", trackmap=" << m_trackMapName
0154             << std::endl;
0155 }
0156
0157 int lwtrackntuplizer::GetNodes(PHCompositeNode* topNode)
0158 {
0159   m_trackMap = findNode::getClass<SvtxTrackMap>(topNode, m_trackMapName);
0160   if (!m_trackMap)
0161   {
0162     std::cout << PHWHERE << " missing required node " << m_trackMapName << std::endl;
0163     return Fun4AllReturnCodes::ABORTEVENT;
0164   }
0165
0166   m_clusterMap = findNode::getClass<TrkrClusterContainer>(topNode, m_clusterContainerName);
0167   if (!m_clusterMap)
0168   {
0169     std::cout << PHWHERE << " missing required node " << m_clusterContainerName << std::endl;
0170     return Fun4AllReturnCodes::ABORTEVENT;
0171   }
0172
0173   m_tGeometry = findNode::getClass<ActsGeometry>(topNode, m_geometryNodeName);
0174   if (!m_tGeometry)
0175   {
0176     std::cout << PHWHERE << " missing required node " << m_geometryNodeName << std::endl;
0177     return Fun4AllReturnCodes::ABORTEVENT;
0178   }
0179
0180   m_tpcGeomContainer = findNode::getClass<PHG4TpcGeomContainer>(topNode, m_tpcGeomNodeName);
0181   if (!m_tpcGeomContainer)
0182   {
0183     std::cout << PHWHERE << " missing required node " << m_tpcGeomNodeName << std::endl;
0184     return Fun4AllReturnCodes::ABORTEVENT;
0185   }
0186
0187   m_vertexMap = findNode::getClass<SvtxVertexMap>(topNode, m_vertexMapName);
0188
0189   return Fun4AllReturnCodes::EVENT_OK;
0190 }
0191
0192 void lwtrackntuplizer::SetupTree()
0193 {
0194   if (!m_outFile)
0195   {
0196     return;
0197   }
0198
0199   m_outFile->cd();
0200   m_outTree = new TTree(m_treeName.c_str(), m_treeName.c_str());
0201
0202   m_outTree->Branch("event", &m_event, "event/i");
0203   m_outTree->Branch("nTracks", &m_nTracks, "nTracks/i");
0204
0205   m_outTree->Branch("trackID", &m_trackID);
0206   m_outTree->Branch("crossing", &m_crossing);
0207   m_outTree->Branch("px", &m_px);
0208   m_outTree->Branch("py", &m_py);
0209   m_outTree->Branch("pz", &m_pz);
0210   m_outTree->Branch("pt", &m_pt);
0211   m_outTree->Branch("eta", &m_eta);
0212   m_outTree->Branch("phi", &m_phi);
0213   m_outTree->Branch("deltapt", &m_deltapt);
0214   m_outTree->Branch("deltaeta", &m_deltaeta);
0215   m_outTree->Branch("deltaphi", &m_deltaphi);
0216   m_outTree->Branch("charge", &m_charge);
0217   m_outTree->Branch("quality", &m_quality);
0218   m_outTree->Branch("chisq", &m_chisq);
0219   m_outTree->Branch("ndf", &m_ndf);
0220   m_outTree->Branch("nhits", &m_nhits);
0221   m_outTree->Branch("nmaps", &m_nmaps);
0222   m_outTree->Branch("nintt", &m_nintt);
0223   m_outTree->Branch("ntpc", &m_ntpc);
0224   m_outTree->Branch("nmms", &m_nmms);
0225   m_outTree->Branch("ntpc1", &m_ntpc1);
0226   m_outTree->Branch("ntpc11", &m_ntpc11);
0227   m_outTree->Branch("ntpc2", &m_ntpc2);
0228   m_outTree->Branch("ntpc3", &m_ntpc3);
0229   m_outTree->Branch("dedx", &m_dedx);
0230   m_outTree->Branch("pidedx", &m_pidedx);
0231   m_outTree->Branch("kdedx", &m_kdedx);
0232   m_outTree->Branch("prdedx", &m_prdedx);
0233   m_outTree->Branch("vertexID", &m_vertexID);
0234   m_outTree->Branch("vx", &m_vx);
0235   m_outTree->Branch("vy", &m_vy);
0236   m_outTree->Branch("vz", &m_vz);
0237   m_outTree->Branch("dca2d", &m_dca2d);
0238   m_outTree->Branch("dca2dsigma", &m_dca2dsigma);
0239   m_outTree->Branch("dca3dxy", &m_dca3dxy);
0240   m_outTree->Branch("dca3dxysigma", &m_dca3dxysigma);
0241   m_outTree->Branch("dca3dz", &m_dca3dz);
0242   m_outTree->Branch("dca3dzsigma", &m_dca3dzsigma);
0243   m_outTree->Branch("pcax", &m_pcax);
0244   m_outTree->Branch("pcay", &m_pcay);
0245   m_outTree->Branch("pcaz", &m_pcaz);
0246   m_outTree->Branch("hlxpt", &m_hlxpt);
0247   m_outTree->Branch("hlxeta", &m_hlxeta);
0248   m_outTree->Branch("hlxphi", &m_hlxphi);
0249   m_outTree->Branch("hlxX0", &m_hlxX0);
0250   m_outTree->Branch("hlxY0", &m_hlxY0);
0251   m_outTree->Branch("hlxZ0", &m_hlxZ0);
0252   m_outTree->Branch("hlxcharge", &m_hlxcharge);
0253 }
0254
0255 void lwtrackntuplizer::FillTrack(const SvtxTrack* track)
0256 {
0257   const float nan = nan_value();
0258
0259   const auto* tpcseed = track->get_tpc_seed();
0260   const auto* silseed = track->get_silicon_seed();
0261
0262   int nhits = 0;
0263   int nmaps = 0;
0264   int nintt = 0;
0265   int ntpc = 0;
0266   int nmms = 0;
0267   int ntpc1 = 0;
0268   int ntpc11 = 0;
0269   int ntpc2 = 0;
0270   int ntpc3 = 0;
0271
0272   const auto count_seed = [&](const TrackSeed* seed)
0273   {
0274     if (!seed)
0275     {
0276       return;
0277     }
0278
0279     nhits += seed->size_cluster_keys();
0280     for (auto iter = seed->begin_cluster_keys(); iter != seed->end_cluster_keys(); ++iter)
0281     {
0282       const auto clusterKey = *iter;
0283       const unsigned int layer = TrkrDefs::getLayer(clusterKey);
0284       switch (TrkrDefs::getTrkrId(clusterKey))
0285       {
0286       case TrkrDefs::TrkrId::mvtxId:
0287         ++nmaps;
0288         break;
0289       case TrkrDefs::TrkrId::inttId:
0290         ++nintt;
0291         break;
0292       case TrkrDefs::TrkrId::tpcId:
0293         ++ntpc;
0294         if ((layer - 7U) < 8U)
0295         {
0296           ++ntpc11;
0297         }
0298         if ((layer - 7U) < 16U)
0299         {
0300           ++ntpc1;
0301         }
0302         else if ((layer - 7U) < 32U)
0303         {
0304           ++ntpc2;
0305         }
0306         else if ((layer - 7U) < 48U)
0307         {
0308           ++ntpc3;
0309         }
0310         break;
0311       case TrkrDefs::TrkrId::micromegasId:
0312         ++nmms;
0313         break;
0314       default:
0315         break;
0316       }
0317     }
0318   };
0319
0320   count_seed(tpcseed);
0321   count_seed(silseed);
0322
0323   float dedx = nan;
0324   if (tpcseed)
0325   {
0326     auto* inner1 = m_tpcGeomContainer->GetLayerCellGeom(7);
0327     auto* inner2 = m_tpcGeomContainer->GetLayerCellGeom(8);
0328     auto* middle = m_tpcGeomContainer->GetLayerCellGeom(27);
0329     auto* outer = m_tpcGeomContainer->GetLayerCellGeom(50);
0330
0331     if (inner1 && inner2 && middle && outer)
0332     {
0333       float layerThicknesses[4] = {
0334         static_cast<float>(inner1->get_thickness()),
0335         static_cast<float>(inner2->get_thickness()),
0336         static_cast<float>(middle->get_thickness()),
0337         static_cast<float>(outer->get_thickness())};
0338       dedx = TrackAnalysisUtils::calc_dedx(
0339           const_cast<TrackSeed*>(tpcseed), m_clusterMap, m_tGeometry, layerThicknesses);
0340     }
0341   }
0342
0343   const float px = track->get_px();
0344   const float py = track->get_py();
0345   const float pz = track->get_pz();
0346   const TVector3 momentum(px, py, pz);
0347   const float pt = momentum.Pt();
0348   const float eta = momentum.Eta();
0349   const float phi = momentum.Phi();
0350
0351   const float cvxx = track->get_error(3, 3);
0352   const float cvxy = track->get_error(3, 4);
0353   const float cvxz = track->get_error(3, 5);
0354   const float cvyy = track->get_error(4, 4);
0355   const float cvyz = track->get_error(4, 5);
0356   const float cvzz = track->get_error(5, 5);
0357
0358   const double pt2 = static_cast<double>(px) * px + static_cast<double>(py) * py;
0359   const double p2 = pt2 + static_cast<double>(pz) * pz;
0360
0361   float deltapt = nan;
0362   if (pt2 > 0.)
0363   {
0364     const double arg = (static_cast<double>(cvxx) * px * px + 2. * static_cast<double>(cvxy) * px * py + static_cast<double>(cvyy) * py * py) / pt2;
0365     if (std::isfinite(arg) && arg >= 0.)
0366     {
0367       deltapt = std::sqrt(arg);
0368     }
0369   }
0370
0371   float deltaeta = nan;
0372   if (pt2 > 0. && p2 > 0.)
0373   {
0374     const double numerator =
0375         static_cast<double>(cvzz) * pt2 * pt2 +
0376         static_cast<double>(pz) *
0377             (-2. * (static_cast<double>(cvxz) * px + static_cast<double>(cvyz) * py) * pt2 +
0378              static_cast<double>(cvxx) * px * px * pz +
0379              static_cast<double>(cvyy) * py * py * pz +
0380              2. * static_cast<double>(cvxy) * px * py * pz);
0381     const double denominator = pt2 * pt2 * p2;
0382     const double arg = numerator / denominator;
0383     if (std::isfinite(arg) && arg >= 0.)
0384     {
0385       deltaeta = std::sqrt(arg);
0386     }
0387   }
0388
0389   float deltaphi = nan;
0390   if (pt2 > 0.)
0391   {
0392     const double denominator = pt2 * pt2;
0393     const double arg =
0394         (static_cast<double>(cvyy) * px * px - 2. * static_cast<double>(cvxy) * px * py + static_cast<double>(cvxx) * py * py) /
0395         denominator;
0396     if (std::isfinite(arg) && arg >= 0.)
0397     {
0398       deltaphi = std::sqrt(arg);
0399     }
0400   }
0401
0402   const int vertexID = track->get_vertex_id();
0403   float vx = nan;
0404   float vy = nan;
0405   float vz = nan;
0406   if (vertexID >= 0 && m_vertexMap)
0407   {
0408     auto vertexIter = m_vertexMap->find(vertexID);
0409     if (vertexIter != m_vertexMap->end() && vertexIter->second)
0410     {
0411       vx = vertexIter->second->get_x();
0412       vy = vertexIter->second->get_y();
0413       vz = vertexIter->second->get_z();
0414     }
0415   }
0416
0417   float hlxpt = nan;
0418   float hlxeta = nan;
0419   float hlxphi = nan;
0420   float hlxX0 = nan;
0421   float hlxY0 = nan;
0422   float hlxZ0 = nan;
0423   int hlxcharge = 0;
0424   if (tpcseed)
0425   {
0426     // The TPC seed stores the helix-fit parameters that TrkrNtuplizer labels as hlx* fields
0427     hlxpt = tpcseed->get_pt();
0428     hlxeta = tpcseed->get_eta();
0429     hlxphi = tpcseed->get_phi();
0430     hlxX0 = tpcseed->get_X0();
0431     hlxY0 = tpcseed->get_Y0();
0432     hlxZ0 = tpcseed->get_Z0();
0433     if (std::isfinite(tpcseed->get_qOverR()) && tpcseed->get_qOverR() != 0.)
0434     {
0435       hlxcharge = (tpcseed->get_qOverR() > 0.) ? 1 : -1;
0436     }
0437   }
0438
0439   m_trackID.push_back(track->get_id());
0440   m_crossing.push_back(track->get_crossing() == SHRT_MAX ? nan : static_cast<float>(track->get_crossing()));
0441   m_px.push_back(px);
0442   m_py.push_back(py);
0443   m_pz.push_back(pz);
0444   m_pt.push_back(pt);
0445   m_eta.push_back(eta);
0446   m_phi.push_back(phi);
0447   m_deltapt.push_back(deltapt);
0448   m_deltaeta.push_back(deltaeta);
0449   m_deltaphi.push_back(deltaphi);
0450   m_charge.push_back(track->get_charge());
0451   m_quality.push_back(track->get_quality());
0452   m_chisq.push_back(track->get_chisq());
0453   m_ndf.push_back(track->get_ndf());
0454   m_nhits.push_back(nhits);
0455   m_nmaps.push_back(nmaps);
0456   m_nintt.push_back(nintt);
0457   m_ntpc.push_back(ntpc);
0458   m_nmms.push_back(nmms);
0459   m_ntpc1.push_back(ntpc1);
0460   m_ntpc11.push_back(ntpc11);
0461   m_ntpc2.push_back(ntpc2);
0462   m_ntpc3.push_back(ntpc3);
0463   m_dedx.push_back(dedx);
0464   m_pidedx.push_back(nan);
0465   m_kdedx.push_back(nan);
0466   m_prdedx.push_back(nan);
0467   m_vertexID.push_back(vertexID);
0468   m_vx.push_back(vx);
0469   m_vy.push_back(vy);
0470   m_vz.push_back(vz);
0471   m_dca2d.push_back(track->get_dca2d());
0472   m_dca2dsigma.push_back(track->get_dca2d_error());
0473   m_dca3dxy.push_back(track->get_dca3d_xy());
0474   m_dca3dxysigma.push_back(track->get_dca3d_xy_error());
0475   m_dca3dz.push_back(track->get_dca3d_z());
0476   m_dca3dzsigma.push_back(track->get_dca3d_z_error());
0477   m_pcax.push_back(track->get_x());
0478   m_pcay.push_back(track->get_y());
0479   m_pcaz.push_back(track->get_z());
0480   m_hlxpt.push_back(hlxpt);
0481   m_hlxeta.push_back(hlxeta);
0482   m_hlxphi.push_back(hlxphi);
0483   m_hlxX0.push_back(hlxX0);
0484   m_hlxY0.push_back(hlxY0);
0485   m_hlxZ0.push_back(hlxZ0);
0486   m_hlxcharge.push_back(hlxcharge);
0487 }
0488
0489 void lwtrackntuplizer::Cleanup()
0490 {
0491   m_nTracks = 0;
0492
0493   Clean(m_trackID);
0494   Clean(m_crossing);
0495   Clean(m_px);
0496   Clean(m_py);
0497   Clean(m_pz);
0498   Clean(m_pt);
0499   Clean(m_eta);
0500   Clean(m_phi);
0501   Clean(m_deltapt);
0502   Clean(m_deltaeta);
0503   Clean(m_deltaphi);
0504   Clean(m_charge);
0505   Clean(m_quality);
0506   Clean(m_chisq);
0507   Clean(m_ndf);
0508   Clean(m_nhits);
0509   Clean(m_nmaps);
0510   Clean(m_nintt);
0511   Clean(m_ntpc);
0512   Clean(m_nmms);
0513   Clean(m_ntpc1);
0514   Clean(m_ntpc11);
0515   Clean(m_ntpc2);
0516   Clean(m_ntpc3);
0517   Clean(m_dedx);
0518   Clean(m_pidedx);
0519   Clean(m_kdedx);
0520   Clean(m_prdedx);
0521   Clean(m_vertexID);
0522   Clean(m_vx);
0523   Clean(m_vy);
0524   Clean(m_vz);
0525   Clean(m_dca2d);
0526   Clean(m_dca2dsigma);
0527   Clean(m_dca3dxy);
0528   Clean(m_dca3dxysigma);
0529   Clean(m_dca3dz);
0530   Clean(m_dca3dzsigma);
0531   Clean(m_pcax);
0532   Clean(m_pcay);
0533   Clean(m_pcaz);
0534   Clean(m_hlxpt);
0535   Clean(m_hlxeta);
0536   Clean(m_hlxphi);
0537   Clean(m_hlxX0);
0538   Clean(m_hlxY0);
0539   Clean(m_hlxZ0);
0540   Clean(m_hlxcharge);
0541 }