LightFlavorRatios/unbinnedMaxLikelihood_fitter/getV0RawYield.C

0001 #include <algorithm>
0002 #include <cmath>
0003 #include <cctype>
0004 #include <iomanip>
0005 #include <iostream>
0006 #include <memory>
0007 #include <sstream>
0008 #include <string>
0009 #include <vector>
0010
0011 #include <RooArgList.h>
0012 #include <RooFitResult.h>
0013 #include <RooRealVar.h>
0014 #include <TCanvas.h>
0015 #include <TFile.h>
0016 #include <TGaxis.h>
0017 #include <TH1.h>
0018 #include <TH1D.h>
0019 #include <TLatex.h>
0020 #include <TSystem.h>
0021
0022 struct V0Config
0023 {
0024     std::string name;
0025     std::string input_dir;
0026     std::string cut_prefix;
0027     std::string output_prefix;
0028     std::string plot_label;
0029 };
0030
0031 void draw1Dhistogram(TH1 *hist, bool logx, bool logy, std::string xtitle, std::string ytitle, std::vector<std::string> addinfo, std::string filename)
0032 {
0033     TCanvas *c = new TCanvas("c", "c", 800, 700);
0034     gPad->SetTopMargin(0.07);
0035     c->cd();
0036     c->SetLogx(logx);
0037     c->SetLogy(logy);
0038     hist->GetXaxis()->SetTitle(xtitle.c_str());
0039     hist->GetYaxis()->SetTitle(ytitle.c_str());
0040     hist->GetYaxis()->SetTitleOffset(1.6);
0041     hist->GetYaxis()->SetRangeUser((logy) ? hist->GetMinimum(0) * 0.5 : 0, (logy) ? hist->GetMaximum() * 50 : hist->GetMaximum() * 1.4);
0042     hist->Draw("e1");
0043
0044     TLatex *latex = new TLatex();
0045     latex->SetTextSize(0.04);
0046     latex->SetTextAlign(12);
0047     latex->SetNDC();
0048     for (size_t i = 0; i < addinfo.size(); ++i)
0049     {
0050         latex->DrawLatex(0.2, 0.88 - i * 0.05, addinfo[i].c_str());
0051     }
0052     c->SaveAs(Form("%s.png", filename.c_str()));
0053     c->SaveAs(Form("%s.pdf", filename.c_str()));
0054     delete latex;
0055     delete c;
0056 }
0057
0058 void replace_all(std::string &s, const std::string &from, const std::string &to)
0059 {
0060     size_t pos = 0;
0061     while ((pos = s.find(from, pos)) != std::string::npos)
0062     {
0063         s.replace(pos, from.size(), to);
0064         pos += to.size();
0065     }
0066 }
0067
0068 std::string cut_to_dir_tag(std::string cut)
0069 {
0070     replace_all(cut, ">=", "geq");
0071     replace_all(cut, "<=", "leq");
0072     replace_all(cut, "==", "eq");
0073     replace_all(cut, "&&", "AND");
0074     replace_all(cut, "||", "OR");
0075     replace_all(cut, ">", "gt");
0076     replace_all(cut, "<", "lt");
0077
0078     std::replace(cut.begin(), cut.end(), '.', 'p');
0079     std::replace(cut.begin(), cut.end(), '-', 'm');
0080     std::replace(cut.begin(), cut.end(), '+', 'p');
0081     return cut;
0082 }
0083
0084 std::string pt_cut(const std::string &cut_prefix, double pt_min, double pt_max)
0085 {
0086     std::ostringstream oss;
0087     oss << std::fixed << std::setprecision(2)
0088         << cut_prefix << "_pT>=" << pt_min << "&&" << cut_prefix << "_pT<" << pt_max;
0089     return oss.str();
0090 }
0091
0092 bool read_nsig(const std::string &filename, double &yield, double &error)
0093 {
0094     std::unique_ptr<TFile> file(TFile::Open(filename.c_str(), "READ"));
0095     if (!file || file->IsZombie())
0096     {
0097         std::cerr << "WARNING: cannot open " << filename << std::endl;
0098         return false;
0099     }
0100
0101     RooFitResult *fitres = dynamic_cast<RooFitResult *>(file->Get("fitres"));
0102     if (!fitres)
0103     {
0104         std::cerr << "WARNING: fitres not found in " << filename << std::endl;
0105         return false;
0106     }
0107
0108     const RooRealVar *nSig = dynamic_cast<const RooRealVar *>(fitres->floatParsFinal().find("nSig"));
0109     if (!nSig)
0110         nSig = dynamic_cast<const RooRealVar *>(fitres->constPars().find("nSig"));
0111     if (!nSig)
0112     {
0113         std::cerr << "WARNING: nSig not found in " << filename << std::endl;
0114         return false;
0115     }
0116
0117     yield = nSig->getVal();
0118     error = nSig->getError();
0119     return true;
0120 }
0121
0122 bool make_v0_config(std::string v0, V0Config &config)
0123 {
0124     std::transform(v0.begin(), v0.end(), v0.begin(), [](unsigned char c) { return std::tolower(c); });
0125
0126     if (v0 == "lambda" || v0 == "lambda0")
0127     {
0128         config = {"lambda", "./figure_lambda_crosscheck", "Lambda0", "lambda", "#Lambda^{0}"};
0129         return true;
0130     }
0131     if (v0 == "kshort" || v0 == "ks" || v0 == "ks0" || v0 == "k_s0")
0132     {
0133         config = {"kshort", "./figure_kshort_crosscheck", "K_S0", "kshort", "K_{S}^{0}"};
0134         return true;
0135     }
0136
0137     std::cerr << "ERROR: unsupported V0 species \"" << v0 << "\". Use \"lambda\", \"kshort\", or \"both\"." << std::endl;
0138     return false;
0139 }
0140
0141 void get_one_v0_raw_yield(const V0Config &config)
0142 {
0143     const std::vector<double> pt_bins = {0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.8, 2.1, 2.4, 2.7, 3.0, 4.0};
0144
0145     std::vector<std::string> subdirs;
0146     void *dir = gSystem->OpenDirectory(config.input_dir.c_str());
0147     if (!dir)
0148     {
0149         std::cerr << "ERROR: cannot open " << config.input_dir << std::endl;
0150         return;
0151     }
0152
0153     const char *entry = nullptr;
0154     while ((entry = gSystem->GetDirEntry(dir)))
0155     {
0156         const std::string dirname = entry;
0157         if (dirname == "." || dirname == "..")
0158             continue;
0159
0160         const std::string rootfile = config.input_dir + "/" + dirname + "/fitresults.root";
0161         if (!gSystem->AccessPathName(rootfile.c_str()))
0162             subdirs.push_back(dirname);
0163     }
0164     gSystem->FreeDirectory(dir);
0165     std::sort(subdirs.begin(), subdirs.end());
0166
0167     TGaxis::SetMaxDigits(3);
0168     TH1D *hYield = new TH1D(Form("h%sRawYieldPt", config.output_prefix.c_str()), Form(";%s p_{T} [GeV];Raw %s yield", config.plot_label.c_str(), config.plot_label.c_str()), pt_bins.size() - 1, pt_bins.data());
0169     hYield->Sumw2();
0170     std::vector<int> fit_counts(pt_bins.size() - 1, 0);
0171
0172     std::cout << "\nRaw-yield summary for " << config.name << std::endl;
0173     for (size_t i = 0; i < pt_bins.size() - 1; ++i)
0174     {
0175         const double pt_min = pt_bins[i];
0176         const double pt_max = pt_bins[i + 1];
0177         const std::string pt_tag = cut_to_dir_tag(pt_cut(config.cut_prefix, pt_min, pt_max));
0178
0179         std::vector<std::string> matching_phi_dirs;
0180         std::vector<std::string> matching_pt_dirs;
0181         for (const std::string &dirname : subdirs)
0182         {
0183             if (dirname.find(pt_tag) != 0)
0184                 continue;
0185
0186             if (dirname.find(config.cut_prefix + "_phi") != std::string::npos)
0187                 matching_phi_dirs.push_back(dirname);
0188             else
0189                 matching_pt_dirs.push_back(dirname);
0190         }
0191
0192         const std::vector<std::string> &dirs_to_sum = matching_phi_dirs.empty() ? matching_pt_dirs : matching_phi_dirs;
0193         double yield_sum = 0.0;
0194         double error2_sum = 0.0;
0195
0196         for (const std::string &dirname : dirs_to_sum)
0197         {
0198             double yield = 0.0;
0199             double error = 0.0;
0200             const std::string rootfile = config.input_dir + "/" + dirname + "/fitresults.root";
0201             if (!read_nsig(rootfile, yield, error))
0202                 continue;
0203
0204             yield_sum += yield;
0205             error2_sum += error * error;
0206         }
0207
0208         hYield->SetBinContent(i + 1, yield_sum);
0209         hYield->SetBinError(i + 1, std::sqrt(error2_sum));
0210         fit_counts[i] = dirs_to_sum.size();
0211
0212         std::cout << std::fixed << std::setprecision(2)
0213                   << pt_min << " <= pT < " << pt_max << ": "
0214                   << yield_sum << " +/- " << std::sqrt(error2_sum)
0215                   << " from " << dirs_to_sum.size() << " fit result file(s)" << std::endl;
0216     }
0217
0218     std::cout << "\n| V0 | pT bin [GeV] | Raw yield | Fit uncertainty |\n"
0219               << "|---|---:|---:|---:|\n";
0220     for (size_t i = 0; i < pt_bins.size() - 1; ++i)
0221     {
0222         std::cout << std::fixed << std::setprecision(2)
0223                   << "| " << config.name << " | [" << pt_bins[i] << ", " << pt_bins[i + 1] << ") | "
0224                   << hYield->GetBinContent(i + 1) << " | "
0225                   << hYield->GetBinError(i + 1) << " |\n ";
0226     }
0227     std::cout << std::endl;
0228
0229     const std::string output_name = config.output_prefix + "_yield_pt";
0230     TFile outfile((config.input_dir + "/" + output_name + ".root").c_str(), "RECREATE");
0231     hYield->Write();
0232     outfile.Close();
0233
0234     draw1Dhistogram(hYield, false, false, config.plot_label + " p_{T} [GeV]", "Raw " + config.plot_label + " yield", {}, config.input_dir + "/" + output_name);
0235 }
0236
0237 void getV0RawYield(std::string v0 = "lambda")
0238 {
0239     std::transform(v0.begin(), v0.end(), v0.begin(), [](unsigned char c) { return std::tolower(c); });
0240     if (v0 == "both" || v0 == "all")
0241     {
0242         V0Config lambda_config;
0243         V0Config kshort_config;
0244         if (make_v0_config("lambda", lambda_config))
0245             get_one_v0_raw_yield(lambda_config);
0246         if (make_v0_config("kshort", kshort_config))
0247             get_one_v0_raw_yield(kshort_config);
0248         return;
0249     }
0250
0251     V0Config config;
0252     if (!make_v0_config(v0, config))
0253         return;
0254
0255     get_one_v0_raw_yield(config);
0256 }