packages/tpctrackreco/Tpc_PolyTrackVertexer.cc

0001 #include "Tpc_PolyTrackVertexer.h"
0002
0003 #include "Tpc_PolyTrack.h"
0004 #include "Tpc_PolyTrackContainer.h"
0005 #include "Tpc_PolyTrackVertexContainerv1.h"
0006 #include "Tpc_PolyTrackVertexv1.h"
0007
0008 #include <fun4all/Fun4AllReturnCodes.h>
0009
0010 #include <phool/PHCompositeNode.h>
0011 #include <phool/PHIODataNode.h>
0012 #include <phool/PHNodeIterator.h>
0013 #include <phool/PHObject.h>
0014 #include <phool/getClass.h>
0015
0016 #include <TMath.h>
0017
0018 #include <algorithm>
0019 #include <cmath>
0020 #include <iostream>
0021 #include <vector>
0022
0023 namespace
0024 {
0025   double wrap_phi(double phi)
0026   {
0027     while (phi > TMath::Pi())
0028     {
0029       phi -= 2.0 * TMath::Pi();
0030     }
0031     while (phi <= -TMath::Pi())
0032     {
0033       phi += 2.0 * TMath::Pi();
0034     }
0035     return phi;
0036   }
0037
0038   bool good(double x)
0039   {
0040     return std::isfinite(x) && std::fabs(x) < 1.0e30;
0041   }
0042 }  // namespace
0043
0044 Tpc_PolyTrackVertexer::Tpc_PolyTrackVertexer(const std::string& name)
0045   : SubsysReco(name)
0046   , m_inputNodeName("TPC_POLYTRACKS")
0047   , m_outputNodeName("TPC_POLYTRACKVERTICES")
0048 {
0049 }
0050
0051 int Tpc_PolyTrackVertexer::InitRun(PHCompositeNode* topNode)
0052 {
0053   if (!createNodes(topNode))
0054   {
0055     return Fun4AllReturnCodes::ABORTRUN;
0056   }
0057   return Fun4AllReturnCodes::EVENT_OK;
0058 }
0059
0060 bool Tpc_PolyTrackVertexer::getNodes(PHCompositeNode* topNode)
0061 {
0062   m_polyTracks = findNode::getClass<Tpc_PolyTrackContainer>(topNode, m_inputNodeName);
0063   if (!m_polyTracks)
0064   {
0065     const char* candidate_names[] = {
0066         "TPC_POLYTRACKS",
0067         "Tpc_PolyTracks",
0068         "Tpc_PolyTrackContainer",
0069         "TPC_POLYTRACKCONTAINER"};
0070
0071     for (unsigned int i = 0;
0072          i < sizeof(candidate_names) / sizeof(candidate_names[0]) && !m_polyTracks;
0073          ++i)
0074     {
0075       m_polyTracks = findNode::getClass<Tpc_PolyTrackContainer>(topNode, candidate_names[i]);
0076       if (m_polyTracks)
0077       {
0078         m_inputNodeName = candidate_names[i];
0079       }
0080     }
0081   }
0082
0083   if (!m_polyTracks)
0084   {
0085     std::cerr << Name() << "::getNodes - missing Tpc_PolyTrackContainer node" << std::endl;
0086     return false;
0087   }
0088
0089   return true;
0090 }
0091
0092 bool Tpc_PolyTrackVertexer::createNodes(PHCompositeNode* topNode)
0093 {
0094   PHNodeIterator iter(topNode);
0095   PHCompositeNode* dstNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0096   if (!dstNode)
0097   {
0098     dstNode = new PHCompositeNode("DST");
0099     topNode->addNode(dstNode);
0100   }
0101
0102   m_vertices = findNode::getClass<Tpc_PolyTrackVertexContainer>(topNode, m_outputNodeName);
0103   if (!m_vertices)
0104   {
0105     m_vertices = new Tpc_PolyTrackVertexContainerv1();
0106     PHIODataNode<PHObject>* node = new PHIODataNode<PHObject>(m_vertices, m_outputNodeName, "PHObject");
0107     dstNode->addNode(node);
0108     std::cout << Name() << "::createNodes - created " << m_outputNodeName << " node" << std::endl;
0109   }
0110
0111   return true;
0112 }
0113
0114 Tpc_PolyTrackVertexer::TrackVertexFit
0115 Tpc_PolyTrackVertexer::fitTrack(const Tpc_PolyTrack* trk) const
0116 {
0117   TrackVertexFit fit;
0118   if (!trk || trk->get_fit_status() <= 0)
0119   {
0120     return fit;
0121   }
0122
0123   const double x = trk->get_x();
0124   const double y = trk->get_y();
0125   const double z = trk->get_z();
0126   const double px = trk->get_px();
0127   const double py = trk->get_py();
0128   const double pz = trk->get_pz();
0129   const double charge = trk->get_charge();
0130   if (!good(x) || !good(y) || !good(z) ||
0131       !good(px) || !good(py) || !good(pz) || !good(charge))
0132   {
0133     return fit;
0134   }
0135
0136   const double pt2 = px * px + py * py;
0137   if (pt2 <= 1.0e-20)
0138   {
0139     return fit;
0140   }
0141   const double pt = std::sqrt(pt2);
0142
0143   if (std::fabs(charge * m_magneticFieldTesla) < 1.0e-12)
0144   {
0145     const double scale = -(x * px + y * py) / pt2;
0146     fit.pca_x = x + scale * px;
0147     fit.pca_y = y + scale * py;
0148     fit.pca_z = z + scale * pz;
0149     fit.dca2d = (x * py - y * px) / pt;
0150   }
0151   else
0152   {
0153     const double signed_radius = pt / (0.003 * charge * m_magneticFieldTesla);
0154     const double radius = std::fabs(signed_radius);
0155     if (!good(radius) || radius <= 0.0)
0156     {
0157       return fit;
0158     }
0159
0160     const double tx = px / pt;
0161     const double ty = py / pt;
0162     const double sign = signed_radius > 0.0 ? 1.0 : -1.0;
0163     const double xc = x + sign * radius * ty;
0164     const double yc = y - sign * radius * tx;
0165     const double dc = std::hypot(xc, yc);
0166     if (!good(dc) || dc <= 1.0e-12)
0167     {
0168       return fit;
0169     }
0170
0171     fit.pca_x = xc * (1.0 - radius / dc);
0172     fit.pca_y = yc * (1.0 - radius / dc);
0173
0174     const double phi0 = std::atan2(y - yc, x - xc);
0175     const double pca_phi_on_circle = std::atan2(fit.pca_y - yc, fit.pca_x - xc);
0176     double best_arc = 0.0;
0177     double best_abs_arc = 1.0e30;
0178     for (int k = -4; k <= 4; ++k)
0179     {
0180       const double dphi = pca_phi_on_circle - phi0 + 2.0 * TMath::Pi() * static_cast<double>(k);
0181       const double arc = -sign * radius * dphi;
0182       const double abs_arc = std::fabs(arc);
0183       if (abs_arc < best_abs_arc)
0184       {
0185         best_abs_arc = abs_arc;
0186         best_arc = arc;
0187       }
0188     }
0189
0190     fit.pca_z = z + (pz / pt) * best_arc;
0191     fit.dca2d = dc - radius;
0192   }
0193
0194   fit.pca_radius = std::hypot(fit.pca_x, fit.pca_y);
0195   fit.pca_phi = wrap_phi(std::atan2(fit.pca_y, fit.pca_x));
0196   fit.z0 = fit.pca_z;
0197   fit.track_id = trk->get_track_id();
0198   fit.source_assembled_track_id = trk->get_source_assembled_track_id();
0199   fit.nclusters = trk->get_nclusters();
0200   fit.ok = good(fit.dca2d) && good(fit.z0) && good(fit.pca_radius) && good(fit.pca_phi);
0201   return fit;
0202 }
0203
0204 Tpc_PolyTrackVertexer::CollisionFit
0205 Tpc_PolyTrackVertexer::fitCollision(const std::vector<TrackVertexFit>& tracks) const
0206 {
0207   CollisionFit fit;
0208   if (tracks.size() < 2)
0209   {
0210     return fit;
0211   }
0212
0213   double sum_w = 0.0;
0214   double sum_x = 0.0;
0215   double sum_y = 0.0;
0216   double sum_z = 0.0;
0217   for (const TrackVertexFit& trk : tracks)
0218   {
0219     if (!trk.ok)
0220     {
0221       continue;
0222     }
0223     const double w = std::max(1.0, static_cast<double>(trk.nclusters));
0224     sum_w += w;
0225     sum_x += w * trk.pca_x;
0226     sum_y += w * trk.pca_y;
0227     sum_z += w * trk.z0;
0228   }
0229
0230   if (sum_w <= 0.0)
0231   {
0232     return fit;
0233   }
0234   fit.x = sum_x / sum_w;
0235   fit.y = sum_y / sum_w;
0236   fit.z = sum_z / sum_w;
0237
0238   double sum_res2 = 0.0;
0239   unsigned int nvalid = 0;
0240   for (const TrackVertexFit& trk : tracks)
0241   {
0242     if (!trk.ok)
0243     {
0244       continue;
0245     }
0246     const double w = std::max(1.0, static_cast<double>(trk.nclusters));
0247     const double dz = trk.z0 - fit.z;
0248     sum_res2 += w * dz * dz;
0249     ++nvalid;
0250   }
0251
0252   if (nvalid < 2)
0253   {
0254     return fit;
0255   }
0256   fit.z_rms = std::sqrt(sum_res2 / sum_w);
0257   fit.ntracks = nvalid;
0258   fit.ok = good(fit.x) && good(fit.y) && good(fit.z) && good(fit.z_rms);
0259   return fit;
0260 }
0261
0262 std::vector<Tpc_PolyTrackVertexer::CollisionFit>
0263 Tpc_PolyTrackVertexer::fitCollisions(std::vector<TrackVertexFit> tracks) const
0264 {
0265   std::vector<CollisionFit> collisions;
0266   tracks.erase(std::remove_if(tracks.begin(), tracks.end(),
0267                               [](const TrackVertexFit& trk)
0268                               {
0269                                 return !trk.ok;
0270                               }),
0271                tracks.end());
0272   if (tracks.size() < 2)
0273   {
0274     return collisions;
0275   }
0276
0277   std::sort(tracks.begin(), tracks.end(),
0278             [](const TrackVertexFit& a, const TrackVertexFit& b)
0279             {
0280               return a.z0 < b.z0;
0281             });
0282
0283   const double separation = std::max(0.0, m_collisionZSeparation);
0284   std::vector<TrackVertexFit> cluster;
0285   cluster.reserve(tracks.size());
0286   double cluster_sum_w = 0.0;
0287   double cluster_sum_z = 0.0;
0288
0289   for (const TrackVertexFit& trk : tracks)
0290   {
0291     const double w = std::max(1.0, static_cast<double>(trk.nclusters));
0292     const double cluster_z = cluster_sum_w > 0.0 ? cluster_sum_z / cluster_sum_w : trk.z0;
0293
0294     if (!cluster.empty() && trk.z0 - cluster_z >= separation)
0295     {
0296       const CollisionFit collision = fitCollision(cluster);
0297       if (collision.ok)
0298       {
0299         collisions.push_back(collision);
0300       }
0301       cluster.clear();
0302       cluster_sum_w = 0.0;
0303       cluster_sum_z = 0.0;
0304     }
0305
0306     cluster.push_back(trk);
0307     cluster_sum_w += w;
0308     cluster_sum_z += w * trk.z0;
0309   }
0310
0311   const CollisionFit collision = fitCollision(cluster);
0312   if (collision.ok)
0313   {
0314     collisions.push_back(collision);
0315   }
0316   return collisions;
0317 }
0318
0319 int Tpc_PolyTrackVertexer::process_event(PHCompositeNode* topNode)
0320 {
0321   if (!m_polyTracks || !m_vertices)
0322   {
0323     if (!getNodes(topNode) || !createNodes(topNode))
0324     {
0325       return Fun4AllReturnCodes::EVENT_OK;
0326     }
0327   }
0328
0329   m_vertices->Reset();
0330
0331   std::vector<TrackVertexFit> collision_tracks;
0332   const unsigned int ntracks = m_polyTracks ? m_polyTracks->size() : 0;
0333   for (unsigned int itrk = 0; itrk < ntracks; ++itrk)
0334   {
0335     const Tpc_PolyTrack* trk = m_polyTracks->get_track(itrk);
0336     const TrackVertexFit fit = fitTrack(trk);
0337     if (!fit.ok)
0338     {
0339       continue;
0340     }
0341
0342     Tpc_PolyTrackVertexv1* out = new Tpc_PolyTrackVertexv1();
0343     out->set_track_id(fit.track_id);
0344     out->set_source_assembled_track_id(fit.source_assembled_track_id);
0345     out->set_dca2d(fit.dca2d);
0346     out->set_z0(fit.z0);
0347     out->set_pca_valid(1);
0348     out->set_pca_x(fit.pca_x);
0349     out->set_pca_y(fit.pca_y);
0350     out->set_pca_z(fit.pca_z);
0351     out->set_pca_radius(fit.pca_radius);
0352     out->set_pca_phi(fit.pca_phi);
0353     m_vertices->add_vertex(out);
0354
0355     if (fit.nclusters >= m_collisionMinClusters)
0356     {
0357       collision_tracks.push_back(fit);
0358     }
0359   }
0360
0361   const std::vector<CollisionFit> collisions = fitCollisions(collision_tracks);
0362   m_vertices->set_collision_min_clusters(m_collisionMinClusters);
0363   m_vertices->clear_collision_vertices();
0364   for (const CollisionFit& collision : collisions)
0365   {
0366     m_vertices->add_collision_vertex(collision.x, collision.y, collision.z,
0367                                      collision.z_rms, collision.ntracks);
0368   }
0369   m_vertices->set_collision_vertex_valid(collisions.empty() ? 0 : 1);
0370
0371   if (Verbosity() > 0)
0372   {
0373     std::cout << Name() << "::process_event - input poly_tracks=" << ntracks
0374               << " vertices=" << m_vertices->size()
0375               << " collision_vertices=" << m_vertices->get_collision_vertex_count()
0376               << " first_collision_ntracks=" << m_vertices->get_collision_ntracks()
0377               << " first_collision_x=" << m_vertices->get_collision_x()
0378               << " first_collision_y=" << m_vertices->get_collision_y()
0379               << " first_collision_z=" << m_vertices->get_collision_z()
0380               << std::endl;
0381   }
0382
0383   return Fun4AllReturnCodes::EVENT_OK;
0384 }