Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:09:54

0001 #include "RooGaussian.h"
0002 #include "RooChebychev.h"
0003 #include "RooVoigtian.h"
0004 #include "RooExponential.h"
0005 #include "RooRealVar.h"
0006 #include "RooFormulaVar.h"
0007 #include "RooDataHist.h"
0008 #include "RooAddPdf.h"
0009 #include "RooGlobalFunc.h"
0010 #include "RooCategory.h"
0011 #include "RooSimultaneous.h"
0012 #include "RooFitResult.h"
0013 #include "TFile.h"
0014 #include "TH1F.h"
0015 #include "TH2F.h"
0016 #include "TH3F.h"
0017 #include "TProfile2D.h"
0018 #include "TCanvas.h"
0019 #include "RooPlot.h"
0020 
0021 namespace dqmTnP {
0022 
0023   class AbstractFitter {
0024   protected:
0025     RooRealVar mass;
0026     RooRealVar mean;
0027     double expectedMean;
0028     RooRealVar sigma;
0029     double expectedSigma;
0030     RooRealVar efficiency;
0031     RooRealVar nSignalAll;
0032     RooFormulaVar nSignalPass;
0033     RooFormulaVar nSignalFail;
0034     RooRealVar nBackgroundFail;
0035     RooRealVar nBackgroundPass;
0036     RooCategory category;
0037     RooSimultaneous simPdf;
0038     RooDataHist* data;
0039     double chi2;
0040     bool verbose;
0041 
0042   public:
0043     AbstractFitter(bool verbose_ = false)
0044         : mass("mass", "mass", 0., 100., "GeV"),
0045           mean("mean", "mean", 0., 100., "GeV"),
0046           sigma("sigma", "sigma", 0., 100., "GeV"),
0047           efficiency("efficiency", "efficiency", 0.5, 0.0, 1.0),
0048           nSignalAll("nSignalAll", "nSignalAll", 0., 1e10),
0049           nSignalPass("nSignalPass", "nSignalAll*efficiency", RooArgList(nSignalAll, efficiency)),
0050           nSignalFail("nSignalFail", "nSignalAll*(1-efficiency)", RooArgList(nSignalAll, efficiency)),
0051           nBackgroundFail("nBackgroundFail", "nBackgroundFail", 0., 1e10),
0052           nBackgroundPass("nBackgroundPass", "nBackgroundPass", 0., 1e10),
0053           category("category", "category"),
0054           simPdf("simPdf", "simPdf", category),
0055           data(nullptr),
0056           verbose(verbose_) {
0057       //turn on/off default messaging of roofit
0058       RooMsgService::instance().setSilentMode(!verbose ? kTRUE : kFALSE);
0059       for (int i = 0; i < RooMsgService::instance().numStreams(); i++) {
0060         RooMsgService::instance().setStreamStatus(i, verbose ? kTRUE : kFALSE);
0061       }
0062       category.defineType("pass");
0063       category.defineType("fail");
0064     };
0065     virtual ~AbstractFitter() = default;
0066     ;
0067     void setup(double expectedMean_, double massLow, double massHigh, double expectedSigma_) {
0068       expectedMean = expectedMean_;
0069       expectedSigma = expectedSigma_;
0070       mass.setRange(massLow, massHigh);
0071       mean.setRange(massLow, massHigh);
0072     }
0073     virtual void fit(TH1* num, TH1* den) = 0;
0074     double getEfficiency() { return efficiency.getVal(); }
0075     double getEfficiencyError() { return efficiency.getError(); }
0076     double getChi2() { return chi2; }
0077     void savePlot(const TString& name) {
0078       using namespace RooFit;
0079       RooPlot* frame = mass.frame(Name(name), Title("Failing and Passing Probe Distributions"));
0080       data->plotOn(frame, Cut("category==category::pass"), LineColor(kGreen), MarkerColor(kGreen));
0081       data->plotOn(frame, Cut("category==category::fail"), LineColor(kRed), MarkerColor(kRed));
0082       simPdf.plotOn(frame, Slice(category, "pass"), ProjWData(category, *data), LineColor(kGreen));
0083       simPdf.plotOn(frame, Slice(category, "fail"), ProjWData(category, *data), LineColor(kRed));
0084       simPdf.paramOn(frame, Layout(0.58, 0.99, 0.99));
0085       data->statOn(frame, Layout(0.70, 0.99, 0.5));
0086       frame->Write();
0087       delete frame;
0088     }
0089 
0090     TString calculateEfficiency(
0091         TH3* pass, TH3* all, int massDimension, TProfile2D*& eff, TProfile2D*& effChi2, const TString& plotName = "") {
0092       //sort out the TAxis
0093       TAxis *par1Axis, *par2Axis, *massAxis;
0094       int par1C, par2C, massC;
0095       if (massDimension == 1) {
0096         massAxis = all->GetXaxis();
0097         massC = 1;
0098         par1Axis = all->GetYaxis();
0099         par1C = all->GetXaxis()->GetNbins() + 2;
0100         par2Axis = all->GetZaxis();
0101         par2C = (all->GetXaxis()->GetNbins() + 2) * (all->GetYaxis()->GetNbins() + 2);
0102       } else if (massDimension == 2) {
0103         par1Axis = all->GetXaxis();
0104         par1C = 1;
0105         massAxis = all->GetYaxis();
0106         massC = all->GetXaxis()->GetNbins() + 2;
0107         par2Axis = all->GetZaxis();
0108         par2C = (all->GetXaxis()->GetNbins() + 2) * (all->GetYaxis()->GetNbins() + 2);
0109       } else if (massDimension == 3) {
0110         par1Axis = all->GetXaxis();
0111         par1C = 1;
0112         par2Axis = all->GetYaxis();
0113         par2C = all->GetXaxis()->GetNbins() + 2;
0114         massAxis = all->GetZaxis();
0115         massC = (all->GetXaxis()->GetNbins() + 2) * (all->GetYaxis()->GetNbins() + 2);
0116       } else {
0117         return "massDimension > 3 !, skipping...";
0118       }
0119 
0120       //create eff and effChi2 TProfiles
0121       if (!par1Axis || !par2Axis)
0122         return "No par1Axis or par2Axis!";
0123       if (par1Axis->GetXbins()->GetSize() == 0 && par2Axis->GetXbins()->GetSize() == 0) {
0124         eff = new TProfile2D("efficiency",
0125                              "efficiency",
0126                              par1Axis->GetNbins(),
0127                              par1Axis->GetXmin(),
0128                              par1Axis->GetXmax(),
0129                              par2Axis->GetNbins(),
0130                              par2Axis->GetXmin(),
0131                              par2Axis->GetXmax());
0132       } else if (par1Axis->GetXbins()->GetSize() == 0) {
0133         eff = new TProfile2D("efficiency",
0134                              "efficiency",
0135                              par1Axis->GetNbins(),
0136                              par1Axis->GetXmin(),
0137                              par1Axis->GetXmax(),
0138                              par2Axis->GetNbins(),
0139                              par2Axis->GetXbins()->GetArray());
0140       } else if (par2Axis->GetXbins()->GetSize() == 0) {
0141         eff = new TProfile2D("efficiency",
0142                              "efficiency",
0143                              par1Axis->GetNbins(),
0144                              par1Axis->GetXbins()->GetArray(),
0145                              par2Axis->GetNbins(),
0146                              par2Axis->GetXmin(),
0147                              par2Axis->GetXmax());
0148       } else {
0149         eff = new TProfile2D("efficiency",
0150                              "efficiency",
0151                              par1Axis->GetNbins(),
0152                              par1Axis->GetXbins()->GetArray(),
0153                              par2Axis->GetNbins(),
0154                              par2Axis->GetXbins()->GetArray());
0155       }
0156       eff->SetTitle("");
0157       eff->SetXTitle(par1Axis->GetTitle());
0158       eff->SetYTitle(par2Axis->GetTitle());
0159       eff->SetStats(kFALSE);
0160       effChi2 = (TProfile2D*)eff->Clone("efficiencyChi2");
0161       eff->SetZTitle("Efficiency");
0162       eff->SetOption("colztexte");
0163       eff->GetZaxis()->SetRangeUser(-0.001, 1.001);
0164       effChi2->SetZTitle("Chi^2/NDF");
0165       effChi2->SetOption("colztext");
0166 
0167       //create the 1D mass distribution container histograms
0168       TH1D* all1D = (massAxis->GetXbins()->GetSize() == 0)
0169                         ? new TH1D("all1D", "all1D", massAxis->GetNbins(), massAxis->GetXmin(), massAxis->GetXmax())
0170                         : new TH1D("all1D", "all1D", massAxis->GetNbins(), massAxis->GetXbins()->GetArray());
0171       auto* pass1D = (TH1D*)all1D->Clone("pass1D");
0172 
0173       //for each parameter bin fit the mass distributions
0174       for (int par1 = 1; par1 <= par1Axis->GetNbins(); par1++) {
0175         for (int par2 = 1; par2 <= par2Axis->GetNbins(); par2++) {
0176           for (int mass = 1; mass <= massAxis->GetNbins(); mass++) {
0177             int index = par1 * par1C + par2 * par2C + mass * massC;
0178             all1D->SetBinContent(mass, all->GetBinContent(index));
0179             pass1D->SetBinContent(mass, pass->GetBinContent(index));
0180           }
0181           fit(pass1D, all1D);
0182           int index = par1 + par2 * (par1Axis->GetNbins() + 2);
0183           eff->SetBinContent(index, getEfficiency());
0184           eff->SetBinEntries(index, 1);
0185           eff->SetBinError(index,
0186                            sqrt(getEfficiency() * getEfficiency() + getEfficiencyError() * getEfficiencyError()));
0187           effChi2->SetBinContent(index, getChi2());
0188           effChi2->SetBinEntries(index, 1);
0189           if (plotName != "") {
0190             savePlot(TString::Format("%s_%d_%d", plotName.Data(), par1, par2));
0191           }
0192         }
0193       }
0194       delete all1D;
0195       delete pass1D;
0196       return "";  //OK
0197     }
0198 
0199     TString calculateEfficiency(
0200         TH2* pass, TH2* all, int massDimension, TProfile*& eff, TProfile*& effChi2, const TString& plotName = "") {
0201       //sort out the TAxis
0202       TAxis *par1Axis, *massAxis;
0203       int par1C, massC;
0204       if (massDimension == 1) {
0205         massAxis = all->GetXaxis();
0206         massC = 1;
0207         par1Axis = all->GetYaxis();
0208         par1C = all->GetXaxis()->GetNbins() + 2;
0209       } else if (massDimension == 2) {
0210         par1Axis = all->GetXaxis();
0211         par1C = 1;
0212         massAxis = all->GetYaxis();
0213         massC = all->GetXaxis()->GetNbins() + 2;
0214       } else {
0215         return "massDimension > 2 !, skipping...";
0216       }
0217 
0218       //create eff and effChi2 TProfiles
0219       if (!par1Axis)
0220         return "No par1Axis!";
0221       eff =
0222           (par1Axis->GetXbins()->GetSize() == 0)
0223               ? new TProfile("efficiency", "efficiency", par1Axis->GetNbins(), par1Axis->GetXmin(), par1Axis->GetXmax())
0224               : new TProfile("efficiency", "efficiency", par1Axis->GetNbins(), par1Axis->GetXbins()->GetArray());
0225       eff->SetTitle("");
0226       eff->SetXTitle(par1Axis->GetTitle());
0227       eff->SetLineColor(2);
0228       eff->SetLineWidth(2);
0229       eff->SetMarkerStyle(20);
0230       eff->SetMarkerSize(0.8);
0231       eff->SetStats(kFALSE);
0232       effChi2 = (TProfile*)eff->Clone("efficiencyChi2");
0233       eff->SetYTitle("Efficiency");
0234       eff->SetOption("PE");
0235       eff->GetYaxis()->SetRangeUser(-0.001, 1.001);
0236       effChi2->SetYTitle("Chi^2/NDF");
0237       effChi2->SetOption("HIST");
0238 
0239       //create the 1D mass distribution container histograms
0240       TH1D* all1D = (massAxis->GetXbins()->GetSize() == 0)
0241                         ? new TH1D("all1D", "all1D", massAxis->GetNbins(), massAxis->GetXmin(), massAxis->GetXmax())
0242                         : new TH1D("all1D", "all1D", massAxis->GetNbins(), massAxis->GetXbins()->GetArray());
0243       auto* pass1D = (TH1D*)all1D->Clone("pass1D");
0244 
0245       //for each parameter bin fit the mass distributions
0246       for (int par1 = 1; par1 <= par1Axis->GetNbins(); par1++) {
0247         for (int mass = 1; mass <= massAxis->GetNbins(); mass++) {
0248           int index = par1 * par1C + mass * massC;
0249           all1D->SetBinContent(mass, all->GetBinContent(index));
0250           pass1D->SetBinContent(mass, pass->GetBinContent(index));
0251         }
0252         fit(pass1D, all1D);
0253         int index = par1;
0254         eff->SetBinContent(index, getEfficiency());
0255         eff->SetBinEntries(index, 1);
0256         eff->SetBinError(index, sqrt(getEfficiency() * getEfficiency() + getEfficiencyError() * getEfficiencyError()));
0257         effChi2->SetBinContent(index, getChi2());
0258         effChi2->SetBinEntries(index, 1);
0259         if (plotName != "") {
0260           savePlot(TString::Format("%s_%d", plotName.Data(), par1));
0261         }
0262       }
0263       delete all1D;
0264       delete pass1D;
0265       return "";  //OK
0266     }
0267   };
0268 
0269   //concrete fitter: Gaussian signal plus linear background
0270   class GaussianPlusLinearFitter : public AbstractFitter {
0271   protected:
0272     RooGaussian gaussian;
0273     RooRealVar slopeFail;
0274     RooChebychev linearFail;
0275     RooRealVar slopePass;
0276     RooChebychev linearPass;
0277     RooAddPdf pdfFail;
0278     RooAddPdf pdfPass;
0279 
0280   public:
0281     GaussianPlusLinearFitter(bool verbose = false)
0282         : AbstractFitter(verbose),
0283           gaussian("gaussian", "gaussian", mass, mean, sigma),
0284           slopeFail("slopeFail", "slopeFail", 0., -1., 1.),
0285           linearFail("linearFail", "linearFail", mass, slopeFail),
0286           slopePass("slopePass", "slopePass", 0., -1., 1.),
0287           linearPass("linearPass", "linearPass", mass, slopePass),
0288           pdfFail("pdfFail", "pdfFail", RooArgList(gaussian, linearFail), RooArgList(nSignalFail, nBackgroundFail)),
0289           pdfPass("pdfPass", "pdfPass", RooArgList(gaussian, linearPass), RooArgList(nSignalPass, nBackgroundPass)) {
0290       simPdf.addPdf(pdfFail, "fail");
0291       simPdf.addPdf(pdfPass, "pass");
0292     };
0293     ~GaussianPlusLinearFitter() override = default;
0294     ;
0295     void fit(TH1* pass, TH1* all) override {
0296       using namespace RooFit;
0297       all->Add(pass, -1);
0298       TH1*& fail = all;
0299       if (!data)
0300         delete data;
0301       data = new RooDataHist("data", "data", mass, Index(category), Import("fail", *fail), Import("pass", *pass));
0302       if (pass->Integral() + fail->Integral() < 5) {
0303         efficiency.setVal(0.5);
0304         efficiency.setError(0.5);
0305         chi2 = 0;
0306         return;
0307       }
0308       mean.setVal(expectedMean);
0309       sigma.setVal(expectedSigma);
0310       efficiency.setVal(pass->Integral() / (pass->Integral() + fail->Integral()));
0311       nSignalAll.setVal(0.5 * (fail->Integral() + pass->Integral()));
0312       nBackgroundFail.setVal(0.5 * fail->Integral());
0313       nBackgroundPass.setVal(0.5 * pass->Integral());
0314       slopeFail.setVal(0.);
0315       slopePass.setVal(0.);
0316       if (verbose) {
0317         simPdf.fitTo(*data);
0318       } else {
0319         simPdf.fitTo(*data, Verbose(kFALSE), PrintLevel(-1), Warnings(kFALSE), PrintEvalErrors(-1));
0320       }
0321       RooDataHist dataFail("fail", "fail", mass, fail);
0322       RooDataHist dataPass("pass", "pass", mass, pass);
0323       using AbsRealPtr = std::unique_ptr<RooAbsReal>;
0324       const double chi2Fail = AbsRealPtr(pdfFail.createChi2(dataFail, DataError(RooAbsData::Poisson)))->getVal();
0325       const double chi2Pass = AbsRealPtr(pdfPass.createChi2(dataPass, DataError(RooAbsData::Poisson)))->getVal();
0326       chi2 = (chi2Fail + chi2Pass) / (2 * pass->GetNbinsX() - 8);
0327       if (chi2 > 3) {
0328         efficiency.setVal(0.5);
0329         efficiency.setError(0.5);
0330       }
0331     }
0332   };
0333 
0334   //concrete fitter: voigtian signal plus exponential background
0335   class VoigtianPlusExponentialFitter : public AbstractFitter {
0336   protected:
0337     RooRealVar width;
0338     RooVoigtian voigtian;
0339     RooRealVar slopeFail;
0340     RooExponential exponentialFail;
0341     RooRealVar slopePass;
0342     RooExponential exponentialPass;
0343     RooAddPdf pdfFail;
0344     RooAddPdf pdfPass;
0345 
0346   public:
0347     VoigtianPlusExponentialFitter(bool verbose = false)
0348         : AbstractFitter(verbose),
0349           width("width", "width", 2.5, "GeV"),
0350           voigtian("voigtian", "voigtian", mass, mean, width, sigma),
0351           slopeFail("slopeFail", "slopeFail", 0., -1., 0.),
0352           exponentialFail("linearFail", "linearFail", mass, slopeFail),
0353           slopePass("slopePass", "slopePass", 0., -1., 0.),
0354           exponentialPass("linearPass", "linearPass", mass, slopePass),
0355           pdfFail(
0356               "pdfFail", "pdfFail", RooArgList(voigtian, exponentialFail), RooArgList(nSignalFail, nBackgroundFail)),
0357           pdfPass(
0358               "pdfPass", "pdfPass", RooArgList(voigtian, exponentialPass), RooArgList(nSignalPass, nBackgroundPass)) {
0359       width.setConstant(kTRUE);
0360       simPdf.addPdf(pdfFail, "fail");
0361       simPdf.addPdf(pdfPass, "pass");
0362     };
0363     ~VoigtianPlusExponentialFitter() override = default;
0364     ;
0365     void setup(double expectedMean_, double massLow, double massHigh, double expectedSigma_, double width_) {
0366       expectedMean = expectedMean_;
0367       expectedSigma = expectedSigma_;
0368       mass.setRange(massLow, massHigh);
0369       mean.setRange(massLow, massHigh);
0370       width.setVal(width_);
0371     }
0372     void fit(TH1* pass, TH1* all) override {
0373       using namespace RooFit;
0374       all->Add(pass, -1);
0375       TH1*& fail = all;
0376       if (!data)
0377         delete data;
0378       data = new RooDataHist("data", "data", mass, Index(category), Import("fail", *fail), Import("pass", *pass));
0379       if (pass->Integral() + fail->Integral() < 5) {
0380         efficiency.setVal(0.5);
0381         efficiency.setError(0.5);
0382         chi2 = 0;
0383         return;
0384       }
0385       mean.setVal(expectedMean);
0386       sigma.setVal(expectedSigma);
0387       efficiency.setVal(pass->Integral() / (pass->Integral() + fail->Integral()));
0388       nSignalAll.setVal(0.5 * (fail->Integral() + pass->Integral()));
0389       nBackgroundFail.setVal(0.5 * fail->Integral());
0390       nBackgroundPass.setVal(0.5 * pass->Integral());
0391       slopeFail.setVal(0.);
0392       slopePass.setVal(0.);
0393       if (verbose) {
0394         simPdf.fitTo(*data);
0395       } else {
0396         simPdf.fitTo(*data, Verbose(kFALSE), PrintLevel(-1), Warnings(kFALSE), PrintEvalErrors(-1));
0397       }
0398       RooDataHist dataFail("fail", "fail", mass, fail);
0399       RooDataHist dataPass("pass", "pass", mass, pass);
0400       using AbsRealPtr = std::unique_ptr<RooAbsReal>;
0401       const double chi2Fail = AbsRealPtr(pdfFail.createChi2(dataFail, DataError(RooAbsData::Poisson)))->getVal();
0402       const double chi2Pass = AbsRealPtr(pdfPass.createChi2(dataPass, DataError(RooAbsData::Poisson)))->getVal();
0403       chi2 = (chi2Fail + chi2Pass) / (2 * all->GetNbinsX() - 8);
0404       if (chi2 > 3) {
0405         efficiency.setVal(0.5);
0406         efficiency.setError(0.5);
0407       }
0408     }
0409   };
0410 
0411 }  //namespace dqmTnP