Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2026-07-16 08:11:26

0001 #include <any>
0002 #include <cmath>
0003 
0004 using namespace std;
0005 
0006 bool saveFinalPlots = true;
0007 bool printMarkdownTables = true;
0008 bool printLatexTables = true;
0009 bool printArrays = true;
0010 
0011 bool global_isSim = true;
0012 
0013 //Custom binning scheme for LF analysis
0014 float lower_bin_bounds[] = {0.5, 0.8, 1.1, 1.4, 1.8, 2.2, 3, 4};
0015 const unsigned int n_variable_bins = sizeof(lower_bin_bounds)/sizeof(lower_bin_bounds[0]) - 1; 
0016 
0017 template <typename T>
0018 string to_string_with_precision(const T a_value, const int n = 0)
0019 {
0020     ostringstream out;
0021     out.precision(n);
0022     out << fixed << a_value;
0023     return out.str();
0024 }
0025 
0026 //Initialise histograms
0027 TH1F makeHisto(int nBins, float min, float max, string type, string xAxisTitle, string unit, int precision, string yAxisTitle = "p_{T} Accept.")
0028 {
0029   string histo_name = type + "_" + xAxisTitle;
0030   TH1F myHisto(histo_name.c_str(), histo_name.c_str(), nBins, min, max);
0031   
0032   if (unit != "") xAxisTitle += " [" + unit +  "]";  
0033   myHisto.GetXaxis()->SetTitle(xAxisTitle.c_str());
0034 
0035   float binWidth = (float) (max - min)/nBins;
0036   if (unit != "") yAxisTitle += " / " + to_string_with_precision(binWidth, precision) + " " + unit;
0037   myHisto.GetYaxis()->SetTitle(yAxisTitle.c_str());
0038   
0039   return myHisto;
0040 }
0041 
0042 TH1F makeHisto(int nBins, float* min, string type, string xAxisTitle, string unit, int precision, string yAxisTitle = "p_{T} Accept.")
0043 {
0044   string histo_name = type + "_" + xAxisTitle;
0045   TH1F myHisto(histo_name.c_str(), histo_name.c_str(), nBins, min);
0046 
0047   if (unit != "") xAxisTitle += " [" + unit +  "]";
0048   myHisto.GetXaxis()->SetTitle(xAxisTitle.c_str());
0049   myHisto.GetYaxis()->SetTitle(yAxisTitle.c_str());
0050 
0051   return myHisto;
0052 }
0053 
0054 //What branches to read into the analysis
0055 struct variableMap
0056 {
0057   float floats[21];
0058   double doubles[1];
0059   unsigned int ints[10];
0060 
0061   //any and variant are a pain here with TBranch 
0062   map<string, float> float_vars{
0063     {"Lambda0_mass", floats[0]},
0064     {"Lambda0_DIRA", floats[1]},
0065     {"Lambda0_chi2", floats[2]},
0066     {"K_S0_mass", floats[3]},
0067     {"K_S0_DIRA", floats[4]},
0068     {"K_S0_chi2", floats[5]},
0069     {"track_1_chi2", floats[6]},
0070     {"track_1_IP_xy", floats[7]},
0071     {"track_2_chi2", floats[8]},
0072     {"track_2_IP_xy", floats[9]},
0073     {"track_1_track_2_DCA", floats[10]},
0074     {"track_1_track_2_DCA_xy", floats[11]},
0075     {"track_1_pT", floats[12]},
0076     {"track_2_pT", floats[13]},
0077     {"Lambda0_pT", floats[13]},
0078     {"Lambda0_pseudorapidity", floats[14]},
0079     {"Lambda0_rapidity", floats[15]},
0080     {"Lambda0_phi", floats[16]},
0081     {"K_S0_pT", floats[17]},
0082     {"K_S0_pseudorapidity", floats[18]},
0083     {"K_S0_rapidity", floats[19]},
0084     {"K_S0_phi", floats[20]},
0085   };
0086     
0087   map<string, unsigned int> int_vars{
0088     {"Lambda0_nDoF", ints[0]},
0089     {"K_S0_nDoF", ints[1]},
0090     {"track_1_nDoF", ints[2]},
0091     {"track_2_nDoF", ints[3]},
0092     {"track_1_MVTX_nStates", ints[4]},
0093     {"track_1_INTT_nStates", ints[5]},
0094     {"track_1_TPC_nStates", ints[6]},
0095     {"track_2_MVTX_nStates", ints[7]},
0096     {"track_2_INTT_nStates", ints[8]},
0097     {"track_2_TPC_nStates", ints[9]},
0098   };
0099 
0100   map<string, double> double_vars{
0101     {"sWeight", doubles[0]},
0102   };
0103 };
0104 
0105 //What histograms to build for each ratio check
0106 class histos
0107 {
0108   public:
0109     TH1F Lambda0_all;
0110     TH1F K_S0_all;
0111     TH1F Lambda0_pT_accept;
0112     TH1F K_S0_pT_accept;
0113     TH1F ratio;
0114     TH1F inv_ratio;
0115     bool variable_bins = false;
0116     int nBins = 15;
0117 
0118   histos(bool constructor_variable_bins = false, float range = 1., string variable = "", string unit = "")
0119   {
0120     variable_bins = constructor_variable_bins;
0121     if (constructor_variable_bins)
0122     {
0123       Lambda0_all       = makeHisto(n_variable_bins, lower_bin_bounds, "Lambda0_all", variable.c_str(), unit.c_str(), 1);
0124       K_S0_all          = makeHisto(n_variable_bins, lower_bin_bounds, "K_S0_all", variable.c_str(), unit.c_str(), 1);
0125       Lambda0_pT_accept = makeHisto(n_variable_bins, lower_bin_bounds, "Lambda0_pT_accept", variable.c_str(), unit.c_str(), 1);
0126       K_S0_pT_accept    = makeHisto(n_variable_bins, lower_bin_bounds, "K_S0_pT_accept", variable.c_str(), unit.c_str(), 1);
0127       ratio             = makeHisto(n_variable_bins, lower_bin_bounds, "ratio", variable.c_str(), unit.c_str(), 1, "#Lambda^{0}/K_{S}^{0} p_{T} Accept.");
0128       inv_ratio         = makeHisto(n_variable_bins, lower_bin_bounds, "inv_ratio", variable.c_str(), unit.c_str(), 1, "#Lambda^{0}/K_{S}^{0} p_{T} Accept.");
0129     }
0130     else
0131     {
0132       Lambda0_all       = makeHisto(nBins, -1*range, range, "Lambda0_all", variable.c_str(), unit.c_str(), 1);
0133       K_S0_all          = makeHisto(nBins, -1*range, range, "K_S0_all", variable.c_str(), unit.c_str(), 1);
0134       Lambda0_pT_accept = makeHisto(nBins, -1*range, range, "Lambda0_pT_accept", variable.c_str(), unit.c_str(), 1);
0135       K_S0_pT_accept    = makeHisto(nBins, -1*range, range, "K_S0_pT_accept", variable.c_str(), unit.c_str(), 1);
0136       ratio             = makeHisto(nBins, -1*range, range, "ratio", variable.c_str(), unit.c_str(), 1, "#Lambda^{0}/K_{S}^{0} p_{T} Accept.");
0137       inv_ratio         = makeHisto(nBins, -1*range, range, "inv_ratio", variable.c_str(), unit.c_str(), 1, "#Lambda^{0}/K_{S}^{0} p_{T} Accept.");
0138     }
0139   }
0140 };
0141 
0142 histos pT(true, 0, "pT", "GeV");
0143 histos eta(false, 1.1, "#eta", "");
0144 histos phi(false, M_PI, "#phi", "");
0145 histos rap(false, 1.0, "y", "");
0146 
0147 template <typename T>
0148 void savePlots(T myPlot, string plotName, bool logY = false, float yMin = 0, float yMax = 1)
0149 {
0150   TGaxis::SetMaxDigits(3);
0151   string plotPath = "plots/";
0152   string makeDirectory = "mkdir -p " + plotPath;
0153   system(makeDirectory.c_str());
0154 
0155   TCanvas *c1  = new TCanvas("myCanvas", "myCanvas",800,800);
0156 
0157   myPlot.GetYaxis()->SetRangeUser(yMin, yMax);
0158 
0159   if (strncmp(typeid(myPlot).name(), "4TH2F", 5) == 0)
0160   {
0161     myPlot.Draw("COLZ");
0162   }
0163   else
0164   {
0165     myPlot.Sumw2();
0166     if (logY) gPad->SetLogy();
0167     myPlot.Draw("PE1");
0168   }
0169 
0170   TPaveText *pt;
0171   pt = new TPaveText(0.15,0.9,0.95,1., "NDC");
0172   pt->SetFillColor(0);
0173   pt->SetFillStyle(0);
0174   pt->SetTextFont(42);
0175   TText *pt_LaTex;
0176   if (global_isSim) pt->AddText("#it{#bf{sPHENIX}} Simulation, #sqrt{s} = 200 GeV, pp");
0177   else pt->AddText("#it{#bf{sPHENIX}} Internal, #sqrt{s} = 200 GeV, pp");
0178   pt->SetBorderSize(0);
0179   pt->Draw();
0180   gPad->Modified();
0181 
0182   string simAddition = global_isSim ? "_sim" : "";
0183 
0184   string extensions[] = {".C", ".pdf", ".png", ".root"};
0185   for (auto extension : extensions)
0186   {
0187     string output = plotPath + plotName + simAddition + extension;
0188     c1->SaveAs(output.c_str());
0189   }
0190 }
0191 
0192 bool isInRange(float min, float value, float max)
0193 {
0194   return min <= value && value <= max;
0195 }
0196 
0197 bool isAccepted(variableMap inputMap, bool isKshort = false)
0198 {
0199   if (isKshort)
0200   {
0201     if (!isInRange(0.4, inputMap.float_vars["K_S0_mass"], 0.6)) return false;
0202     if (inputMap.float_vars["K_S0_DIRA"] < 0.99) return false;
0203     if (inputMap.float_vars["K_S0_chi2"]/inputMap.int_vars["K_S0_nDoF"] > 20) return false;
0204   }
0205   else
0206   {
0207     if (!isInRange(1.08, inputMap.float_vars["Lambda0_mass"], 1.15)) return false;
0208     if (inputMap.float_vars["Lambda0_DIRA"] < 0.99) return false;
0209     if (inputMap.float_vars["Lambda0_chi2"]/inputMap.int_vars["Lambda0_nDoF"] > 20) return false;
0210   }
0211 
0212   if (inputMap.int_vars["track_1_MVTX_nStates"] < 1) return false;
0213   if (inputMap.int_vars["track_1_INTT_nStates"] < 1) return false;
0214   if (inputMap.int_vars["track_1_TPC_nStates"] < 20) return false;
0215   if (inputMap.int_vars["track_2_MVTX_nStates"] < 1) return false;
0216   if (inputMap.int_vars["track_2_INTT_nStates"] < 1) return false;
0217   if (inputMap.int_vars["track_2_TPC_nStates"] < 20) return false;
0218 
0219   if (inputMap.float_vars["track_1_chi2"]/inputMap.int_vars["track_1_nDoF"] > 300) return false;
0220   if (inputMap.float_vars["track_2_chi2"]/inputMap.int_vars["track_2_nDoF"] > 300) return false;
0221 
0222   if (abs(inputMap.float_vars["track_1_IP_xy"]) < 0.05) return false;
0223   if (abs(inputMap.float_vars["track_2_IP_xy"]) < 0.05) return false;
0224 
0225   if (abs(inputMap.float_vars["track_1_track_2_DCA"]) > 0.5) return false;
0226   if (abs(inputMap.float_vars["track_1_track_2_DCA_xy"]) > 1) return false;
0227 
0228   return true;
0229 }
0230 
0231 void processData(string type = "Kshort2pipi")
0232 {
0233   bool processingKshort = type == "Kshort2pipi" ? true : false;
0234 
0235   string dir = global_isSim ? "/sphenix/user/cdean/software/analysis/LightFlavorRatios/geometric_acceptance/simulation/"
0236                             : "/sphenix/user/cdean/scripts/fitters/files/";
0237   string fileName = processingKshort ? dir + "KShort6RunCombined_weighted.root"
0238                                           : dir + "Lambda6RunCombined_weighted.root";
0239   if (global_isSim)
0240   {
0241     fileName = processingKshort ? dir + "outputKFParticle_Kshort_reco.root"
0242                                     : dir + "outputKFParticle_Lambda_reco.root";
0243   }
0244 
0245   TFile *file = new TFile(fileName.c_str());
0246   TTree* data = (TTree*)file->Get("DecayTree");
0247 
0248   variableMap inputMap;
0249   for (auto &[branch, var] : inputMap.float_vars) data->SetBranchAddress(branch.c_str(), &var);
0250   for (auto &[branch, var] : inputMap.int_vars)   data->SetBranchAddress(branch.c_str(), &var);
0251   if (!global_isSim) for (auto &[branch, var] : inputMap.double_vars) data->SetBranchAddress(branch.c_str(), &var);
0252 
0253   int tmp = 0;
0254   int barWidth = 50;
0255 
0256   //Acceptance
0257   int num_entries = data->GetEntries();
0258   for (int  l = 0; l < num_entries; ++l)
0259   {
0260     if (tmp != (int)100*l/num_entries)
0261     {
0262       tmp = (int)100*l/num_entries;
0263       if ((tmp%1)  == 0)
0264       {
0265         cout << "[";
0266         int pos = barWidth * tmp/100;
0267         for (int i = 0; i < barWidth; ++i)
0268         {
0269           if (i < pos) cout << "=";
0270           else if (i == pos) cout << ">";
0271           else cout << " ";
0272         }
0273         cout << "] " << tmp << " %\r";
0274         cout.flush();
0275       }
0276     }
0277 
0278     data->GetEntry(l);
0279 
0280     double weight = global_isSim ? 1 : inputMap.double_vars["sWeight"];
0281 
0282     if (global_isSim)
0283     {
0284       bool checkCandidate = isAccepted(inputMap, processingKshort);
0285       if (!checkCandidate) continue;
0286     }
0287  
0288     if (processingKshort)
0289     {
0290       pT.K_S0_all.Fill(inputMap.float_vars["K_S0_pT"], weight);
0291       eta.K_S0_all.Fill(inputMap.float_vars["K_S0_pseudorapidity"], weight);
0292       rap.K_S0_all.Fill(inputMap.float_vars["K_S0_rapidity"], weight);
0293       phi.K_S0_all.Fill(inputMap.float_vars["K_S0_phi"], weight);
0294     }
0295     else
0296     {
0297       pT.Lambda0_all.Fill(inputMap.float_vars["Lambda0_pT"], weight);
0298       eta.Lambda0_all.Fill(inputMap.float_vars["Lambda0_pseudorapidity"], weight);
0299       rap.Lambda0_all.Fill(inputMap.float_vars["Lambda0_rapidity"], weight);
0300       phi.Lambda0_all.Fill(inputMap.float_vars["Lambda0_phi"], weight);
0301     }
0302 
0303     bool accepted = min(inputMap.float_vars["track_1_pT"], inputMap.float_vars["track_2_pT"]) >= 0.25 ? true : false;
0304 
0305     if (accepted)
0306     {
0307       if (processingKshort)
0308       {
0309         pT.K_S0_pT_accept.Fill(inputMap.float_vars["K_S0_pT"], weight);
0310         eta.K_S0_pT_accept.Fill(inputMap.float_vars["K_S0_pseudorapidity"], weight);
0311         rap.K_S0_pT_accept.Fill(inputMap.float_vars["K_S0_rapidity"], weight);
0312         phi.K_S0_pT_accept.Fill(inputMap.float_vars["K_S0_phi"], weight);
0313       }
0314       else
0315       {
0316         pT.Lambda0_pT_accept.Fill(inputMap.float_vars["Lambda0_pT"], weight);
0317         eta.Lambda0_pT_accept.Fill(inputMap.float_vars["Lambda0_pseudorapidity"], weight);
0318         rap.Lambda0_pT_accept.Fill(inputMap.float_vars["Lambda0_rapidity"], weight);
0319         phi.Lambda0_pT_accept.Fill(inputMap.float_vars["Lambda0_phi"], weight);
0320       }
0321     }
0322 
0323   }
0324 
0325   cout << "[";
0326   int pos = barWidth * tmp;
0327   for (int i = 0; i < barWidth; ++i)
0328   {
0329     if (i < pos) cout << "=";
0330     else if (i == pos) cout << ">";
0331     else cout << " ";
0332   }
0333   cout << "] 100 %\r";
0334   cout.flush();
0335   cout<<endl;
0336 }
0337 
0338 void makeRatios(histos &histoSet, string trailer)
0339 {
0340   string K_S0_saveName = "KS0_geometric_acceptance_ratio_" + trailer;
0341   string Lambda0_saveName = "Lambda0_geometric_acceptance_ratio_" + trailer;
0342   string ratio_saveName = "Lambda0_to_KS0_geometric_acceptance_ratio_" + trailer;
0343   string inv_ratio_saveName = "K_S0_to_Lambda0_geometric_acceptance_ratio_" + trailer;
0344 
0345   histoSet.Lambda0_all.Sumw2();
0346   histoSet.K_S0_all.Sumw2();
0347 
0348   histoSet.Lambda0_pT_accept.Sumw2();
0349   histoSet.K_S0_pT_accept.Sumw2();
0350 
0351   histoSet.Lambda0_pT_accept.Divide(&histoSet.Lambda0_all);
0352   histoSet.K_S0_pT_accept.Divide(&histoSet.K_S0_all);
0353 
0354   histoSet.ratio = histoSet.Lambda0_pT_accept;
0355   histoSet.inv_ratio = histoSet.K_S0_pT_accept;
0356 
0357   histoSet.ratio.Divide(&histoSet.K_S0_pT_accept);
0358   histoSet.inv_ratio.Divide(&histoSet.Lambda0_pT_accept);
0359 
0360   if (saveFinalPlots)
0361   {
0362     savePlots(histoSet.K_S0_pT_accept, K_S0_saveName.c_str(), false, 0, 1.1);
0363     savePlots(histoSet.Lambda0_pT_accept, Lambda0_saveName.c_str(), false, 0, 1.1);
0364 
0365     savePlots(histoSet.ratio, ratio_saveName.c_str(), false, 0, 1.1);
0366     savePlots(histoSet.inv_ratio, inv_ratio_saveName.c_str(), false, 0, 1.1);
0367   }
0368 }
0369 
0370 void makeTable(string var, histos theseHistos, string type = "Markdown")
0371 {
0372   cout << endl;
0373  
0374   string startline = type == "Markdown" ? "| " : "";
0375   string separator = type == "Markdown" ? " | " : " & ";
0376   string endline  = type == "Markdown" ? " | " : " \\\\ ";
0377 
0378   int nBins = theseHistos.variable_bins ? n_variable_bins : theseHistos.nBins;
0379 
0380   if (type != "Markdown") cout << "\\begin{tabular}{@{}ccccc@{}}" << endl << "\\toprule[1pt]" << endl;
0381   
0382   cout << startline << var << separator << "$K_S^0$" << separator << "$\\Lambda^0$" << separator << "$\\Lambda^0 / K_S^0$" << separator << "$K_S^0 / \\Lambda^0$" << endline << endl;
0383 
0384   if (type == "Markdown") cout << "|:--:|:--:|:--:|:--:|:--:|" << endl;
0385   else cout << "\\midrule[0.2pt]" << endl;
0386 
0387   for (int i = 1; i <= nBins; ++i)
0388   {
0389     string low = to_string_with_precision(theseHistos.ratio.GetXaxis()->GetBinLowEdge(i), 2);
0390     string high = to_string_with_precision(theseHistos.ratio.GetXaxis()->GetBinUpEdge(i), 2);
0391 
0392     string Ks0_content = to_string_with_precision(theseHistos.K_S0_pT_accept.GetBinContent(i), 4);
0393     string Ks0_error = to_string_with_precision(theseHistos.K_S0_pT_accept.GetBinError(i), 4);
0394 
0395     string Lambda_content = to_string_with_precision(theseHistos.Lambda0_pT_accept.GetBinContent(i), 4);
0396     string Lambda_error = to_string_with_precision(theseHistos.Lambda0_pT_accept.GetBinError(i), 4);
0397 
0398     string content = to_string_with_precision(theseHistos.ratio.GetBinContent(i), 2);
0399     string error = to_string_with_precision(theseHistos.ratio.GetBinError(i), 2);
0400 
0401     string inv_content = to_string_with_precision(theseHistos.inv_ratio.GetBinContent(i), 2);
0402     string inv_error = to_string_with_precision(theseHistos.inv_ratio.GetBinError(i), 2);
0403     cout << startline << low << " $\\rightarrow$ " << high << separator << Ks0_content << " $\\pm$ " << Ks0_error << separator << Lambda_content << " $\\pm$ " << Lambda_error << separator << content << " $\\pm$ " << error << separator << inv_content << " $\\pm$ " << inv_error << endline << endl;
0404   }
0405 
0406   if (type != "Markdown") cout << "\\bottomrule[1pt]" << endl << "\\end{tabular}" << endl;
0407 }
0408 
0409 void makeArray(string type, histos theseHistos)
0410 {
0411   cout << "Printing arrays for type " << type << endl;
0412   int nBins = theseHistos.variable_bins ? n_variable_bins : theseHistos.nBins;
0413   float xVal[nBins], xErr[nBins], yVal[nBins], yErr[nBins];
0414 
0415   for (int i = 0; i <= nBins; ++i)
0416   {
0417     xVal[i] = theseHistos.inv_ratio.GetXaxis()->GetBinCenter(i+1);
0418     xErr[i] = xVal[i] - theseHistos.inv_ratio.GetXaxis()->GetBinLowEdge(i+1);
0419     yVal[i] = theseHistos.inv_ratio.GetBinContent(i+1);
0420     yErr[i] = theseHistos.inv_ratio.GetBinError(i+1);
0421   }
0422 
0423   cout << "  float x[] = {";
0424   for (auto &val : xVal) cout << to_string_with_precision(val, 3) << ", ";
0425   cout << "};" << endl;
0426 
0427   cout << "  float ex[] = {";
0428   for (auto &val : xErr) cout << to_string_with_precision(val, 3) << ", ";
0429   cout << "};" << endl;
0430 
0431   cout << "  float y[] = {";
0432   for (auto &val : yVal) cout << to_string_with_precision(val, 3) << ", ";
0433   cout << "};" << endl;
0434 
0435   cout << "  float ey[] = {";
0436   for (auto &val : yErr) cout << to_string_with_precision(val, 3) << ", ";
0437   cout << "};" << endl;
0438 
0439   cout << "const int n = sizeof(x)/sizeof(x[0]);" << endl;
0440 }
0441 
0442 void calcpTCutRetention(bool isSim = true)
0443 {
0444   global_isSim = isSim;
0445 
0446   cout << "Processing K-short data" << endl;
0447   processData();
0448   cout << "Processing Lambda0 data" << endl;
0449   processData("Lambda02ppi");
0450 
0451   makeRatios(pT, "pT");
0452   makeRatios(eta, "eta");
0453   makeRatios(phi, "phi");
0454   makeRatios(rap, "rap");
0455 
0456   if (printMarkdownTables)
0457   {
0458     makeTable("$p_{T}$ [GeV]", pT, "Markdown");
0459     //makeTable("$\\eta$", eta, "Markdown");
0460     //makeTable("$y$", rap, "Markdown");
0461     //makeTable("$\\phi$", phi, "Markdown");
0462   }
0463 
0464   if (printLatexTables)
0465   {
0466     makeTable("$p_{T}$ [GeV]", pT, "LaTex");
0467     makeTable("$\\eta$", eta, "LaTex");
0468     makeTable("$y$", rap, "LaTex");
0469     makeTable("$\\phi$", phi, "LaTex");
0470   }
0471 
0472   if (printArrays)
0473   {
0474     makeArray("pT", pT);
0475   }
0476 }