analysis_INTT/plot/compStrange.C

0001 float TickSize = 0.03;
0002 float AxisTitleSize = 0.05;
0003 float AxisLabelSize = 0.045;
0004 float LeftMargin = 0.15;
0005 float RightMargin = 0.05;
0006 float TopMargin = 0.08;
0007 float BottomMargin = 0.13;
0008 float textscale = 2.6;
0009
0010 std::string plotdir = "./countLepton";
0011
0012 bool checkfile(const std::string &filename)
0013 {
0014     TFile *f = TFile::Open(filename.c_str());
0015     if (!f || f->IsZombie() || f->GetSize() <= 0)
0016     {
0017         std::cout << "Skipping invalid or empty file: " << filename << std::endl;
0018         if (f)
0019             f->Close();
0020         return false;
0021     }
0022
0023     if (!f->GetListOfKeys()->Contains("minitree"))
0024     {
0025         std::cout << "Skipping file without minitree: " << filename << std::endl;
0026         f->Close();
0027         return false;
0028     }
0029
0030     // Check if PrimaryG4P_PID branch exists
0031     TTree *tree = (TTree*)f->Get("minitree");
0032     if (!tree || !tree->GetBranch("PrimaryG4P_PID"))
0033     {
0034         std::cout << "Skipping file without PrimaryG4P_PID branch: " << filename << std::endl;
0035         f->Close();
0036         return false;
0037     }
0038
0039     f->Close();
0040     return true;
0041 }
0042
0043 // return std::tuple<float, float, float> for the strange fraction of HIJING, EPOS, and AMPT
0044 std::tuple<float, float, float> getStrangeFrac(int fileidx = 0, bool makeplot = false)
0045 {
0046     // make sure all the three files exist and is valid, and contains the required tree
0047     std::string file_HIJING = "/sphenix/tg/tg01/hf/hjheng/ppg02/minitree/TrackletMinitree_Sim_HIJING_MDC2_ana467_20250225/dRcut0p5_NominalVtxZ_RandomClusSet0_clusAdcCutSet0_clusPhiSizeCutSet0/minitree_" + std::string(TString::Format("%05d", fileidx).Data()) + ".root";
0048     std::string file_EPOS = "/sphenix/tg/tg01/hf/hjheng/ppg02/minitree/TrackletMinitree_Sim_EPOS_MDC2_ana467_20250301/dRcut0p5_NominalVtxZ_RandomClusSet0_clusAdcCutSet0_clusPhiSizeCutSet0/minitree_" + std::string(TString::Format("%05d", fileidx).Data()) + ".root";
0049     std::string file_AMPT = "/sphenix/tg/tg01/hf/hjheng/ppg02/minitree/TrackletMinitree_Sim_AMPT_MDC2_ana467_20250301/dRcut0p5_NominalVtxZ_RandomClusSet0_clusAdcCutSet0_clusPhiSizeCutSet0/minitree_" + std::string(TString::Format("%05d", fileidx).Data()) + ".root";
0050
0051     bool valid_HIJING = checkfile(file_HIJING);
0052     bool valid_EPOS = checkfile(file_EPOS);
0053     bool valid_AMPT = checkfile(file_AMPT);
0054
0055     if (!valid_HIJING || !valid_EPOS || !valid_AMPT)
0056     {
0057         return std::make_tuple(std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN());
0058     }
0059
0060     ROOT::EnableImplicitMT();
0061     ROOT::RDataFrame df_HIJING("minitree", file_HIJING);
0062     ROOT::RDataFrame df_EPOS("minitree", file_EPOS);
0063     ROOT::RDataFrame df_AMPT("minitree", file_AMPT);
0064
0065     // Define a helper function to calculate NStrange and NPrimaryG4P
0066     auto defineNStrange = [](ROOT::RDF::RNode df)
0067     {
0068         return df.Define("NPrimaryG4P", [](const std::vector<int> &PrimaryG4P_PID) { return static_cast<int>(PrimaryG4P_PID.size()); }, {"PrimaryG4P_PID"})
0069             .Define("NStrange",
0070                     [](const std::vector<int> &PrimaryG4P_PID, const std::vector<float> &PrimaryG4P_eta)
0071                     {
0072                         int NStrange = 0;
0073                         for (size_t i = 0; i < PrimaryG4P_PID.size(); i++)
0074                         {
0075                             if ((PrimaryG4P_PID[i] == 310 || abs(PrimaryG4P_PID[i]) == 3122) && fabs(PrimaryG4P_eta[i]) < 1.5)
0076                             {
0077                                 NStrange++;
0078                             }
0079                         }
0080                         return NStrange;
0081                     },
0082                     {"PrimaryG4P_PID", "PrimaryG4P_eta"});
0083     };
0084
0085     // Apply the function to each dataframe
0086     auto df_HIJING_NStrange = defineNStrange(df_HIJING);
0087     auto df_EPOS_NStrange = defineNStrange(df_EPOS);
0088     auto df_AMPT_NStrange = defineNStrange(df_AMPT);
0089
0090     // Get the vector of NStrange and NPrimaryG4P
0091     auto v_NStrange_HIJING = df_HIJING_NStrange.Take<int>("NStrange");
0092     auto v_NStrange_EPOS = df_EPOS_NStrange.Take<int>("NStrange");
0093     auto v_NStrange_AMPT = df_AMPT_NStrange.Take<int>("NStrange");
0094     auto v_NPrimaryG4P_HIJING = df_HIJING_NStrange.Take<int>("NPrimaryG4P");
0095     auto v_NPrimaryG4P_EPOS = df_EPOS_NStrange.Take<int>("NPrimaryG4P");
0096     auto v_NPrimaryG4P_AMPT = df_AMPT_NStrange.Take<int>("NPrimaryG4P");
0097
0098     float avg_strange_HIJING = 0, avg_strange_EPOS = 0, avg_strange_AMPT = 0;
0099     for (size_t i = 0; i < v_NStrange_HIJING->size(); i++)
0100     {
0101         avg_strange_HIJING += (float)(*v_NStrange_HIJING)[i] / (*v_NPrimaryG4P_HIJING)[i] * 100.;
0102         avg_strange_EPOS += (float)(*v_NStrange_EPOS)[i] / (*v_NPrimaryG4P_EPOS)[i] * 100.;
0103         avg_strange_AMPT += (float)(*v_NStrange_AMPT)[i] / (*v_NPrimaryG4P_AMPT)[i] * 100.;
0104     }
0105     avg_strange_HIJING /= v_NStrange_HIJING->size();
0106     avg_strange_EPOS /= v_NStrange_EPOS->size();
0107     avg_strange_AMPT /= v_NStrange_AMPT->size();
0108
0109     if (makeplot)
0110     {
0111         // Create 3 histograms, one for the number of strange particles, one for the number of primary G4P, and one for the ratio of the two. All these three are as a function of the event index
0112         auto h_NStrange_HIJING = new TH1F("h_NStrange_HIJING", "h_NStrange_HIJING;Event index;NStrange", v_NStrange_HIJING->size(), 0, v_NStrange_HIJING->size());
0113         auto h_NStrange_EPOS = new TH1F("h_NStrange_EPOS", "h_NStrange_EPOS;Event index;NStrange", v_NStrange_EPOS->size(), 0, v_NStrange_EPOS->size());
0114         auto h_NStrange_AMPT = new TH1F("h_NStrange_AMPT", "h_NStrange_AMPT;Event index;NStrange", v_NStrange_AMPT->size(), 0, v_NStrange_AMPT->size());
0115         auto h_NPrimaryG4P_HIJING = new TH1F("h_NPrimaryG4P_HIJING", "h_NPrimaryG4P_HIJING;Event index;NPrimaryG4P", v_NPrimaryG4P_HIJING->size(), 0, v_NPrimaryG4P_HIJING->size());
0116         auto h_NPrimaryG4P_EPOS = new TH1F("h_NPrimaryG4P_EPOS", "h_NPrimaryG4P_EPOS;Event index;NPrimaryG4P", v_NPrimaryG4P_EPOS->size(), 0, v_NPrimaryG4P_EPOS->size());
0117         auto h_NPrimaryG4P_AMPT = new TH1F("h_NPrimaryG4P_AMPT", "h_NPrimaryG4P_AMPT;Event index;NPrimaryG4P", v_NPrimaryG4P_AMPT->size(), 0, v_NPrimaryG4P_AMPT->size());
0118         auto h_NStrange_ratio_HIJING = new TH1F("h_NStrange_ratio_HIJING", "h_NStrange_ratio_HIJING;Event index;NStrange/NPrimaryG4P", v_NStrange_HIJING->size(), 0, v_NStrange_HIJING->size());
0119         auto h_NStrange_ratio_EPOS = new TH1F("h_NStrange_ratio_EPOS", "h_NStrange_ratio_EPOS;Event index;NStrange/NPrimaryG4P", v_NStrange_EPOS->size(), 0, v_NStrange_EPOS->size());
0120         auto h_NStrange_ratio_AMPT = new TH1F("h_NStrange_ratio_AMPT", "h_NStrange_ratio_AMPT;Event index;NStrange/NPrimaryG4P", v_NStrange_AMPT->size(), 0, v_NStrange_AMPT->size());
0121
0122         float avg_strange_HIJING = 0, avg_strange_EPOS = 0, avg_strange_AMPT = 0;
0123         for (size_t i = 0; i < v_NStrange_HIJING->size(); i++)
0124         {
0125             h_NStrange_HIJING->SetBinContent(i + 1, (*v_NStrange_HIJING)[i]);
0126             h_NStrange_EPOS->SetBinContent(i + 1, (*v_NStrange_EPOS)[i]);
0127             h_NStrange_AMPT->SetBinContent(i + 1, (*v_NStrange_AMPT)[i]);
0128             h_NPrimaryG4P_HIJING->SetBinContent(i + 1, (*v_NPrimaryG4P_HIJING)[i]);
0129             h_NPrimaryG4P_EPOS->SetBinContent(i + 1, (*v_NPrimaryG4P_EPOS)[i]);
0130             h_NPrimaryG4P_AMPT->SetBinContent(i + 1, (*v_NPrimaryG4P_AMPT)[i]);
0131             h_NStrange_ratio_HIJING->SetBinContent(i + 1, (float)(*v_NStrange_HIJING)[i] / (*v_NPrimaryG4P_HIJING)[i] * 100.);
0132             h_NStrange_ratio_EPOS->SetBinContent(i + 1, (float)(*v_NStrange_EPOS)[i] / (*v_NPrimaryG4P_EPOS)[i] * 100.);
0133             h_NStrange_ratio_AMPT->SetBinContent(i + 1, (float)(*v_NStrange_AMPT)[i] / (*v_NPrimaryG4P_AMPT)[i] * 100.);
0134         }
0135
0136         // directory to save the output plots; if the directory does not exist, create it using
0137         system(Form("mkdir -p %s", plotdir.c_str()));
0138
0139         float pad1_ymax = std::max({h_NPrimaryG4P_HIJING->GetMaximum(), h_NPrimaryG4P_EPOS->GetMaximum(), h_NPrimaryG4P_AMPT->GetMaximum()}) * 1.2;
0140         float pad2_ymax = std::max({h_NStrange_ratio_HIJING->GetMaximum(), h_NStrange_ratio_EPOS->GetMaximum(), h_NStrange_ratio_AMPT->GetMaximum()}) * 1.2;
0141         float pad1_ymin = std::min({h_NStrange_HIJING->GetMinimum(0) * 0.3, h_NStrange_EPOS->GetMinimum(0) * 0.3, h_NStrange_AMPT->GetMinimum(0) * 0.3});
0142         float pad2_ymin = 0;
0143         TCanvas *c = new TCanvas("c", "c", 1000, 700);
0144         TPad *pad1 = new TPad("pad1", "pad1", 0, 0.3, 1, 1);
0145         pad1->SetBottomMargin(0);
0146         pad1->SetTopMargin(0.33);
0147         pad1->SetRightMargin(RightMargin);
0148         pad1->Draw();
0149         TPad *pad2 = new TPad("pad2", "pad2", 0, 0, 1, 0.3);
0150         pad2->SetRightMargin(RightMargin);
0151         pad2->SetTopMargin(0);
0152         pad2->SetBottomMargin(0.33);
0153         pad2->Draw();
0154         pad1->cd();
0155         pad1->SetLogy(1);
0156         h_NPrimaryG4P_HIJING->SetLineColor(30);
0157         h_NPrimaryG4P_HIJING->SetLineWidth(2);
0158         h_NPrimaryG4P_HIJING->GetYaxis()->SetTitle("Counts");
0159         h_NPrimaryG4P_HIJING->GetYaxis()->SetTitleOffset(1.4);
0160         h_NPrimaryG4P_HIJING->GetYaxis()->SetRangeUser(pad1_ymin, pad1_ymax);
0161         h_NPrimaryG4P_HIJING->GetYaxis()->SetTitleSize(AxisTitleSize);
0162         h_NPrimaryG4P_HIJING->GetYaxis()->SetLabelSize(AxisLabelSize);
0163         h_NPrimaryG4P_HIJING->GetXaxis()->SetTitleSize(0);
0164         h_NPrimaryG4P_HIJING->GetXaxis()->SetLabelSize(0);
0165         h_NPrimaryG4P_HIJING->Draw("hist");
0166         h_NPrimaryG4P_EPOS->SetLineColor(38);
0167         h_NPrimaryG4P_EPOS->SetLineWidth(2);
0168         h_NPrimaryG4P_EPOS->Draw("hist same");
0169         h_NPrimaryG4P_AMPT->SetLineColor(46);
0170         h_NPrimaryG4P_AMPT->SetLineWidth(2);
0171         h_NPrimaryG4P_AMPT->Draw("hist same");
0172         h_NStrange_HIJING->SetLineColor(30);
0173         h_NStrange_HIJING->SetLineWidth(2);
0174         h_NStrange_HIJING->SetLineStyle(2);
0175         h_NStrange_HIJING->Draw("hist same");
0176         h_NStrange_EPOS->SetLineColor(38);
0177         h_NStrange_EPOS->SetLineWidth(2);
0178         h_NStrange_EPOS->SetLineStyle(2);
0179         h_NStrange_EPOS->Draw("hist same");
0180         h_NStrange_AMPT->SetLineColor(46);
0181         h_NStrange_AMPT->SetLineWidth(2);
0182         h_NStrange_AMPT->SetLineStyle(2);
0183         h_NStrange_AMPT->Draw("hist same");
0184         TLegend *leg = new TLegend(pad1->GetLeftMargin(),           //
0185                                    1 - pad1->GetTopMargin() + 0.01, //
0186                                    pad1->GetLeftMargin() + 0.25,    //
0187                                    0.98);
0188         leg->AddEntry(h_NPrimaryG4P_HIJING, "HIJING", "l");
0189         leg->AddEntry(h_NPrimaryG4P_EPOS, "EPOS", "l");
0190         leg->AddEntry(h_NPrimaryG4P_AMPT, "AMPT", "l");
0191         leg->AddEntry((TObject *)0, "Solid - Total number of primary particles", "");
0192         leg->AddEntry((TObject *)0, "Dash - Number of K_{s}^{0} and #Lambda", "");
0193         leg->SetBorderSize(0);
0194         leg->Draw();
0195         pad2->cd();
0196         h_NStrange_ratio_HIJING->SetLineColor(30);
0197         h_NStrange_ratio_HIJING->SetLineWidth(1);
0198         h_NStrange_ratio_HIJING->GetYaxis()->SetTitle("Fraction (%)");
0199         h_NStrange_ratio_HIJING->GetYaxis()->SetTitleOffset(0.5);
0200         h_NStrange_ratio_HIJING->GetYaxis()->SetRangeUser(pad2_ymin, pad2_ymax);
0201         h_NStrange_ratio_HIJING->GetYaxis()->SetNdivisions(505);
0202         h_NStrange_ratio_HIJING->GetYaxis()->SetTitleSize(AxisTitleSize * textscale);
0203         h_NStrange_ratio_HIJING->GetYaxis()->SetLabelSize(AxisLabelSize * textscale);
0204         h_NStrange_ratio_HIJING->GetXaxis()->SetTitleSize(AxisTitleSize * textscale);
0205         h_NStrange_ratio_HIJING->GetXaxis()->SetLabelSize(AxisLabelSize * textscale);
0206         h_NStrange_ratio_HIJING->GetXaxis()->SetTitle("Event index");
0207         h_NStrange_ratio_HIJING->GetXaxis()->SetTitleOffset(1.2);
0208         h_NStrange_ratio_HIJING->Draw("l");
0209         h_NStrange_ratio_EPOS->SetLineColor(38);
0210         h_NStrange_ratio_EPOS->SetLineWidth(1);
0211         h_NStrange_ratio_EPOS->Draw("l same");
0212         h_NStrange_ratio_AMPT->SetLineColor(46);
0213         h_NStrange_ratio_AMPT->SetLineWidth(1);
0214         h_NStrange_ratio_AMPT->Draw("l same");
0215         TLine *l_strange_HIJING = new TLine(0, avg_strange_HIJING, v_NStrange_HIJING->size(), avg_strange_HIJING);
0216         l_strange_HIJING->SetLineColor(30);
0217         l_strange_HIJING->SetLineStyle(2);
0218         l_strange_HIJING->Draw();
0219         TLine *l_strange_EPOS = new TLine(0, avg_strange_EPOS, v_NStrange_EPOS->size(), avg_strange_EPOS);
0220         l_strange_EPOS->SetLineColor(38);
0221         l_strange_EPOS->SetLineStyle(2);
0222         l_strange_EPOS->Draw();
0223         TLine *l_strange_AMPT = new TLine(0, avg_strange_AMPT, v_NStrange_AMPT->size(), avg_strange_AMPT);
0224         l_strange_AMPT->SetLineColor(46);
0225         l_strange_AMPT->SetLineStyle(2);
0226         l_strange_AMPT->Draw();
0227         TLatex *l_avg = new TLatex();
0228         l_avg->SetTextAlign(kHAlignLeft + kVAlignTop);
0229         l_avg->SetTextSize(AxisLabelSize * textscale * 0.8);
0230         l_avg->DrawLatex(0.1, pad2_ymax - 0.02, Form("HIJING: %.3g%%, EPOS: %.3g%%, AMPT: %.3g%%", avg_strange_HIJING, avg_strange_EPOS, avg_strange_AMPT));
0231         c->RedrawAxis();
0232         c->SaveAs(Form("%s/compStrange.pdf", plotdir.c_str()));
0233         c->SaveAs(Form("%s/compStrange.png", plotdir.c_str()));
0234     }
0235
0236     return std::make_tuple(avg_strange_HIJING, avg_strange_EPOS, avg_strange_AMPT);
0237 }
0238
0239 void compStrange()
0240 {
0241     system(Form("mkdir -p %s", plotdir.c_str()));
0242
0243     std::vector<float> avg_strange_HIJING, avg_strange_EPOS, avg_strange_AMPT;
0244     float minY = std::numeric_limits<float>::max();
0245     float maxY = std::numeric_limits<float>::min();
0246
0247     for (int i = 0; i < 5000; i++)
0248     {
0249         if (avg_strange_HIJING.size() >= 1000)
0250         {
0251             break;
0252         }
0253
0254         if (i == 0)
0255         {
0256             std::tuple<float, float, float> avg_strange = getStrangeFrac(i, true);
0257
0258             // if either of the elements is NaN, skip this event
0259             if (std::isnan(std::get<0>(avg_strange)) || std::isnan(std::get<1>(avg_strange)) || std::isnan(std::get<2>(avg_strange)))
0260             {
0261                 continue;
0262             }
0263             avg_strange_HIJING.push_back(std::get<0>(avg_strange));
0264             avg_strange_EPOS.push_back(std::get<1>(avg_strange));
0265             avg_strange_AMPT.push_back(std::get<2>(avg_strange));
0266         }
0267         else
0268         {
0269             std::tuple<float, float, float> avg_strange = getStrangeFrac(i, false);
0270             if (std::isnan(std::get<0>(avg_strange)) || std::isnan(std::get<1>(avg_strange)) || std::isnan(std::get<2>(avg_strange)))
0271             {
0272                 continue;
0273             }
0274             avg_strange_HIJING.push_back(std::get<0>(avg_strange));
0275             avg_strange_EPOS.push_back(std::get<1>(avg_strange));
0276             avg_strange_AMPT.push_back(std::get<2>(avg_strange));
0277         }
0278
0279         minY = std::min({minY, avg_strange_HIJING.back(), avg_strange_EPOS.back(), avg_strange_AMPT.back()});
0280         maxY = std::max({maxY, avg_strange_HIJING.back(), avg_strange_EPOS.back(), avg_strange_AMPT.back()});
0281
0282         std::cout << "File idx: " << i << ", HIJING: " << avg_strange_HIJING.back() << ", EPOS: " << avg_strange_EPOS.back() << ", AMPT: " << avg_strange_AMPT.back() << std::endl;
0283     }
0284
0285     // create TGraphs for the average strange fraction
0286     std::vector<float> x_vals(avg_strange_HIJING.size());
0287     for (size_t i = 0; i < x_vals.size(); ++i)
0288     {
0289         x_vals[i] = i;
0290     }
0291     TGraph *g_avg_strange_HIJING = new TGraph(avg_strange_HIJING.size(), x_vals.data(), avg_strange_HIJING.data());
0292     TGraph *g_avg_strange_EPOS = new TGraph(avg_strange_EPOS.size(), x_vals.data(), avg_strange_EPOS.data());
0293     TGraph *g_avg_strange_AMPT = new TGraph(avg_strange_AMPT.size(), x_vals.data(), avg_strange_AMPT.data());
0294     TCanvas *c = new TCanvas("c", "c", 800, 600);
0295     gPad->SetTopMargin(0.17);
0296     c->cd();
0297     g_avg_strange_HIJING->SetLineColor(30);
0298     g_avg_strange_HIJING->SetLineWidth(2);
0299     g_avg_strange_HIJING->GetYaxis()->SetTitle("Strange fraction (%)");
0300     g_avg_strange_HIJING->GetYaxis()->SetTitleOffset(1.4);
0301     g_avg_strange_HIJING->GetYaxis()->SetTitleSize(AxisTitleSize);
0302     g_avg_strange_HIJING->GetYaxis()->SetLabelSize(AxisLabelSize);
0303     g_avg_strange_HIJING->GetXaxis()->SetTitle("File index");
0304     g_avg_strange_HIJING->GetXaxis()->SetTitleSize(AxisTitleSize);
0305     g_avg_strange_HIJING->GetXaxis()->SetLabelSize(AxisLabelSize);
0306     g_avg_strange_HIJING->GetYaxis()->SetRangeUser(minY * 0.9, maxY * 1.1);
0307     g_avg_strange_HIJING->Draw("al");
0308     g_avg_strange_EPOS->SetLineColor(38);
0309     g_avg_strange_EPOS->SetLineWidth(2);
0310     g_avg_strange_EPOS->Draw("l same");
0311     g_avg_strange_AMPT->SetLineColor(46);
0312     g_avg_strange_AMPT->SetLineWidth(2);
0313     g_avg_strange_AMPT->Draw("l same");
0314     TLegend *leg = new TLegend(gPad->GetLeftMargin(), 1 - gPad->GetTopMargin() + 0.01, gPad->GetLeftMargin() + 0.25, 0.98);
0315     leg->AddEntry(g_avg_strange_HIJING, Form("HIJING (avg. strange fraction = %.3g)",g_avg_strange_HIJING->GetMean(2)), "l");
0316     leg->AddEntry(g_avg_strange_EPOS, Form("EPOS (avg. strange fraction = %.3g)",g_avg_strange_EPOS->GetMean(2)), "l");
0317     leg->AddEntry(g_avg_strange_AMPT, Form("AMPT (avg. strange fraction = %.3g)",g_avg_strange_AMPT->GetMean(2)), "l");
0318     leg->SetTextSize(0.04);
0319     leg->SetBorderSize(0);
0320     leg->SetFillStyle(0);
0321     leg->Draw();
0322     c->RedrawAxis();
0323     c->SaveAs(Form("%s/compStrange_avgfileidx.pdf", plotdir.c_str()));
0324     c->SaveAs(Form("%s/compStrange_avgfileidx.png", plotdir.c_str()));
0325 }