File indexing completed on 2026-08-31 08:21:42
0001 #include "Tpc_PolyClusterDisplay.h"
0002
0003 #include "tpctrackreco/Tpc_PolyCluster.h"
0004 #include "tpctrackreco/Tpc_PolyClusterContainer.h"
0005 #include "tpctrackreco/Tpc_PolyTrack.h"
0006 #include "tpctrackreco/Tpc_PolyTrackContainer.h"
0007 #include "tpctrackreco/Tpc_PolyTrackVertex.h"
0008 #include "tpctrackreco/Tpc_PolyTrackVertexContainer.h"
0009
0010 #include <fun4all/Fun4AllReturnCodes.h>
0011 #include <phool/PHCompositeNode.h>
0012 #include <phool/getClass.h>
0013
0014 #include <TCanvas.h>
0015 #include <TDirectory.h>
0016 #include <TFile.h>
0017 #include <TH3D.h>
0018 #include <TPolyLine3D.h>
0019 #include <TPolyMarker3D.h>
0020
0021 #include <algorithm>
0022 #include <cmath>
0023 #include <format>
0024 #include <iostream>
0025 #include <limits>
0026 #include <map>
0027 #include <vector>
0028
0029 namespace
0030 {
0031 int cluster_color(const unsigned int icluster)
0032 {
0033 static const int colors[] = {
0034 kRed + 1, kBlue + 1, kGreen + 2, kMagenta + 1, kCyan + 2,
0035 kOrange + 7, kViolet + 1, kAzure + 1, kPink + 7, kTeal + 3};
0036 return colors[icluster % (sizeof(colors) / sizeof(colors[0]))];
0037 }
0038
0039 bool cluster_group_z_range(const std::vector<const Tpc_PolyCluster*>& clusters,
0040 const double display_zmin,
0041 const double display_zmax,
0042 double& zmin,
0043 double& zmax)
0044 {
0045 if (clusters.empty())
0046 {
0047 return false;
0048 }
0049
0050 zmin = std::numeric_limits<double>::max();
0051 zmax = -std::numeric_limits<double>::max();
0052
0053 auto update_range = [&](const double z)
0054 {
0055 if (!std::isfinite(z))
0056 {
0057 return;
0058 }
0059 zmin = std::min(zmin, z);
0060 zmax = std::max(zmax, z);
0061 };
0062
0063 for (const Tpc_PolyCluster* cluster : clusters)
0064 {
0065 if (!cluster || !cluster->isValid())
0066 {
0067 continue;
0068 }
0069 update_range(cluster->get_centroid_z());
0070 }
0071
0072 if (zmin == std::numeric_limits<double>::max() || zmax == -std::numeric_limits<double>::max())
0073 {
0074 return false;
0075 }
0076
0077 zmin = std::max(zmin, display_zmin);
0078 zmax = std::min(zmax, display_zmax);
0079 if (zmin > zmax)
0080 {
0081 return false;
0082 }
0083
0084 if (zmin == zmax)
0085 {
0086 zmin = std::max(zmin - 0.1, display_zmin);
0087 zmax = std::min(zmax + 0.1, display_zmax);
0088 }
0089
0090 return zmin < zmax;
0091 }
0092
0093 bool track_vertex_z_selected(const Tpc_PolyTrackVertex* vtx,
0094 const double vertex_zmin,
0095 const double vertex_zmax)
0096 {
0097 if (!vtx)
0098 {
0099 return false;
0100 }
0101 const double z0 = vtx->get_z0();
0102 return std::isfinite(z0) && z0 >= vertex_zmin && z0 <= vertex_zmax;
0103 }
0104
0105 bool poly_track_xy_at_z(const Tpc_PolyTrack* trk,
0106 const double z,
0107 const double magnetic_field_tesla,
0108 const double arc_direction,
0109 const bool use_straight_line,
0110 double& x,
0111 double& y)
0112 {
0113 if (!trk || trk->get_fit_status() == 0 || !std::isfinite(z))
0114 {
0115 return false;
0116 }
0117
0118 const double x0 = trk->get_x();
0119 const double y0 = trk->get_y();
0120 const double z0 = trk->get_z();
0121 const double px = trk->get_px();
0122 const double py = trk->get_py();
0123 const double pz = trk->get_pz();
0124 const double charge = trk->get_charge();
0125 if (!std::isfinite(x0) || !std::isfinite(y0) || !std::isfinite(z0) ||
0126 !std::isfinite(px) || !std::isfinite(py) || !std::isfinite(pz) ||
0127 !std::isfinite(charge))
0128 {
0129 return false;
0130 }
0131
0132 const double dz = z - z0;
0133 if (use_straight_line || std::fabs(charge * magnetic_field_tesla) < 1.0e-12)
0134 {
0135 if (std::fabs(pz) < 1.0e-12)
0136 {
0137 return false;
0138 }
0139 x = x0 + arc_direction * px / pz * dz;
0140 y = y0 + arc_direction * py / pz * dz;
0141 return std::isfinite(x) && std::isfinite(y);
0142 }
0143
0144 const double pt = std::hypot(px, py);
0145 if (pt <= 0.0 || std::fabs(pz) < 1.0e-12)
0146 {
0147 return false;
0148 }
0149
0150 const double signed_radius = pt / (0.003 * charge * magnetic_field_tesla);
0151 const double radius = std::fabs(signed_radius);
0152 if (!std::isfinite(radius) || radius <= 0.0)
0153 {
0154 return false;
0155 }
0156
0157 const double tx = px / pt;
0158 const double ty = py / pt;
0159 const double sign = signed_radius > 0.0 ? 1.0 : -1.0;
0160 const double xc = x0 + sign * radius * ty;
0161 const double yc = y0 - sign * radius * tx;
0162 const double phi0 = std::atan2(y0 - yc, x0 - xc);
0163 const double dzds = pz / pt;
0164 if (std::fabs(dzds) < 1.0e-12)
0165 {
0166 return false;
0167 }
0168
0169 const double arc = arc_direction * dz / dzds;
0170 const double phi = phi0 - sign * arc / radius;
0171 x = xc + radius * std::cos(phi);
0172 y = yc + radius * std::sin(phi);
0173 return std::isfinite(x) && std::isfinite(y);
0174 }
0175
0176 double cluster_line_residual2(const Tpc_PolyTrack* poly_track,
0177 const std::vector<const Tpc_PolyCluster*>& clusters,
0178 const double magnetic_field_tesla,
0179 const double arc_direction,
0180 const bool use_straight_line)
0181 {
0182 if (!poly_track || clusters.empty())
0183 {
0184 return std::numeric_limits<double>::max();
0185 }
0186
0187 double sum = 0.0;
0188 unsigned int n = 0;
0189 for (const Tpc_PolyCluster* cluster : clusters)
0190 {
0191 if (!cluster || !cluster->isValid())
0192 {
0193 continue;
0194 }
0195 const double cx = cluster->get_centroid_x();
0196 const double cy = cluster->get_centroid_y();
0197 const double cz = cluster->get_centroid_z();
0198 if (!std::isfinite(cx) || !std::isfinite(cy) || !std::isfinite(cz))
0199 {
0200 continue;
0201 }
0202
0203 double x = 0.0;
0204 double y = 0.0;
0205 if (!poly_track_xy_at_z(poly_track, cz, magnetic_field_tesla, arc_direction, use_straight_line, x, y))
0206 {
0207 continue;
0208 }
0209
0210 const double dx = x - cx;
0211 const double dy = y - cy;
0212 sum += dx * dx + dy * dy;
0213 ++n;
0214 }
0215
0216 return n > 0 ? sum / static_cast<double>(n) : std::numeric_limits<double>::max();
0217 }
0218
0219 TPolyLine3D* make_poly_track_line(const Tpc_PolyTrack* trk,
0220 const double zmin,
0221 const double zmax,
0222 const double xymax,
0223 const double magnetic_field_tesla,
0224 const double arc_direction,
0225 const bool use_straight_line,
0226 const int color)
0227 {
0228 if (!trk || trk->get_fit_status() == 0)
0229 {
0230 return nullptr;
0231 }
0232
0233 std::vector<double> zs;
0234 std::vector<double> xs;
0235 std::vector<double> ys;
0236 const unsigned int nsteps = 80;
0237 zs.reserve(nsteps + 1);
0238 xs.reserve(nsteps + 1);
0239 ys.reserve(nsteps + 1);
0240
0241 for (unsigned int istep = 0; istep <= nsteps; ++istep)
0242 {
0243 const double f = static_cast<double>(istep) / static_cast<double>(nsteps);
0244 const double z = zmin + f * (zmax - zmin);
0245 double x = 0.0;
0246 double y = 0.0;
0247 if (!poly_track_xy_at_z(trk, z, magnetic_field_tesla, arc_direction, use_straight_line, x, y))
0248 {
0249 continue;
0250 }
0251 if (std::fabs(x) > xymax || std::fabs(y) > xymax)
0252 {
0253 continue;
0254 }
0255 zs.push_back(z);
0256 xs.push_back(x);
0257 ys.push_back(y);
0258 }
0259
0260 if (zs.size() < 2)
0261 {
0262 return nullptr;
0263 }
0264 TPolyLine3D* line = new TPolyLine3D(static_cast<int>(zs.size()));
0265 for (unsigned int i = 0; i < zs.size(); ++i)
0266 {
0267 line->SetPoint(static_cast<int>(i), zs[i], xs[i], ys[i]);
0268 }
0269 line->SetLineColor(color);
0270 line->SetLineWidth(3);
0271 return line;
0272 }
0273
0274 TPolyMarker3D* make_cluster_marker(const double z,
0275 const double x,
0276 const double y,
0277 const int color)
0278 {
0279 TPolyMarker3D* marker = new TPolyMarker3D(1);
0280 marker->SetPoint(0, z, x, y);
0281 marker->SetMarkerColor(color);
0282 marker->SetMarkerStyle(20);
0283 marker->SetMarkerSize(1.4);
0284 return marker;
0285 }
0286
0287 TPolyMarker3D* make_pca_marker(const Tpc_PolyTrackVertex* vtx,
0288 const double zmin,
0289 const double zmax,
0290 const double xymax,
0291 const int color)
0292 {
0293 if (!vtx || !vtx->get_pca_valid())
0294 {
0295 return nullptr;
0296 }
0297
0298 const double x = vtx->get_pca_x();
0299 const double y = vtx->get_pca_y();
0300 const double z = vtx->get_pca_z();
0301 if (!std::isfinite(x) || !std::isfinite(y) || !std::isfinite(z))
0302 {
0303 return nullptr;
0304 }
0305 if (z < zmin || z > zmax)
0306 {
0307 return nullptr;
0308 }
0309 if (std::fabs(x) > xymax || std::fabs(y) > xymax)
0310 {
0311 return nullptr;
0312 }
0313
0314 TPolyMarker3D* marker = new TPolyMarker3D(1);
0315 marker->SetPoint(0, z, x, y);
0316 marker->SetMarkerColor(color);
0317 marker->SetMarkerStyle(20);
0318 marker->SetMarkerSize(1.2);
0319 return marker;
0320 }
0321
0322 TPolyMarker3D* make_collision_vertex_marker(const double x,
0323 const double y,
0324 const double z,
0325 const double zmin,
0326 const double zmax,
0327 const double xymax)
0328 {
0329 if (!std::isfinite(x) || !std::isfinite(y) || !std::isfinite(z))
0330 {
0331 return nullptr;
0332 }
0333 if (z < zmin || z > zmax)
0334 {
0335 return nullptr;
0336 }
0337 if (std::fabs(x) > xymax || std::fabs(y) > xymax)
0338 {
0339 return nullptr;
0340 }
0341
0342 TPolyMarker3D* marker = new TPolyMarker3D(1);
0343 marker->SetPoint(0, z, x, y);
0344 marker->SetMarkerColor(kBlack);
0345 marker->SetMarkerStyle(29);
0346 marker->SetMarkerSize(2.0);
0347 return marker;
0348 }
0349 }
0350
0351 Tpc_PolyClusterDisplay::Tpc_PolyClusterDisplay(const std::string& name,
0352 const std::string& outfilename,
0353 const std::string& clusterNodeName,
0354 const unsigned int maxEventDisplays)
0355 : SubsysReco(name)
0356 , m_outfilename(outfilename)
0357 , m_clusterNodeName(clusterNodeName)
0358 , m_finalTrackNodeName("TPC_POLYTRACKS")
0359 , m_finalTrackVertexNodeName("TPC_POLYTRACKVERTICES")
0360 , m_maxEventDisplays(maxEventDisplays)
0361 , m_evt(0)
0362 , m_eventsSaved(0)
0363 , m_zmin(-102.0)
0364 , m_zmax(102.0)
0365 , m_trackVertexZMin(-20.0)
0366 , m_trackVertexZMax(20.0)
0367 , m_xymax(85.0)
0368 , m_magneticFieldTesla(1.4)
0369 , m_useStraightLineTracks(false)
0370 , m_outfile(nullptr)
0371 , m_clusters(nullptr)
0372 , m_finalTracks(nullptr)
0373 , m_finalTrackVertices(nullptr)
0374 {
0375 }
0376
0377 Tpc_PolyClusterDisplay::~Tpc_PolyClusterDisplay()
0378 {
0379 if (m_outfile)
0380 {
0381 delete m_outfile;
0382 m_outfile = nullptr;
0383 }
0384 }
0385
0386 int Tpc_PolyClusterDisplay::Init(PHCompositeNode* )
0387 {
0388
0389 m_outfile = new TFile(m_outfilename.c_str(), "RECREATE");
0390 if (!m_outfile || m_outfile->IsZombie())
0391 {
0392 std::cerr << Name() << "::Init - cannot open output file " << m_outfilename << std::endl;
0393 return Fun4AllReturnCodes::ABORTRUN;
0394 }
0395 m_outfile->mkdir("events");
0396 return Fun4AllReturnCodes::EVENT_OK;
0397 }
0398
0399 bool Tpc_PolyClusterDisplay::get_nodes(PHCompositeNode* topNode)
0400 {
0401 m_clusters = findNode::getClass<Tpc_PolyClusterContainer>(topNode, m_clusterNodeName);
0402 if (!m_clusters)
0403 {
0404 std::cerr << Name() << " - missing " << m_clusterNodeName << std::endl;
0405 return false;
0406 }
0407 m_finalTracks = findNode::getClass<Tpc_PolyTrackContainer>(topNode, m_finalTrackNodeName);
0408 if (!m_finalTracks)
0409 {
0410 std::cerr << Name() << " - missing " << m_finalTrackNodeName << ", drawing clusters without fitted lines" << std::endl;
0411 }
0412 m_finalTrackVertices = findNode::getClass<Tpc_PolyTrackVertexContainer>(topNode, m_finalTrackVertexNodeName);
0413 if (!m_finalTrackVertices)
0414 {
0415 std::cerr << Name() << " - missing " << m_finalTrackVertexNodeName
0416 << ", drawing without final-track PCA/collision vertices" << std::endl;
0417 }
0418 return true;
0419 }
0420
0421 int Tpc_PolyClusterDisplay::process_event(PHCompositeNode* topNode)
0422 {
0423 ++m_evt;
0424 if (!get_nodes(topNode))
0425 {
0426 return Fun4AllReturnCodes::EVENT_OK;
0427 }
0428 if (!m_outfile || m_eventsSaved >= m_maxEventDisplays)
0429 {
0430 return Fun4AllReturnCodes::EVENT_OK;
0431 }
0432
0433 TDirectory* eventsTop = m_outfile->GetDirectory("events");
0434 if (!eventsTop)
0435 {
0436 eventsTop = m_outfile->mkdir("events");
0437 }
0438 eventsTop->cd();
0439
0440 TDirectory* eventDir = eventsTop->mkdir(std::format("event_{:06}", m_evt).c_str());
0441 if (!eventDir)
0442 {
0443 return Fun4AllReturnCodes::EVENT_OK;
0444 }
0445 eventDir->cd();
0446
0447 TH3D* h3 = new TH3D(std::format("h3_evt{:06}_tpc_polycluster_centroids_bothsides", m_evt).c_str(),
0448 std::format("event {} both sides Tpc_PolyCluster centroids;z [cm];x [cm];y [cm]", m_evt).c_str(),
0449 204, m_zmin, m_zmax,
0450 170, -m_xymax, m_xymax,
0451 170, -m_xymax, m_xymax);
0452 h3->SetStats(false);
0453 h3->SetDirectory(nullptr);
0454
0455 std::vector<TPolyMarker3D*> markers;
0456 std::vector<TPolyMarker3D*> pca_markers;
0457 std::vector<TPolyMarker3D*> collision_vertex_markers;
0458 std::vector<TPolyLine3D*> lines;
0459 std::map<unsigned int, const Tpc_PolyTrackVertex*> poly_track_vertices_by_assembled_track_id;
0460 std::map<unsigned int, std::vector<const Tpc_PolyCluster*> > clusters_by_assembled_track_id;
0461 const unsigned int nvertices = m_finalTrackVertices ? m_finalTrackVertices->size() : 0;
0462 for (unsigned int ivtx = 0; ivtx < nvertices; ++ivtx)
0463 {
0464 const Tpc_PolyTrackVertex* vtx = m_finalTrackVertices->get_vertex(ivtx);
0465 if (!vtx)
0466 {
0467 continue;
0468 }
0469 poly_track_vertices_by_assembled_track_id[vtx->get_source_assembled_track_id()] = vtx;
0470 }
0471
0472 const bool applyTrackVertexZRange = m_finalTrackVertices != nullptr;
0473 const unsigned int nclusters_total = m_clusters->size();
0474 unsigned int nclusters = 0;
0475 unsigned int nclustersPlotted = 0;
0476 for (unsigned int icluster = 0; icluster < nclusters_total; ++icluster)
0477 {
0478 const Tpc_PolyCluster* cluster = m_clusters->get_cluster(icluster);
0479 if (!cluster || !cluster->isValid())
0480 {
0481 continue;
0482 }
0483 const unsigned int source_id = cluster->get_source_assembled_track_id();
0484 const auto vertex_iter = poly_track_vertices_by_assembled_track_id.find(source_id);
0485 if (applyTrackVertexZRange &&
0486 (vertex_iter == poly_track_vertices_by_assembled_track_id.end() ||
0487 !track_vertex_z_selected(vertex_iter->second, m_trackVertexZMin, m_trackVertexZMax)))
0488 {
0489 continue;
0490 }
0491 clusters_by_assembled_track_id[source_id].push_back(cluster);
0492
0493 ++nclusters;
0494 const double x = cluster->get_centroid_x();
0495 const double y = cluster->get_centroid_y();
0496 const double z = cluster->get_centroid_z();
0497 if (!std::isfinite(x) || !std::isfinite(y) || !std::isfinite(z))
0498 {
0499 continue;
0500 }
0501 if (z < m_zmin || z > m_zmax)
0502 {
0503 continue;
0504 }
0505 markers.push_back(make_cluster_marker(z, x, y, cluster_color(source_id)));
0506
0507 ++nclustersPlotted;
0508 }
0509
0510 const unsigned int npoly_tracks = m_finalTracks ? m_finalTracks->size() : 0;
0511 for (unsigned int ifinal = 0; ifinal < npoly_tracks; ++ifinal)
0512 {
0513 const Tpc_PolyTrack* trk = m_finalTracks->get_track(ifinal);
0514 if (!trk)
0515 {
0516 continue;
0517 }
0518
0519 const auto cluster_iter = clusters_by_assembled_track_id.find(trk->get_source_assembled_track_id());
0520 if (cluster_iter == clusters_by_assembled_track_id.end())
0521 {
0522 continue;
0523 }
0524
0525 double line_zmin = m_zmin;
0526 double line_zmax = m_zmax;
0527 if (!cluster_group_z_range(cluster_iter->second, m_zmin, m_zmax, line_zmin, line_zmax))
0528 {
0529 continue;
0530 }
0531
0532 const bool use_straight_line = m_useStraightLineTracks || std::fabs(trk->get_charge() * m_magneticFieldTesla) < 1.0e-12;
0533 const double forward_residual2 = cluster_line_residual2(trk, cluster_iter->second, m_magneticFieldTesla, 1.0, use_straight_line);
0534 const double reverse_residual2 = cluster_line_residual2(trk, cluster_iter->second, m_magneticFieldTesla, -1.0, use_straight_line);
0535 const double arc_direction = forward_residual2 <= reverse_residual2 ? 1.0 : -1.0;
0536
0537 TPolyLine3D* line = make_poly_track_line(trk, line_zmin, line_zmax, m_xymax,
0538 m_magneticFieldTesla, arc_direction,
0539 use_straight_line,
0540 cluster_color(trk->get_source_assembled_track_id()));
0541 if (line)
0542 {
0543 lines.push_back(line);
0544 }
0545 }
0546
0547 for (unsigned int ivtx = 0; ivtx < nvertices; ++ivtx)
0548 {
0549 const Tpc_PolyTrackVertex* vtx = m_finalTrackVertices->get_vertex(ivtx);
0550 if (!vtx)
0551 {
0552 continue;
0553 }
0554 if (!track_vertex_z_selected(vtx, m_trackVertexZMin, m_trackVertexZMax))
0555 {
0556 continue;
0557 }
0558
0559 int color = cluster_color(vtx->get_track_id());
0560 const auto cluster_iter = clusters_by_assembled_track_id.find(vtx->get_source_assembled_track_id());
0561 if (cluster_iter != clusters_by_assembled_track_id.end())
0562 {
0563 color = cluster_color(vtx->get_source_assembled_track_id());
0564 }
0565
0566 TPolyMarker3D* marker = make_pca_marker(vtx, m_zmin, m_zmax, m_xymax, color);
0567 if (marker)
0568 {
0569 pca_markers.push_back(marker);
0570 }
0571 }
0572
0573 const unsigned int ncollision_vertices = (m_finalTrackVertices && m_finalTrackVertices->get_collision_vertex_valid())
0574 ? m_finalTrackVertices->get_collision_vertex_count()
0575 : 0;
0576 for (unsigned int ivtx = 0; ivtx < ncollision_vertices; ++ivtx)
0577 {
0578 TPolyMarker3D* marker = make_collision_vertex_marker(m_finalTrackVertices->get_collision_x(ivtx),
0579 m_finalTrackVertices->get_collision_y(ivtx),
0580 m_finalTrackVertices->get_collision_z(ivtx),
0581 m_zmin, m_zmax, m_xymax);
0582 if (marker)
0583 {
0584 collision_vertex_markers.push_back(marker);
0585 }
0586 }
0587
0588 TCanvas* c3 = new TCanvas(std::format("c3_evt{:06}_tpc_polycluster_z_x_y_centroids_bothsides", m_evt).c_str(),
0589 std::format("event {} both sides Tpc_PolyCluster centroids", m_evt).c_str(),
0590 1200, 900);
0591 h3->Draw();
0592 for (TPolyLine3D* line : lines)
0593 {
0594 if (line)
0595 {
0596 line->Draw("same");
0597 }
0598 }
0599 for (TPolyMarker3D* marker : markers)
0600 {
0601 if (marker)
0602 {
0603 marker->Draw("same");
0604 }
0605 }
0606 for (TPolyMarker3D* marker : pca_markers)
0607 {
0608 if (marker)
0609 {
0610 marker->Draw("same");
0611 }
0612 }
0613 for (TPolyMarker3D* marker : collision_vertex_markers)
0614 {
0615 if (marker)
0616 {
0617 marker->Draw("same");
0618 }
0619 }
0620 c3->Modified();
0621 c3->Update();
0622 c3->Write();
0623
0624 std::cout << Name() << " - saved event " << m_evt
0625 << " clusters=" << nclusters_total
0626 << " selected_clusters=" << nclusters
0627 << " plotted=" << nclustersPlotted
0628 << " poly_tracks=" << npoly_tracks
0629 << " lines=" << lines.size()
0630 << " poly_track_vertices=" << nvertices
0631 << " pca_markers=" << pca_markers.size()
0632 << " collision_vertices=" << ncollision_vertices
0633 << " collision_markers=" << collision_vertex_markers.size() << std::endl;
0634
0635 ++m_eventsSaved;
0636 return Fun4AllReturnCodes::EVENT_OK;
0637 }
0638
0639 int Tpc_PolyClusterDisplay::End(PHCompositeNode* )
0640 {
0641 if (m_outfile)
0642 {
0643 m_outfile->Close();
0644 delete m_outfile;
0645 m_outfile = nullptr;
0646 }
0647 return Fun4AllReturnCodes::EVENT_OK;
0648 }
