NewCode2024/KshortMomentum/KshortMomentum.C

0001 void Draw1D(TH1 *h1, TH1 *h2, const char *name, const char *title, const char *xtitle, const char *ytitle)
0002 {
0003     // normalize the histograms
0004     h1->Scale(1.0 / h1->Integral(-1, -1));
0005     h2->Scale(1.0 / h2->Integral(-1, -1));
0006     TCanvas *c = new TCanvas(name, title, 800, 700);
0007     h1->SetMarkerStyle(20);
0008     h1->SetMarkerSize(0.8);
0009     h1->SetLineColor(kBlack);
0010     h2->SetLineColor(kRed);
0011     h1->SetLineWidth(2);
0012     h2->SetLineWidth(2);
0013     h2->SetMarkerSize(0);
0014     h2->Draw("histe");
0015     h1->Draw("PE same");
0016     h2->GetXaxis()->SetTitle(xtitle);
0017     h2->GetYaxis()->SetTitle(ytitle);
0018     h2->GetYaxis()->SetTitleOffset(1.5);
0019     TLegend *leg = new TLegend(0.7, 0.7, 0.9, 0.9);
0020     leg->AddEntry(h1, "Truth", "lep");
0021     leg->AddEntry(h2, "Sampled", "le");
0022     leg->Draw();
0023     c->Draw();
0024     c->SaveAs(Form("%s.png", name));
0025     c->SaveAs(Form("%s.pdf", name));
0026 };
0027
0028 void KshortMomentum()
0029 {
0030     ROOT::EnableImplicitMT();
0031     ROOT::RDataFrame df("HFTrackEfficiency", "/sphenix/user/cdean/public/forHaoRen/outputHFEffFile_Kshort.root");
0032
0033     // a 3D histogram for true_mother_px vs true_mother_py vs true_mother_pz
0034     auto h_motherpx_motherpy_motherpz = df.Histo3D({"h_motherpx_motherpy_motherpz", ";p_{x} [GeV];p_{y} [GeV];p_{z} [GeV]", 200, -10, 10, 200, -10, 10, 200, -10, 10}, "true_mother_px", "true_mother_py", "true_mother_pz");
0035     auto h_motherP_motherPt = df.Histo2D({"h_motherP_motherPt", ";p [GeV];p_{T} [GeV]", 100, 0, 10, 100, 0, 10}, "true_mother_p", "true_mother_pT");
0036     auto h_motherP = df.Histo1D({"h_motherP", ";K_{S}^{0} p [GeV]", 100, 0, 10}, "true_mother_p");
0037     auto h_motherPt = df.Histo1D({"h_motherPt", ";K_{S}^{0} p_{T} [GeV]", 100, 0, 10}, "true_mother_pT");
0038     auto h_motherPz = df.Histo1D({"h_motherPz", ";K_{S}^{0} p_{z} [GeV]", 100, 0, 10}, "true_mother_pz");
0039     auto h_motherEta = df.Histo1D({"h_motherEta", ";K_{S}^{0} #eta", 40, -2, 2}, "true_mother_eta");
0040     auto h_motherPxPy = df.Histo2D({"h_motherPxPy", ";K_{S}^{0} p_{x} [GeV];K_{S}^{0} p_{y} [GeV]", 100, -10, 10, 100, -10, 10}, "true_mother_px", "true_mother_py");
0041
0042     // Randomly generate by sampling from the 2D histogram
0043     TH1F *h_sampled_P = new TH1F("h_sampled_P", ";K_{S}^{0} p [GeV]", 100, 0, 10);
0044     TH1F *h_sampled_pT = new TH1F("h_sampled_pT", ";K_{S}^{0} p_{T} [GeV]", 100, 0, 10);
0045     TH1F *h_sampled_pz = new TH1F("h_sampled_pz", ";K_{S}^{0} p_{z} [GeV]", 100, 0, 10);
0046     TH1F *h_sampled_Eta = new TH1F("h_sampled_Eta", ";K_{S}^{0} #eta", 40, -2, 2);
0047     TH1F *h_sampled_y = new TH1F("h_sampled_y", ";K_{S}^{0} Rapidity y", 40, -2, 2);
0048     TH2F *h_sampled_p_pt = new TH2F("h_sampled_p_pt", ";K_{S}^{0} p [GeV];K_{S}^{0} p_{T} [GeV]", 100, 0, 10, 100, 0, 10);
0049     TH2F *h_sampled_px_py = new TH2F("h_sampled_px_py", ";K_{S}^{0} p_{x} [GeV];K_{S}^{0} p_{y} [GeV]", 100, -10, 10, 100, -10, 10);
0050     int Nrandom = 1E5;
0051     double M_kshort = 0.497611; // GeV
0052     for (int i = 0; i < Nrandom; i++)
0053     {
0054         // GetRandom3 from the 3D histogram
0055         double px, py, pz, p, pT, phi, eta, y, E;
0056         h_motherpx_motherpy_motherpz->GetRandom3(px, py, pz);
0057         // Randomly generate phi between -pi to pi
0058         phi = gRandom->Uniform(-TMath::Pi(), TMath::Pi());
0059         p = std::sqrt(px * px + py * py + pz * pz);
0060         pT = std::sqrt(px * px + py * py);
0061         eta = std::atanh(pz / p);
0062         // energy
0063         E = std::sqrt(p*p + M_kshort*M_kshort);
0064         // rapidity
0065         y = 0.5 * std::log((E + pz) / (E - pz));
0066         // cout << "p: " << p << ", pT: " << pT << ", pz: " << pz << ", eta: " << eta << endl;
0067         TLorentzVector lvec_rand;
0068         lvec_rand.SetPtEtaPhiM(pT, eta, phi, M_kshort);
0069
0070         h_sampled_P->Fill(p);
0071         h_sampled_pT->Fill(pT);
0072         h_sampled_pz->Fill(pz);
0073         h_sampled_Eta->Fill(eta);
0074         h_sampled_y->Fill(y);
0075         h_sampled_px_py->Fill(lvec_rand.Px(), lvec_rand.Py());
0076         h_sampled_p_pt->Fill(p, pT);
0077     }
0078
0079     // Draw
0080     TCanvas *c = new TCanvas("c", "c", 800, 700);
0081     c->SetLogz();
0082     c->cd();
0083     gPad->SetRightMargin(0.13);
0084     h_motherP_motherPt->Draw("colz");
0085     c->Draw();
0086     c->SaveAs("Kshort_motherP_motherPt.png");
0087     c->SaveAs("Kshort_motherP_motherPt.pdf");
0088
0089     c->Clear();
0090     c->cd();
0091     h_motherPxPy->Draw("colz");
0092     c->Draw();
0093     c->SaveAs("Kshort_motherPxPy.png");
0094     c->SaveAs("Kshort_motherPxPy.pdf");
0095
0096     // 2D histograms for the sampled values
0097     c->Clear();
0098     c->cd();
0099     h_sampled_p_pt->Draw("colz");
0100     c->Draw();
0101     c->SaveAs("Kshort_TH3sampled_p_pt.png");
0102     c->SaveAs("Kshort_TH3sampled_p_pt.pdf");
0103
0104     c->Clear();
0105     c->cd();
0106     h_sampled_px_py->Draw("colz");
0107     c->Draw();
0108     c->SaveAs("Kshort_TH3sampled_pxpy.png");
0109     c->SaveAs("Kshort_TH3sampled_pxpy.pdf");
0110
0111     Draw1D(h_motherP.GetPtr(), h_sampled_P, "Kshort_motherP_TH3sampled", "Kshort Mother p", "p [GeV]", "Counts");
0112     Draw1D(h_motherPt.GetPtr(), h_sampled_pT, "Kshort_motherPt_TH3sampled", "Kshort Mother pT", "p_{T} [GeV]", "Counts");
0113     Draw1D(h_motherPz.GetPtr(), h_sampled_pz, "Kshort_motherPz_TH3sampled", "Kshort Mother pz", "p_{z} [GeV]", "Counts");
0114     Draw1D(h_motherEta.GetPtr(), h_sampled_Eta, "Kshort_motherEta_TH3sampled", "Kshort Mother eta", "#eta", "Counts");
0115 }