InttSeedTrackPerformance-Jingyu/macro/GetPtFunc_wTakashi.C

0001 #include <TFile.h>
0002 #include <TTree.h>
0003 #include <TH3D.h>
0004 #include <TCanvas.h>
0005 #include <TVector2.h>
0006 #include <vector>
0007 #include <cmath>
0008
0009 void GetPtFunc_wTakashi(
0010     // const char* filename="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/dataset/merged_last500kNew.root",
0011     const char* filename="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/dataset/positron_last500kNew.root",
0012     // const char* filename="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/ParticleGen/output/electron_1M.root",
0013
0014     // const char* funcfile="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_rough_takashi_e.root",
0015     const char* funcfile="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_rough_takashi_p.root",
0016
0017     const char* treename="tree",
0018
0019     // const char* outputname="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_rough_takashi_p.root"
0020     const char* outputname="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_fitCofEta_takashi_p.root"
0021 )
0022 {
0023     TFile outfile(outputname, "RECREATE");
0024
0025     TFile *file_in = TFile::Open(filename);
0026     if (!file_in || file_in->IsZombie()) {
0027         std::cerr<<"Error: cannot open "<<filename<<std::endl;
0028         return;
0029     }
0030
0031     TTree *tree_data = (TTree*)file_in->Get(treename);
0032     if (!tree_data) {
0033         std::cerr<<"Error: cannot find tree "<<treename<<std::endl;
0034         return;
0035     }
0036
0037     TFile* funcFile = TFile::Open(funcfile, "READ");
0038     if (!funcFile || funcFile->IsZombie()) {
0039         std::cerr << "Error: cannot open function file " << funcfile << std::endl;
0040         return;
0041     }
0042     TF1* rough_funcTP = nullptr;
0043     TF1* rough_funcTG = nullptr;
0044     funcFile->GetObject("poly_func_profile", rough_funcTP);
0045     funcFile->GetObject("poly_func_graph", rough_funcTG);
0046
0047     int NClus;
0048     std::vector <int> *clus_system = nullptr;
0049     std::vector <int> *clus_layer = nullptr;
0050     std::vector <int> *clus_adc = nullptr;
0051     std::vector <double> *clus_X = nullptr;
0052     std::vector <double> *clus_Y = nullptr;
0053     std::vector <double> *clus_Z = nullptr;
0054
0055     std::vector<double> *CEMC_Pr_Hit_x = nullptr;
0056     std::vector<double> *CEMC_Pr_Hit_y = nullptr;
0057     std::vector<double> *CEMC_Pr_Hit_z = nullptr;
0058     std::vector<double> *CEMC_Pr_Hit_R = nullptr;
0059
0060     std::vector<double> *PrimaryG4P_Pt_ = nullptr;
0061     std::vector<double> *PrimaryG4P_Eta_ = nullptr;
0062
0063     tree_data->SetBranchAddress("trk_NClus", & NClus);
0064     tree_data->SetBranchAddress("trk_system", & clus_system);
0065     tree_data->SetBranchAddress("trk_layer", & clus_layer);
0066     tree_data->SetBranchAddress("trk_X", & clus_X);
0067     tree_data->SetBranchAddress("trk_Y", & clus_Y);
0068     tree_data->SetBranchAddress("trk_Z", & clus_Z);
0069
0070     tree_data->SetBranchAddress("CEMC_Pr_Hit_x",  &CEMC_Pr_Hit_x);
0071     tree_data->SetBranchAddress("CEMC_Pr_Hit_y",  &CEMC_Pr_Hit_y);
0072     tree_data->SetBranchAddress("CEMC_Pr_Hit_z",  &CEMC_Pr_Hit_z);
0073     tree_data->SetBranchAddress("CEMC_Pr_Hit_R",  &CEMC_Pr_Hit_R);
0074
0075     tree_data->SetBranchAddress("PrimaryG4P_Pt", &PrimaryG4P_Pt_);
0076     tree_data->SetBranchAddress("PrimaryG4P_Eta", &PrimaryG4P_Eta_);
0077
0078     TProfile* profC = new TProfile("C_vs_eta","C(eta);#eta;C(eta) [GeV·rad]", 120, -1.2, 1.2);
0079
0080     TH2D* h2etac = new TH2D("h2etac","C(eta);#eta;C(eta) [GeV·rad]", 48, -1.2, 1.2, 200, 0, 0.4);
0081
0082     std::vector<double> eta_vals, C_vals;
0083
0084     Long64_t nentries = tree_data->GetEntries();
0085     for (Long64_t i=0; i<nentries; ++i)
0086     {
0087         tree_data->GetEntry(i);
0088
0089         // 选事件：层4/5 恰有1个、层6/7 恰有1个
0090         int idx34 = -1, idx56 = -1;
0091         int cnt34 = 0, cnt56 = 0;
0092         for (int j=0; j<(int)clus_layer->size(); ++j) {
0093             int L = (*clus_layer)[j];
0094             if (L==3 || L==4) { ++cnt34; idx34=j; }
0095             if (L==5 || L==6) { ++cnt56; idx56=j; }
0096         }
0097         if (cnt34!=1 || cnt56!=1) continue;
0098
0099         // 要求 CEMC 记录为1
0100         if (CEMC_Pr_Hit_x->size()!=1) continue;
0101
0102         // 要求 prim particle 为1
0103         if (PrimaryG4P_Pt_->size()!=1) continue;
0104
0105         // 取出所有量
0106         double trk34_X = (*clus_X)[idx34];
0107         double trk34_Y = (*clus_Y)[idx34];
0108         double trk56_X = (*clus_X)[idx56];
0109         double trk56_Y = (*clus_Y)[idx56];
0110         double cemc_Pr_X  = (*CEMC_Pr_Hit_x)[0];
0111         double cemc_Pr_Y  = (*CEMC_Pr_Hit_y)[0];
0112         double Pr_pt = (*PrimaryG4P_Pt_)[0];
0113         double Pr_eta = (*PrimaryG4P_Eta_)[0];
0114
0115         double cemc_Pr_R = std::sqrt(cemc_Pr_X * cemc_Pr_X + cemc_Pr_Y * cemc_Pr_Y);
0116         // if (cemc_Pr_R > 102) continue; // 限制 CEMC 半径范围
0117
0118         // 计算 Δφ
0119         double phi1 = std::atan2(trk56_Y - trk34_Y, trk56_X - trk34_X);
0120         double phi2 = std::atan2(cemc_Pr_Y - trk56_Y, cemc_Pr_X - trk56_X);
0121         double dphi = TVector2::Phi_mpi_pi(phi2 - phi1);
0122
0123         // calculate C(η) = 〈pT·Δφ〉
0124         dphi = abs(dphi); // 确保 Δφ 为正值
0125         // double Cval = Pr_pt * dphi;
0126         double Cval = Pr_pt * std::pow(dphi, 0.986);
0127
0128         double C_rough_tp = rough_funcTP->Eval(Pr_eta);
0129         double C_rough_tg = rough_funcTG->Eval(Pr_eta);
0130         const double Cval_max_tp = C_rough_tp + 0.01;
0131         const double Cval_min_tp = C_rough_tp - 0.01;
0132         const double Cval_max_tg = C_rough_tg + 0.01;
0133         const double Cval_min_tg = C_rough_tg - 0.01;
0134
0135         if (Cval < Cval_max_tp && Cval > Cval_min_tp) {
0136             profC->Fill(Pr_eta, Cval);
0137             h2etac->Fill(Pr_eta, Cval);
0138         }
0139         if (Cval < Cval_max_tg && Cval > Cval_min_tg) {
0140             eta_vals.push_back(Pr_eta);
0141             C_vals.push_back(Cval);
0142         }
0143
0144         // rough fit function
0145         // const double Cval_max = 0.26;
0146         // const double Cval_min = 0.17;
0147
0148         // if (Cval > Cval_max || Cval < Cval_min) {
0149         //     continue;
0150         // }
0151
0152         // profC->Fill(Pr_eta, Cval);
0153         // h2etac->Fill(Pr_eta, Cval);
0154
0155         // eta_vals.push_back(Pr_eta);
0156         // C_vals.push_back(Cval);
0157     }
0158
0159     // 用 TGraph 做无损散点拟合
0160     int N = eta_vals.size();
0161     TGraph* gr = new TGraph(N, eta_vals.data(), C_vals.data());
0162     gr->SetTitle("C(eta);#eta;C=pt*#Delta#phi");
0163
0164     // Fit TGraph
0165     // TF1* poly_func_graph = new TF1("poly_func_graph","([0]+[1]*x+[2]*x*x+[3]*x*x*x+[4]*x*x*x*x)", -1.1, 1.1);
0166     TF1* poly_func_graph = new TF1("poly_func_graph","([0]+[1]*x*x+[2]*x*x*x*x)", -1.1, 1.1);
0167     gr->Fit(poly_func_graph,"");
0168
0169     // Fit TProfile
0170     const double thr = 0.205;
0171     for (int ib = 1; ib <= profC->GetNbinsX(); ++ib) {
0172         if (profC->GetBinContent(ib) > thr) {
0173             profC->SetBinError(ib, 2e-1);
0174         }
0175     }
0176     // TF1* poly_func_profile = new TF1("poly_func_profile","([0]+[1]*x+[2]*x*x+[3]*x*x*x+[4]*x*x*x*x)", -1.1, 1.1);
0177     TF1* poly_func_profile = new TF1("poly_func_profile","([0]+[1]*x*x+[2]*x*x*x*x)", -1.1, 1.1);
0178     profC->Fit(poly_func_profile, "R");
0179
0180     // 画图看效果
0181     TCanvas* c1 = new TCanvas;
0182     gr->Draw("AP");       // A=画轴，P=画点
0183     poly_func_graph->Draw("same");  // 叠加画拟合曲线
0184     c1->SaveAs("/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/fitCofEta.pdf");
0185
0186     TCanvas* c2 = new TCanvas;
0187     poly_func_profile->SetLineColor(kRed);
0188     poly_func_profile->SetLineStyle(1);  // 1=实线
0189
0190     poly_func_graph->SetLineColor(kBlue);
0191     poly_func_graph->SetLineStyle(2);    // 2=虚线
0192
0193     // X 轴：η，范围 [-1.2, 1.2]
0194     profC->GetXaxis()->SetTitle("#eta");
0195     profC->GetXaxis()->SetLimits(-1.2, 1.2);
0196     // Y 轴：C = p_{T} * Δφ，范围 [0.196, 0.206]
0197     profC->GetYaxis()->SetTitle("C = p_{T} #times #Delta#phi");
0198     profC->GetYaxis()->SetRangeUser(0.196, 0.206);
0199
0200     profC->Draw("AP");       // A=画轴，P=画点
0201     poly_func_profile->Draw("same");  // 叠加画拟合曲线
0202     poly_func_graph->Draw("same"); // 叠加画拟合曲线
0203     c2->SaveAs("/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/TPRfitfuncS.pdf");
0204
0205     outfile.cd();
0206     profC->Write();
0207     h2etac->Write();
0208     // gr->Write();
0209     poly_func_profile->Write();
0210     poly_func_graph->Write();
0211
0212     file_in->Close();
0213 }
0214
0215