LightFlavorRatios/unbinnedMaxLikelihood_fitter/fitutil.h

0001 #ifndef FITUTIL_H
0002 #define FITUTIL_H
0003
0004 #include <iostream>
0005 #include <memory>
0006 #include <string>
0007 #include <utility>
0008 #include <vector>
0009
0010 #include <TString.h>
0011
0012 #include <RooAbsArg.h>
0013 #include <RooAbsPdf.h>
0014 #include <RooAddPdf.h>
0015 #include <RooArgList.h>
0016 #include <RooArgSet.h>
0017 #include <RooFormulaVar.h>
0018 #include <RooRealVar.h>
0019
0020 #include <RooGaussian.h>
0021 #include <RooVoigtian.h>
0022 #include <RooCBShape.h>
0023
0024 #include <RooPolynomial.h>
0025 #include <RooArgusBG.h>
0026
0027 struct fitparam_config
0028 {
0029     std::string branch;
0030     std::string decaystring;
0031
0032     int nBins = 60;
0033
0034     std::string sigmodel;
0035     double minMass = 0.0, maxMass = 1.0;
0036
0037     double nSig = 0.0;
0038     double nSig_low = 0.0, nSig_high = 1e12;
0039     double nSig_fracTol = 0.35; // fraction of total entries
0040     double nSig_fracTol_low = 0.0, nSig_fracTol_high = 1.0;
0041
0042     double mean = 0.0, mean_low = 0.0, mean_high = 1.0;
0043     double sigma = 0.01, sigma_low = 1e-6, sigma_high = 0.2;
0044     double width = 0.002, width_low = 1e-6, width_high = 0.1;
0045
0046     double alpha1 = 1.0, alpha1_low = 1e-3, alpha1_high = 100.0;
0047     double n1 = 5.0, n1_low = 0.1, n1_high = 200.0;
0048     double alpha2 = 1.0, alpha2_low = 1e-3, alpha2_high = 100.0;
0049     double n2 = 5.0, n2_low = 0.1, n2_high = 200.0;
0050     double frac = 0.5;
0051
0052     std::string bkgmodel;
0053
0054     double nBkg = 0.0;
0055     double nBkg_low = 0.0, nBkg_high = 1e12;
0056     double nBkg_fracTol = 1.0 - nSig_fracTol; // fraction of total entries
0057     double nBkg_fracTol_low = 0.0, nBkg_fracTol_high = 1.0;
0058
0059     // Polynomial coeffs (up to 10)
0060     double p1 = 0.0, p1_low = -100.0, p1_high = 100.0;
0061     double p2 = 0.0, p2_low = -100.0, p2_high = 100.0;
0062     double p3 = 0.0, p3_low = -100.0, p3_high = 100.0;
0063     double p4 = 0.0, p4_low = -100.0, p4_high = 100.0;
0064     double p5 = 0.0, p5_low = -100.0, p5_high = 100.0;
0065     double p6 = 0.0, p6_low = -100.0, p6_high = 100.0;
0066     double p7 = 0.0, p7_low = -100.0, p7_high = 100.0;
0067     double p8 = 0.0, p8_low = -100.0, p8_high = 100.0;
0068     double p9 = 0.0, p9_low = -100.0, p9_high = 100.0;
0069     double p10 = 0.0, p10_low = -100.0, p10_high = 100.0;
0070
0071     // Argus
0072     double k = -10.0, k_low = -200.0, k_high = -1e-3;
0073     double edp = 0.0, edp_low = -10.0, edp_high = 10.0;
0074     double argusShift = 0.0;
0075     double argusfrac = 0.5;
0076
0077     bool useStagedFit = false;
0078     double stagedGapNSigma = 3.0;
0079 };
0080
0081 namespace FitParams
0082 {
0083 std::string branch = "Lambda0_mass";
0084 double minMass = 2.1;
0085 double maxMass = 2.4;
0086
0087 RooRealVar mass("mass", "mass", 0.0, -10.0, 10.0);
0088
0089 RooRealVar nSig("nSig", "nSig", 10000.0, 0.0, 1e12);
0090 RooRealVar nBkg("nBkg", "nBkg", 100000.0, 0.0, 1e12);
0091
0092 RooRealVar mean("mean", "mean", 0.0, -10.0, 10.0);
0093 RooRealVar sigma("sigma", "sigma", 0.01, 1e-7, 1.0);
0094 RooRealVar width("width", "width", 0.002, 1e-6, 0.2);
0095 RooRealVar alpha1("alpha1", "alpha1", 1.0, 1e-3, 100.0);
0096 RooRealVar n1("n1", "n1", 5.0, 0.1, 200.0);
0097 RooRealVar alpha2("alpha2", "alpha2", 1.0, 1e-3, 100.0);
0098 RooRealVar n2("n2", "n2", 5.0, 0.1, 200.0);
0099 RooRealVar frac("frac", "frac", 0.5, 0.0, 1.0);
0100
0101 RooRealVar p1("p1", "coeff #1", 0.0, -100.0, 100.0);
0102 RooRealVar p2("p2", "coeff #2", 0.0, -100.0, 100.0);
0103 RooRealVar p3("p3", "coeff #3", 0.0, -100.0, 100.0);
0104 RooRealVar p4("p4", "coeff #4", 0.0, -100.0, 100.0);
0105 RooRealVar p5("p5", "coeff #5", 0.0, -100.0, 100.0);
0106 RooRealVar p6("p6", "coeff #6", 0.0, -100.0, 100.0);
0107 RooRealVar p7("p7", "coeff #7", 0.0, -100.0, 100.0);
0108 RooRealVar p8("p8", "coeff #8", 0.0, -100.0, 100.0);
0109 RooRealVar p9("p9", "coeff #9", 0.0, -100.0, 100.0);
0110 RooRealVar p10("p10", "coeff #10", 0.0, -100.0, 100.0);
0111
0112 RooRealVar k("k", "k", -10.0, -200.0, -1e-3);
0113 RooRealVar edp("edp", "edp", 0.0, -10.0, 10.0);
0114 double ArgusShift = 0.0;
0115 RooRealVar frac_argus("frac_argus", "Fraction of Argus", 0.5, 0.0, 1.0);
0116
0117 static std::vector<std::unique_ptr<RooAbsArg>> _owned_nodes;
0118
0119 inline void BeginFit() { _owned_nodes.clear(); }
0120
0121 template <typename T, typename... Args> inline T *Own(Args &&...args)
0122 {
0123     auto up = std::make_unique<T>(std::forward<Args>(args)...);
0124     T *ptr = up.get();
0125     _owned_nodes.emplace_back(std::move(up));
0126     return ptr;
0127 }
0128
0129 inline RooArgList PolyCoeffList(int order)
0130 {
0131     RooArgList coeffs;
0132     if (order >= 1)
0133         coeffs.add(p1);
0134     if (order >= 2)
0135         coeffs.add(p2);
0136     if (order >= 3)
0137         coeffs.add(p3);
0138     if (order >= 4)
0139         coeffs.add(p4);
0140     if (order >= 5)
0141         coeffs.add(p5);
0142     if (order >= 6)
0143         coeffs.add(p6);
0144     if (order >= 7)
0145         coeffs.add(p7);
0146     if (order >= 8)
0147         coeffs.add(p8);
0148     if (order >= 9)
0149         coeffs.add(p9);
0150     if (order >= 10)
0151         coeffs.add(p10);
0152     return coeffs;
0153 }
0154
0155 inline int ParsePolyOrder(const std::string &modelType)
0156 {
0157     const std::string prefix = "Polynomial";
0158     if (modelType.rfind(prefix, 0) != 0)
0159         return -1;
0160     const std::string suffix = modelType.substr(prefix.size());
0161     if (suffix.empty())
0162         return -1;
0163     try
0164     {
0165         int n = std::stoi(suffix);
0166         if (n < 1 || n > 10)
0167             return -1;
0168         return n;
0169     }
0170     catch (...)
0171     {
0172         return -1;
0173     }
0174 }
0175 } // namespace FitParams
0176
0177 inline RooAbsPdf *signalModel(const std::string &modelType, RooRealVar &mass)
0178 {
0179     using namespace FitParams;
0180
0181     if (modelType == "Gaussian")
0182     {
0183         return Own<RooGaussian>("signal", "Gaussian", mass, mean, sigma);
0184     }
0185     else if (modelType == "Voigtian")
0186     {
0187         return Own<RooVoigtian>("signal", "Voigtian", mass, mean, sigma, width);
0188     }
0189     else if (modelType == "CrystalBall")
0190     {
0191         return Own<RooCBShape>("signal", "Crystal Ball", mass, mean, sigma, alpha1, n1);
0192     }
0193     else if (modelType == "DoubleCrystalBall")
0194     {
0195         return Own<RooDoubleCB>("signal", "Double Crystal Ball", mass, mean, sigma, alpha1, n1, alpha2, n2, frac);
0196     }
0197
0198     std::cerr << "[fitutil.h] ERROR: Unknown signal model type: " << modelType << " (returning nullptr)\n";
0199     return nullptr;
0200 }
0201
0202 inline RooAbsPdf *backgroundModel(const std::string &modelType, RooRealVar &mass)
0203 {
0204     using namespace FitParams;
0205
0206     const int polyOrder = ParsePolyOrder(modelType);
0207     if (polyOrder >= 1)
0208     {
0209         RooArgList coeffs = PolyCoeffList(polyOrder);
0210         return Own<RooPolynomial>("background", modelType.c_str(), mass, coeffs);
0211     }
0212
0213     if (modelType == "Argus")
0214     {
0215         const TString formula = TString::Format("-%s+2*%.17g", mass.GetName(), ArgusShift);
0216
0217         RooArgList deps(mass);
0218         auto *shifted = Own<RooFormulaVar>("shifted", "shifted", formula.Data(), deps);
0219
0220         auto *argus = Own<RooArgusBG>("Argus", "Argus", *shifted, edp, k);
0221
0222         RooArgList polyCoeffs = PolyCoeffList(8);
0223         auto *poly = Own<RooPolynomial>("Polynomial", "Polynomial", mass, polyCoeffs);
0224
0225         RooArgList pdfs(*argus, *poly);
0226         RooArgList fracs(frac_argus);
0227         return Own<RooAddPdf>("background", "Composite Background", pdfs, fracs);
0228     }
0229
0230     std::cerr << "[fitutil.h] ERROR: Unknown background model type: " << modelType << " (returning nullptr)\n";
0231     return nullptr;
0232 }
0233
0234 #endif // FITUTIL_H