packages/tpccalib/SiliconDriftEvaluator.cc

0001 #include "SiliconDriftEvaluator.h"
0002
0003 #include <fun4all/Fun4AllReturnCodes.h>
0004 #include <phool/PHCompositeNode.h>
0005 #include <phool/PHNodeIterator.h>
0006 #include <phool/getClass.h>
0007 #include <trackbase/TrkrDefs.h>
0008 #include <trackbase_historic/SvtxTrack.h>
0009 #include <trackbase_historic/SvtxTrackMap.h>
0010 #include <trackbase_historic/TrackSeedHelper.h>
0011
0012 #include <TCanvas.h>
0013 #include <TDirectory.h>
0014 #include <TF1.h>
0015 #include <TF2.h>
0016 #include <TFile.h>
0017 #include <TH1F.h>
0018 #include <TH2F.h>
0019 #include <TH3F.h>
0020 #include <TLegend.h>
0021 #include <TLine.h>
0022 #include <TParameter.h>
0023
0024 #include <cassert>
0025 #include <climits>
0026 #include <cmath>
0027 #include <format>
0028 #include <iostream>
0029 #include <memory>
0030
0031 //_____________________________________________________________________
0032 namespace
0033 {
0034
0035   //! pt
0036   template <class T>
0037   T get_pt(const T& px, const T& py)
0038   {
0039     return std::sqrt(px * px + py * py);
0040   }
0041
0042   //_____________________________________________________________________
0043   //  par[0] = constrained slope
0044   //  par[1] = offset for eta < 0
0045   //  par[2] = offset for eta >= 0
0046   //
0047   double fit_function_2d(double* x, double* par)
0048   {
0049     const int ieta = static_cast<int>(std::floor(x[0]));
0050     const double z = x[1];
0051     if (ieta < 0 || ieta > 1)
0052     {
0053       TF2::RejectPoint();
0054       return 0.;
0055     }
0056     return par[ieta + 1] + par[0] * z;
0057   }
0058
0059   //! 1D version used to draw per-eta overlay lines on QA canvas
0060   // NOLINTNEXTLINE(readability-non-const-parameter): ROOT TF1 requires this exact signature
0061   double linear_function(double* x, double* par)
0062   {
0063     return par[0] * x[0] + par[1];
0064   }
0065
0066   //! human-readable label for each eta bin
0067   const char* k_eta_labels[2] = {"#eta_{TPC} < 0", "#eta_{TPC} #geq 0"};
0068
0069 }  // namespace
0070
0071 //_____________________________________________________________________
0072 SiliconDriftEvaluator::SiliconDriftEvaluator(const std::string& name)
0073   : SubsysReco(name)
0074 {
0075 }
0076
0077 //_____________________________________________________________________
0078 int SiliconDriftEvaluator::Init(PHCompositeNode* topNode)
0079 {
0080   // find DST node
0081   PHNodeIterator iter(topNode);
0082   auto* dstNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0083   if (!dstNode)
0084   {
0085     std::cout << "SiliconDriftEvaluator::Init - DST Node missing" << std::endl;
0086     return Fun4AllReturnCodes::ABORTEVENT;
0087   }
0088
0089   // get EVAL node
0090   iter = PHNodeIterator(dstNode);
0091   auto* evalNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "EVAL"));
0092   if (!evalNode)
0093   {
0094     // create
0095     std::cout << "SiliconDriftEvaluator::Init - EVAL node missing - creating" << std::endl;
0096     evalNode = new PHCompositeNode("EVAL");
0097     dstNode->addNode(evalNode);
0098   }
0099
0100   // add container to output tree
0101   auto* newNode = new PHIODataNode<PHObject>(new Container, "SiliconDriftEvaluator::Container", "PHObject");
0102
0103   // overwrite split level for easier offline browsing
0104   newNode->SplitLevel(99);
0105   evalNode->addNode(newNode);
0106
0107   // book 3D accumulator histogram
0108   // x = eta bin: 0 = eta<0, 1 = eta>=0
0109   // y = z_si (cm)
0110   // z = dz   (cm)
0111   m_hist3D = new TH3F("SiliconDriftEval_hist3D", ";#eta bin;z_{silicon} (cm);#Deltaz_{TPC-silicon} (cm)", 2, 0, 2, 200, -m_max_z, m_max_z, 200, -m_max_dz, m_max_dz);
0112   m_hist3D->SetDirectory(nullptr);
0113
0114   return Fun4AllReturnCodes::EVENT_OK;
0115 }
0116
0117 //_____________________________________________________________________
0118 int SiliconDriftEvaluator::InitRun(PHCompositeNode* topNode)
0119 {
0120   return load_nodes(topNode);
0121 }
0122
0123 //_____________________________________________________________________
0124 int SiliconDriftEvaluator::process_event(PHCompositeNode* topNode)
0125 {
0126   // load nodes
0127   const auto res = load_nodes(topNode);
0128   if (res != Fun4AllReturnCodes::EVENT_OK)
0129   {
0130     return res;
0131   }
0132
0133   // cleanup output
0134   if (m_container)
0135   {
0136     m_container->Reset();
0137   }
0138
0139   evaluate_tracks();
0140
0141   return Fun4AllReturnCodes::EVENT_OK;
0142 }
0143
0144 //_____________________________________________________________________
0145 int SiliconDriftEvaluator::End(PHCompositeNode* /*topNode*/)
0146 {
0147   if (!m_hist3D)
0148   {
0149     std::cerr << Name() << "::End - histogram not found, skipping fit." << std::endl;
0150     return Fun4AllReturnCodes::EVENT_OK;
0151   }
0152
0153   const int nEntries = static_cast<int>(m_hist3D->GetEntries());
0154   std::cout << Name() << "::End - fitting " << nEntries << " entries" << std::endl;
0155
0156   // build mean-dz TH2F via FitSlicesY, one eta bin at a time
0157   // x = eta bin [0,2), y = z_si (cm), content = mean dz (cm)
0158   auto* h_fit = new TH2F("h_fit_silicon", "",
0159                          2, 0, 2,
0160                          200, -m_max_z, m_max_z);
0161   h_fit->SetDirectory(nullptr);
0162
0163   for (int ieta = 0; ieta < 2; ++ieta)
0164   {
0165     m_hist3D->GetXaxis()->SetRange(ieta + 1, ieta + 1);
0166     auto* h2d = static_cast<TH2F*>(m_hist3D->Project3D("zy"));
0167     h2d->SetName(std::format("h2d_etabin_{}", ieta).c_str());
0168     h2d->SetDirectory(nullptr);
0169
0170     // fit vertical slices; require a minimum of m_min_slice_entries per slice
0171     h2d->FitSlicesY(nullptr, 0, -1, m_min_slice_entries);
0172     auto* h_mean = static_cast<TH1F*>(gDirectory->Get(std::format("h2d_etabin_{}_1", ieta).c_str()));
0173
0174     if (!h_mean)
0175     {
0176       delete h2d;
0177       continue;
0178     }
0179
0180     for (int iz = 1; iz <= h_mean->GetNbinsX(); ++iz)
0181     {
0182       const double entries = h2d->Integral(iz, iz, 1, m_hist3D->GetNbinsZ());
0183       if (entries > 0)
0184       {
0185         h_fit->SetBinContent(ieta + 1, iz, h_mean->GetBinContent(iz));
0186       }
0187     }
0188
0189     delete h2d;
0190   }
0191
0192   m_hist3D->GetXaxis()->SetRange(0, 0);
0193
0194   // 2D piecewise fit: shared slope + per-eta offset
0195   auto* fit2d = new TF2("fit2d_silicon", fit_function_2d, 0, 2, -m_max_z, m_max_z, 3);
0196   for (int i = 0; i < 3; ++i)
0197   {
0198     fit2d->SetParameter(i, 0.0);
0199   }
0200   h_fit->Fit(fit2d, "0R");
0201
0202   const double slope = fit2d->GetParameter(0);
0203   const double slope_err = fit2d->GetParError(0);
0204   const double off_neg = fit2d->GetParameter(1);  // ieta=0, eta<0
0205   const double off_pos = fit2d->GetParameter(2);  // ieta=1, eta>=0
0206
0207   const double dv_new = m_drift_velocity / (1.0 + slope);
0208   const double dv_err = m_drift_velocity / std::pow(1.0 + slope, 2) * slope_err;
0209   const double t0_new = (off_pos - off_neg) / (2.0 * dv_new);
0210
0211   std::cout << Name() << "::End" << " slope=" << slope << " dv_in=" << m_drift_velocity << " cm/ns" << " dv_new=" << dv_new << " +/- " << dv_err << " cm/ns" << " t0_new=" << t0_new << " ns" << std::endl;
0212
0213   // draw the plot
0214   auto* canvas = new TCanvas("silicon_drift_calib", "Silicon drift velocity calibration", 1400, 700);
0215   canvas->Divide(2, 1);
0216
0217   for (int ieta = 0; ieta < 2; ++ieta)
0218   {
0219     canvas->cd(ieta + 1);
0220     gPad->SetTopMargin(0.13);
0221     gPad->SetRightMargin(0.18);
0222
0223     // 2D distribution for this eta bin
0224     m_hist3D->GetXaxis()->SetRange(ieta + 1, ieta + 1);
0225     auto* h2d = static_cast<TH2F*>(m_hist3D->Project3D("zy"));
0226     h2d->SetName(std::format("hplot_etabin_{}", ieta).c_str());
0227     h2d->SetTitle(";z_{silicon} (cm);#Deltaz_{TPC-silicon} (cm)");
0228     h2d->SetStats(false);
0229     h2d->Draw("COLZ");
0230
0231     // mean-dz points from FitSlicesY
0232     auto* h_fit_proj = h_fit->ProjectionY(std::format("h_fit_proj_{}", ieta).c_str(), ieta + 1, ieta + 1);
0233     h_fit_proj->SetMarkerStyle(20);
0234     h_fit_proj->SetMarkerSize(0.6);
0235     h_fit_proj->SetMarkerColor(kRed);
0236     h_fit_proj->SetLineColor(kRed);
0237     h_fit_proj->Draw("same P");
0238
0239     // 1D fit line for this eta bin
0240     auto* f1d = new TF1(std::format("f1d_etabin_{}", ieta).c_str(), linear_function, -m_max_z, m_max_z, 2);
0241     f1d->SetParameter(0, slope);
0242     f1d->SetParameter(1, (ieta == 0) ? off_neg : off_pos);
0243     f1d->SetLineColor(kGreen + 2);
0244     f1d->SetLineWidth(2);
0245     f1d->Draw("same");
0246
0247     // reference line at dz = 0
0248     auto* zero = new TLine(-m_max_z, 0, m_max_z, 0);
0249     zero->SetLineStyle(2);
0250     zero->SetLineColor(kGray + 1);
0251     zero->Draw();
0252
0253     auto* leg = new TLegend(0.13, 0.76, 0.82, 0.95);
0254     leg->SetBorderSize(0);
0255     leg->SetFillStyle(0);
0256     leg->SetTextSize(0.033);
0257     leg->SetHeader(std::format("{}   entries: {}   v_{{in}}={:.4f} cm/ns",
0258                                k_eta_labels[ieta], nEntries, m_drift_velocity)
0259                        .c_str(),
0260                    "C");
0261     leg->AddEntry(h_fit_proj, "Gaussian slice mean", "p");
0262     leg->AddEntry(f1d, std::format("slope={:.4f}  v_{{new}}={:.4f}#pm{:.4f} cm/ns  t_{{0}}={:.1f} ns", slope, dv_new, dv_err, t0_new).c_str(), "l");
0263     leg->Draw();
0264   }
0265
0266   m_hist3D->GetXaxis()->SetRange(0, 0);
0267
0268   canvas->SaveAs(m_plot_filename.c_str());
0269   std::cout << Name() << "::End - QA canvas saved to " << m_plot_filename << std::endl;
0270
0271   // write histograms, fit and results to a ROOT file
0272   if (!m_root_filename.empty())
0273   {
0274     std::unique_ptr<TFile> outfile(TFile::Open(m_root_filename.c_str(), "RECREATE"));
0275     if (outfile && !outfile->IsZombie())
0276     {
0277       outfile->cd();
0278       m_hist3D->Write();
0279       h_fit->Write("h_fit_silicon");
0280       fit2d->Write();
0281       canvas->Write();
0282       TParameter<double>("slope", slope).Write();
0283       TParameter<double>("drift_velocity_in", m_drift_velocity).Write();
0284       TParameter<double>("drift_velocity_new", dv_new).Write();
0285       TParameter<double>("drift_velocity_err", dv_err).Write();
0286       TParameter<double>("t0_new", t0_new).Write();
0287       outfile->Close();
0288       std::cout << Name() << "::End - histograms and fit results saved to " << m_root_filename << std::endl;
0289     }
0290     else
0291     {
0292       std::cerr << Name() << "::End - could not open " << m_root_filename << " for writing." << std::endl;
0293     }
0294   }
0295
0296   delete canvas;
0297   delete fit2d;
0298   delete h_fit;
0299
0300   return Fun4AllReturnCodes::EVENT_OK;
0301 }
0302
0303 //_____________________________________________________________________
0304 int SiliconDriftEvaluator::load_nodes(PHCompositeNode* topNode)
0305 {
0306   // track map
0307   m_track_map = findNode::getClass<SvtxTrackMap>(topNode, m_trackmapname);
0308
0309   // local container
0310   m_container = findNode::getClass<Container>(topNode, "SiliconDriftEvaluator::Container");
0311   assert(m_container);
0312
0313   return Fun4AllReturnCodes::EVENT_OK;
0314 }
0315
0316 //_____________________________________________________________________
0317 void SiliconDriftEvaluator::evaluate_tracks()
0318 {
0319   if (!(m_track_map && m_container && m_hist3D))
0320   {
0321     return;
0322   }
0323
0324   // clear array
0325   m_container->clearTracks();
0326
0327   for (const auto& [track_id, track] : *m_track_map)
0328   {
0329     // require valid beam-crossing
0330     const auto crossing = track->get_crossing();
0331     if (crossing == SHRT_MAX)
0332     {
0333       std::cout << "SiliconDriftEvaluator::evaluate_tracks - invalid crossing, track ignored." << std::endl;
0334       continue;
0335     }
0336
0337     // require both seeds
0338     const auto* si_seed = track->get_silicon_seed();
0339     const auto* tpc_seed = track->get_tpc_seed();
0340     if (!si_seed || !tpc_seed)
0341     {
0342       continue;
0343     }
0344
0345     // count clusters per subsystem
0346     unsigned int n_tpc = 0;
0347     unsigned int n_mvtx = 0;
0348     unsigned int n_intt = 0;
0349
0350     for (const auto* seed : {track->get_silicon_seed(), track->get_tpc_seed()})
0351     {
0352       if (!seed)
0353       {
0354         continue;
0355       }
0356       for (auto it = seed->begin_cluster_keys(); it != seed->end_cluster_keys(); ++it)
0357       {
0358         switch (TrkrDefs::getTrkrId(*it))
0359         {
0360         case TrkrDefs::tpcId:
0361           ++n_tpc;
0362           break;
0363         case TrkrDefs::mvtxId:
0364           ++n_mvtx;
0365           break;
0366         case TrkrDefs::inttId:
0367           ++n_intt;
0368           break;
0369         default:
0370           break;
0371         }
0372       }
0373     }
0374
0375     // apply selection cuts
0376     if (n_tpc < m_min_nclusters_tpc)
0377     {
0378       continue;
0379     }
0380     if (n_mvtx < m_min_nclusters_mvtx)
0381     {
0382       continue;
0383     }
0384     if (n_intt < m_min_nclusters_intt)
0385     {
0386       continue;
0387     }
0388
0389     const float eta = tpc_seed->get_eta();
0390     if (std::abs(eta) > m_max_eta)
0391     {
0392       continue;
0393     }
0394
0395     const float pt = get_pt(track->get_px(), track->get_py());
0396     if (pt < m_min_pt)
0397     {
0398       continue;
0399     }
0400
0401     // get seed z positions at POCA
0402     const auto si_pos = TrackSeedHelper::get_xyz(si_seed);
0403     const auto tpc_pos = TrackSeedHelper::get_xyz(tpc_seed);
0404
0405     const float z_si = si_pos.z();
0406     const float z_tpc = tpc_pos.z();
0407
0408     // dz = (z_tpc + sign(eta)*crossing*crossing_interval*dv) - z_si
0409     const double sign_eta = (eta >= 0) ? 1.0 : -1.0;
0410     const float z_tpc_corr = z_tpc + sign_eta * crossing * m_crossing_interval * m_drift_velocity;
0411     const float dz = z_tpc_corr - z_si;
0412
0413     // fill track struct
0414     TrackStruct track_struct;
0415     track_struct._nclusters_mvtx = n_mvtx;
0416     track_struct._nclusters_intt = n_intt;
0417     track_struct._nclusters_tpc = n_tpc;
0418     track_struct._pt = pt;
0419     track_struct._eta = eta;
0420     track_struct._phi = tpc_seed->get_phi();
0421     track_struct._z_tpc = z_tpc;
0422     track_struct._z_si = z_si;
0423     track_struct._crossing = crossing;
0424     track_struct._dz = dz;
0425
0426     // fill histogram
0427     // eta bin centre: 0.5 for eta<0, 1.5 for eta>=0
0428     const double eta_bin = (eta >= 0) ? 1.5 : 0.5;
0429     m_hist3D->Fill(eta_bin, z_si, dz);
0430
0431     m_container->addTrack(track_struct);
0432   }
0433 }