InttSeedTrackPerformance-Jingyu/macro/GetPtFunc.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(
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_sym_positron.root",
0015     // const char* funcfile="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_rough_sym.root",
0016
0017     const char* treename="tree",
0018
0019     // const char* outputname="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_rough_sym_positron.root"
0020     const char* outputname="/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/Ptfunc_fitCofEta_positron.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
0127         double C_rough_tp = rough_funcTP->Eval(Pr_eta);
0128         double C_rough_tg = rough_funcTG->Eval(Pr_eta);
0129         const double Cval_max_tp = C_rough_tp + 0.01;
0130         const double Cval_min_tp = C_rough_tp - 0.01;
0131         const double Cval_max_tg = C_rough_tg + 0.01;
0132         const double Cval_min_tg = C_rough_tg - 0.01;
0133
0134         if (Cval < Cval_max_tp && Cval > Cval_min_tp) {
0135             profC->Fill(Pr_eta, Cval);
0136             h2etac->Fill(Pr_eta, Cval);
0137         }
0138         if (Cval < Cval_max_tg && Cval > Cval_min_tg) {
0139             eta_vals.push_back(Pr_eta);
0140             C_vals.push_back(Pr_pt * dphi);
0141         }
0142
0143         // rough fit function
0144         // const double Cval_max = 0.3;
0145         // const double Cval_min = 0.15;
0146
0147         // if (Cval > Cval_max || Cval < Cval_min) {
0148         //     continue;
0149         // }
0150
0151         // profC->Fill(Pr_eta, Cval);
0152         // h2etac->Fill(Pr_eta, Cval);
0153
0154         // eta_vals.push_back(Pr_eta);
0155         // C_vals.push_back(Pr_pt * dphi);
0156     }
0157
0158     // 用 TGraph 做无损散点拟合
0159     int N = eta_vals.size();
0160     TGraph* gr = new TGraph(N, eta_vals.data(), C_vals.data());
0161     gr->SetTitle("C(eta);#eta;C=pt*#Delta#phi");
0162
0163     // Fit TGraph
0164     // 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);
0165     TF1* poly_func_graph = new TF1("poly_func_graph","([0]+[1]*x*x+[2]*x*x*x*x)", -1.1, 1.1);
0166     gr->Fit(poly_func_graph,"");
0167
0168     // Fit TProfile
0169     const double thr = 0.205;
0170     for (int ib = 1; ib <= profC->GetNbinsX(); ++ib) {
0171         if (profC->GetBinContent(ib) > thr) {
0172             profC->SetBinError(ib, 2e-1);
0173         }
0174     }
0175     // 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);
0176     TF1* poly_func_profile = new TF1("poly_func_profile","([0]+[1]*x*x+[2]*x*x*x*x)", -1.1, 1.1);
0177     profC->Fit(poly_func_profile, "R");
0178
0179     // 画图看效果
0180     TCanvas* c1 = new TCanvas;
0181     gr->Draw("AP");       // A=画轴，P=画点
0182     poly_func_graph->Draw("same");  // 叠加画拟合曲线
0183     c1->SaveAs("/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/fitCofEta.pdf");
0184
0185     TCanvas* c2 = new TCanvas;
0186     poly_func_profile->SetLineColor(kRed);
0187     poly_func_profile->SetLineStyle(1);  // 1=实线
0188
0189     poly_func_graph->SetLineColor(kBlue);
0190     poly_func_graph->SetLineStyle(2);    // 2=虚线
0191
0192     // X 轴：η，范围 [-1.2, 1.2]
0193     profC->GetXaxis()->SetTitle("#eta");
0194     profC->GetXaxis()->SetLimits(-1.2, 1.2);
0195     // Y 轴：C = p_{T} * Δφ，范围 [0.196, 0.206]
0196     profC->GetYaxis()->SetTitle("C = p_{T} #times #Delta#phi");
0197     profC->GetYaxis()->SetRangeUser(0.196, 0.206);
0198
0199     profC->Draw("AP");       // A=画轴，P=画点
0200     poly_func_profile->Draw("same");  // 叠加画拟合曲线
0201     poly_func_graph->Draw("same"); // 叠加画拟合曲线
0202     c2->SaveAs("/mnt/e/sphenix/INTT-EMCAL/InttSeedingTrackDev/InttSeedTrackPerformance/output/TPRfitfuncS.pdf");
0203
0204     outfile.cd();
0205     profC->Write();
0206     h2etac->Write();
0207     // gr->Write();
0208     poly_func_profile->Write();
0209     poly_func_graph->Write();
0210
0211     file_in->Close();
0212 }
0213
0214