Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-03-17 10:58:34

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