Back to home page

sPhenix code displayed by LXR
 
 

    

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

0001 #include <cmath>
0002 #include <iostream>
0003 #include <sstream>
0004 #include <string>
0005 #ifndef __CINT__
0006 #include <RooGlobalFunc.h>
0007 #endif
0008 #include <RooBifurGauss.h>
0009 #include <RooDataHist.h>
0010 #include <RooDataSet.h>
0011 #include <RooGaussian.h>
0012 #include <RooRealVar.h>
0013 // #include <RooDoubleCB.h>
0014 #include <RooAbsDataHelper.h>
0015 #include <RooAddPdf.h>
0016 #include <RooExponential.h>
0017 #include <RooHist.h>
0018 #include <RooPlot.h>
0019 #include <RooStats/SPlot.h>
0020 
0021 #include <TBranch.h>
0022 #include <TCanvas.h>
0023 #include <TChain.h>
0024 #include <TCut.h>
0025 #include <TFile.h>
0026 #include <TH1F.h>
0027 #include <TMath.h>
0028 #include <TPad.h>
0029 #include <TTree.h>
0030 
0031 using namespace RooFit;
0032 using namespace std;
0033 
0034 void RDataframeToRoofit_SOnly(const bool doSnapshot,                 //
0035                           std::string snapshotName,              //
0036                           std::string inputfilename,             //
0037                           TCut selections,                       //
0038                           fitparam_config &fit_conf,             //
0039                           const std::string plotdir = "./figure" //
0040 )
0041 {
0042     RooMsgService::instance().setGlobalKillBelow(RooFit::ERROR);
0043 
0044     system(Form("mkdir -p %s", plotdir.c_str()));
0045 
0046     std::string dirName = snapshotName;
0047     std::string::size_type pos = dirName.find_last_of("/");
0048     if (pos != std::string::npos)
0049     {
0050         dirName = dirName.substr(0, pos);
0051         system(Form("mkdir -p %s", dirName.c_str()));
0052     }
0053 
0054     auto start = std::chrono::high_resolution_clock::now();
0055 
0056     ROOT::EnableImplicitMT();
0057     ROOT::RDataFrame df("DecayTree", inputfilename.c_str());
0058 
0059     auto filtered_df = df.Filter(selections.GetTitle());
0060     if (doSnapshot)
0061     {
0062         filtered_df.Snapshot("DecayTree", Form("%s", snapshotName.c_str()));
0063     }
0064 
0065     // print the number of entries in the dataframe
0066     int nEntries = filtered_df.Count().GetValue();
0067     std::cout << "Number of entries in the dataframe: " << nEntries << std::endl;
0068 
0069     FitParams::branch = fit_conf.branch;
0070     std::string xAxisTitle = fit_conf.decaystring + " candidate mass [GeV]";
0071 
0072     FitParams::minMass = fit_conf.minMass;
0073     FitParams::maxMass = fit_conf.maxMass;
0074     FitParams::mass.SetName(FitParams::branch.c_str());
0075     FitParams::mass.SetTitle("mass");
0076     FitParams::mass.setRange(FitParams::minMass, FitParams::maxMass);
0077 
0078     auto dataset_result = filtered_df.Book<float>(RooDataSetHelper("dataset", "dataset", RooArgSet(FitParams::mass)), {FitParams::branch.c_str()});
0079     RooDataSet *dataset = dataset_result.GetPtr();
0080 
0081     if (fit_conf.sigmodel == "Gaussian")
0082     {
0083         FitParams::mean.setVal(fit_conf.mean);
0084         FitParams::mean.setRange(fit_conf.mean_low, fit_conf.mean_high);
0085         FitParams::sigma.setVal(fit_conf.sigma);
0086         FitParams::sigma.setRange(fit_conf.sigma_low, fit_conf.sigma_high);
0087     }
0088     else if (fit_conf.sigmodel == "Voigtian")
0089     {
0090         FitParams::mean.setVal(fit_conf.mean);
0091         FitParams::mean.setRange(fit_conf.mean_low, fit_conf.mean_high);
0092         FitParams::sigma.setVal(fit_conf.sigma);
0093         FitParams::sigma.setRange(fit_conf.sigma_low, fit_conf.sigma_high);
0094         FitParams::width.setVal(fit_conf.width);
0095         FitParams::width.setRange(fit_conf.width_low, fit_conf.width_high);
0096     }
0097     else if (fit_conf.sigmodel == "CrystalBall")
0098     {
0099         FitParams::mean.setVal(fit_conf.mean);
0100         FitParams::mean.setRange(fit_conf.mean_low, fit_conf.mean_high);
0101         FitParams::sigma.setVal(fit_conf.sigma);
0102         FitParams::sigma.setRange(fit_conf.sigma_low, fit_conf.sigma_high);
0103         FitParams::alpha1.setVal(fit_conf.alpha1);
0104         FitParams::alpha1.setRange(fit_conf.alpha1_low, fit_conf.alpha1_high);
0105         FitParams::n1.setVal(fit_conf.n1);
0106         FitParams::n1.setRange(fit_conf.n1_low, fit_conf.n1_high);
0107     }
0108     else if (fit_conf.sigmodel == "DoubleCrystalBall")
0109     {
0110         FitParams::mean.setVal(fit_conf.mean);
0111         FitParams::mean.setRange(fit_conf.mean_low, fit_conf.mean_high);
0112         FitParams::sigma.setVal(fit_conf.sigma);
0113         FitParams::sigma.setRange(fit_conf.sigma_low, fit_conf.sigma_high);
0114         FitParams::alpha1.setVal(fit_conf.alpha1);
0115         FitParams::alpha1.setRange(fit_conf.alpha1_low, fit_conf.alpha1_high);
0116         FitParams::n1.setVal(fit_conf.n1);
0117         FitParams::n1.setRange(fit_conf.n1_low, fit_conf.n1_high);
0118         FitParams::alpha2.setVal(fit_conf.alpha2);
0119         FitParams::alpha2.setRange(fit_conf.alpha2_low, fit_conf.alpha2_high);
0120         FitParams::n2.setVal(fit_conf.n2);
0121         FitParams::n2.setRange(fit_conf.n2_low, fit_conf.n2_high);
0122         FitParams::frac.setVal(fit_conf.frac);
0123         FitParams::frac.setRange(0, 1);
0124     }
0125     else
0126     {
0127         std::cerr << "Unknown signal model type: " << fit_conf.sigmodel << ", returning nullptr. Please define the model." << std::endl;
0128     }
0129 
0130     RooAbsPdf *signal = signalModel(fit_conf.sigmodel, FitParams::mass);
0131     FitParams::nSig.setVal(float(nEntries));
0132     FitParams::nSig.setRange(float(nEntries)*0.95, float(nEntries)*1.05);
0133     RooAddPdf *model = new RooAddPdf("model", "model", RooArgList(*signal), RooArgList(FitParams::nSig));
0134 
0135     RooFitResult *m_fitres(0);
0136     m_fitres = model->fitTo(*dataset, Save(kTRUE), Extended(kTRUE), PrintLevel(-1));
0137     m_fitres->Print();
0138     // Save the fit result to a file
0139     TFile *fitresfile = new TFile(Form("%s/fitresults_signal.root", plotdir.c_str()), "RECREATE");
0140     fitresfile->cd();
0141     m_fitres->Write("fitres");
0142     fitresfile->Close();
0143 
0144     std::cout << "Making plots..." << std::endl;
0145 
0146     // make plot
0147     std::string plotTitle = "";
0148     RooPlot *frame = FitParams::mass.frame(Title(plotTitle.c_str())); // creating the frame
0149     RooBinning bins(FitParams::minMass, FitParams::maxMass);
0150     bins.addUniform(fit_conf.nBins, FitParams::minMass, FitParams::maxMass);
0151     dataset->plotOn(frame, DrawOption("PE1"), Binning(bins), XErrorSize(0), DataError(RooAbsData::SumW2), RooFit::Name("data")); // plotting the raw unbinned data on the frame
0152     signal->plotOn(frame, LineColor(TColor::GetColor("#FF7F0E")), RooFit::Name("signalpdf"));                                                           // plotting the fit onto the frame
0153 
0154     RooHist *pull = frame->pullHist();
0155     RooPlot *frame2 = FitParams::mass.frame(Title(""));
0156     frame2->addPlotable(pull, "PE1");
0157 
0158     std::cout << "Creating canvas..." << std::endl;
0159 
0160     TCanvas *c = new TCanvas("massFitCanvas", "massFitCanvas", 800, 800);
0161     TPad mainPad("mainPad", "mainPad", 0., 0.3, 1., 1.);
0162     mainPad.SetTopMargin(TopMargin);
0163     mainPad.SetBottomMargin(0);
0164     mainPad.Draw();
0165     TPad pullPad("pullPad", "pullPad", 0., 0.0, 1., 0.3);
0166     pullPad.SetBottomMargin(0.5);
0167     pullPad.SetTopMargin(0);
0168     pullPad.Draw();
0169     mainPad.cd();
0170     frame->SetMarkerStyle(kCircle);
0171     frame->SetMarkerSize(0.02);
0172     frame->SetLineWidth(1);
0173     frame->GetXaxis()->SetTitleSize(0);
0174     frame->GetXaxis()->SetLabelSize(0);
0175     frame->GetYaxis()->SetTitleSize(AxisTitleSize * textscale_pad1);
0176     frame->GetYaxis()->SetLabelSize(AxisLabelSize * textscale_pad1);
0177     frame->GetYaxis()->SetTitleOffset(1.2);
0178     frame->GetYaxis()->SetTitleFont(42);
0179     frame->GetYaxis()->SetLabelFont(42);
0180     TH1 *hdataset = dataset->createHistogram("hdataset", FitParams::mass, Binning(fit_conf.nBins, FitParams::minMass, FitParams::maxMass));
0181     // frame->GetYaxis()->SetRangeUser(hdataset->GetMinimum(0) * 0.85, hdataset->GetMaximum() * 1.4);
0182     frame->GetYaxis()->SetRangeUser(0.1, hdataset->GetMaximum() * 1.8);
0183     float binWidth = 1000. * (FitParams::maxMass - FitParams::minMass) / fit_conf.nBins;
0184     string yAxisTitle = "Candidates / (" + to_string_with_precision(binWidth, 1) + " MeV)";
0185     frame->GetYaxis()->SetTitle(yAxisTitle.c_str());
0186     frame->Draw();
0187     frame->Print();
0188     c->RedrawAxis();
0189 
0190     TLatex *datestamp = new TLatex();
0191     datestamp->SetTextSize(0.06);
0192     datestamp->SetTextAlign(kHAlignRight + kVAlignBottom);
0193     datestamp->SetNDC();
0194     datestamp->DrawLatex(1 - mainPad.GetRightMargin(),      //
0195                          1 - mainPad.GetTopMargin() + 0.01, //
0196                          getTime().c_str()                  //
0197     );
0198 
0199     TLegend *sphnxleg = new TLegend(mainPad.GetLeftMargin() + 0.03,   //
0200                                     1 - mainPad.GetTopMargin() - 0.2, //
0201                                     mainPad.GetLeftMargin() + 0.2,    //
0202                                     1 - mainPad.GetTopMargin() - 0.05 //
0203     );
0204     sphnxleg->SetTextAlign(kHAlignLeft + kVAlignCenter);
0205     sphnxleg->SetTextSize(0.06);
0206     sphnxleg->SetFillStyle(0);
0207     sphnxleg->AddEntry("", Form("#it{#bf{sPHENIX}} %s", prelimtext.c_str()), "");
0208     sphnxleg->AddEntry("", "p+p #sqrt{s_{NN}}=200 GeV", "");
0209     sphnxleg->Draw();
0210 
0211     TLegend *leg = new TLegend(1 - mainPad.GetRightMargin() - 0.27, //
0212                                1 - mainPad.GetTopMargin() - 0.2,   //
0213                                1 - mainPad.GetRightMargin() - 0.1,  //
0214                                1 - mainPad.GetTopMargin() - 0.06    //
0215     );
0216     leg->AddEntry(frame->findObject("data"), "Data", "PE2");
0217     leg->AddEntry(frame->findObject("signalpdf"), "Fit", "L");
0218     leg->SetFillColor(0);
0219     leg->SetFillStyle(0);
0220     leg->SetBorderSize(0);
0221     leg->SetTextSize(0.055);
0222     leg->Draw();
0223 
0224     // plot the significance
0225     float spacing = 0.05;
0226     TLatex *fitparams = new TLatex();
0227     fitparams->SetTextSize(0.045);
0228     fitparams->SetTextAlign(kHAlignLeft + kVAlignCenter);
0229     fitparams->SetNDC();
0230     fitparams->DrawLatex(mainPad.GetLeftMargin() + 0.07,                                                                     //
0231                          1 - mainPad.GetTopMargin() - 0.15 - spacing - 0.06,                                                  //
0232                          Form("#mu = %.0f #pm %.0g MeV", FitParams::mean.getVal() * 1000, FitParams::mean.getError() * 1000) //
0233     );
0234     fitparams->DrawLatex(mainPad.GetLeftMargin() + 0.07,                                                                          //
0235                          1 - mainPad.GetTopMargin() - 0.15 - spacing - 0.06 - 0.055,                                               //
0236                          Form("#sigma = %.2f #pm %.1g MeV", FitParams::sigma.getVal() * 1000, FitParams::sigma.getError() * 1000) //
0237     );
0238 
0239     gPad->Modified();
0240     pullPad.cd();
0241     frame2->SetMarkerStyle(kCircle);
0242     frame2->SetMarkerSize(0.02);
0243     frame2->SetTitle("");
0244     frame2->GetXaxis()->SetTitle(xAxisTitle.c_str());
0245     frame2->GetXaxis()->SetTitleOffset(1.3);
0246     frame2->GetXaxis()->SetTitleFont(42);
0247     frame2->GetXaxis()->SetTitleSize(AxisTitleSize * textscale_pad2);
0248     frame2->GetXaxis()->SetLabelFont(42);
0249     frame2->GetXaxis()->SetLabelSize(AxisLabelSize * textscale_pad2);
0250     frame2->GetYaxis()->SetTitle("Pull");
0251     frame2->GetYaxis()->SetTitleOffset(0.5);
0252     frame2->GetYaxis()->SetTitleFont(42);
0253     frame2->GetYaxis()->SetTitleSize(AxisTitleSize * textscale_pad2);
0254     frame2->GetYaxis()->SetLabelFont(42);
0255     frame2->GetYaxis()->SetLabelSize(AxisLabelSize * textscale_pad2);
0256     frame2->GetYaxis()->SetRangeUser(-6, 6);
0257     frame2->GetYaxis()->SetNdivisions(5);
0258     TLine *plusThreeLine = new TLine(FitParams::minMass, 3, FitParams::maxMass, 3);
0259     plusThreeLine->SetLineColor(1);
0260     plusThreeLine->SetLineStyle(2);
0261     plusThreeLine->SetLineWidth(2);
0262     TLine *zeroLine = new TLine(FitParams::minMass, 0, FitParams::maxMass, 0);
0263     zeroLine->SetLineColor(1);
0264     zeroLine->SetLineStyle(2);
0265     zeroLine->SetLineWidth(2);
0266     TLine *minusThreeLine = new TLine(FitParams::minMass, -3, FitParams::maxMass, -3);
0267     minusThreeLine->SetLineColor(1);
0268     minusThreeLine->SetLineStyle(2);
0269     minusThreeLine->SetLineWidth(2);
0270     frame2->Draw();
0271     plusThreeLine->Draw("same");
0272     zeroLine->Draw("same");
0273     minusThreeLine->Draw("same");
0274 
0275     std::cout << "Saving plots to " << plotdir << std::endl;
0276 
0277     vector<string> extensions = {".C", ".pdf", ".png"};
0278     for (auto &extension : extensions)
0279     {
0280         c->SaveAs(Form("%s/RDataframeToRoofit_%s_signal%s", plotdir.c_str(), FitParams::branch.c_str(), extension.c_str()));
0281     }
0282 
0283     auto end = std::chrono::high_resolution_clock::now();
0284     std::chrono::duration<double> elapsed = end - start;
0285     std::cout << "Elapsed time: " << elapsed.count() / 60. << " minutes" << std::endl;
0286 }
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!