OfflineValidation/macros/analyzeDiMuonBiases.C

0001 #include <TCanvas.h>
0002 #include <TF1.h>
0003 #include <TFile.h>
0004 #include <TH1D.h>
0005 #include <TH1F.h>
0006 #include <TH2F.h>
0007 #include <TLorentzVector.h>
0008 #include <TMath.h>
0009 #include <TProfile.h>
0010 #include <TRegexp.h>
0011 #include <TStopwatch.h>
0012 #include <TString.h>
0013 #include <TTree.h>
0014 #include <iomanip>  // Include the header for setw
0015 #include <map>
0016 #include <string>
0017 #include <iostream>
0018
0019 // analysis types
0020 enum anaKind { sagitta_t = 0, d0_t = 1, dz_t = 2 };
0021
0022 // number of bins for the correction map
0023 unsigned int k_NBINS = 48;
0024
0025 // maximum number of iterations
0026 int k_MAXITERATIONS = 30;
0027
0028 static constexpr const char* const analysis_types[] = {
0029     "sagitta",  //  0 - sagitta bias
0030     "d0",       //  1 - transverse IP bias
0031     "dz",       //  2 - longitudinal IP bias
0032 };
0033
0034 // physics constants
0035 static constexpr double k_muMass = 0.105658;  // GeV
0036 static constexpr double MZ_PDG = 91.1876;     // in GeV
0037
0038 // ansatz from eq. 5 of https://arxiv.org/pdf/2212.07338.pdf
0039 //________________________________________________________________
0040 double pTcorrector(const double& unCorrPt, const double& delta, const float& charge) {
0041   return unCorrPt / (1 - charge * delta * unCorrPt);
0042 }
0043
0044 //________________________________________________________________
0045 float calcDeltaSagitta(float charge, float frac, float avgSagitta, float thisTrackPt) {
0046   return (-charge * 0.5 * frac * ((1 + charge * thisTrackPt * avgSagitta) / thisTrackPt));
0047 }
0048
0049 //________________________________________________________________
0050 std::pair<int, int> findEtaPhiBin(const TH2* h2, const float& eta, const float& phi) {
0051   const auto& etaBin = h2->GetXaxis()->FindBin(eta) - 1;
0052   const auto& phiBin = h2->GetYaxis()->FindBin(phi) - 1;
0053
0054   if (etaBin < 0 || etaBin > int(k_NBINS - 1)) {
0055     std::cout << "etaBin: " << etaBin << " eta: " << eta << std::endl;
0056     return {-1., -1.};
0057   }
0058
0059   if (phiBin < 0 || phiBin > int(k_NBINS - 1)) {
0060     std::cout << "phiBin: " << phiBin << " phi: " << phi << std::endl;
0061     return {-1., -1.};
0062   }
0063   return std::make_pair(etaBin, phiBin);
0064 }
0065
0066 //________________________________________________________________
0067 float updateSagittaMap(TTree* tree,
0068                        std::map<std::pair<int, int>, float>& theMap,
0069                        TH2F*& hSagitta,
0070                        const int iteration,
0071                        float oldMaxCorrection);
0072
0073 //________________________________________________________________
0074 float updateIPMap(
0075     TTree* tree, std::map<std::pair<int, int>, float>& theMap, TH2F*& hIP, const int iteration, anaKind type);
0076
0077 // Enum to specify the formatting type
0078 enum TitleFormatType { MEAN, RMS };
0079
0080 //________________________________________________________________
0081 TString updateTitle(const TString& title, TitleFormatType formatType) {
0082   // Regular expression to match units in square brackets
0083   TRegexp re("\\[.*\\]");
0084   TString units;
0085   TString cleanTitle = title;
0086
0087   // Check if units are present and extract them
0088   if (cleanTitle.Index(re) != kNPOS) {
0089     units = cleanTitle(re);                         // Extract the unit part
0090     cleanTitle(re) = "";                            // Remove the unit from the cleanTitle
0091     cleanTitle = cleanTitle.Strip(TString::kBoth);  // Remove extra spaces
0092   }
0093
0094   // Construct the new title based on the formatType
0095   TString newTitle;
0096   switch (formatType) {
0097     case MEAN:
0098       if (!units.IsNull()) {
0099         newTitle = TString::Format("#LT%s#GT %s", cleanTitle.Data(), units.Data());
0100       } else {
0101         newTitle = TString::Format("#LT%s#GT", cleanTitle.Data());
0102       }
0103       break;
0104     case RMS:
0105       if (!units.IsNull()) {
0106         newTitle = TString::Format("RMS (%s) %s", cleanTitle.Data(), units.Data());
0107       } else {
0108         newTitle = TString::Format("RMS (%s)", cleanTitle.Data());
0109       }
0110       break;
0111   }
0112
0113   // Set the new title to the x-axis of h_prof
0114   return newTitle;
0115 }
0116
0117 //________________________________________________________________
0118 std::pair<TH1F*, TH1F*> makeProfileVsEta(TH2F* h2) {
0119   const auto& nBinsX = h2->GetNbinsX();
0120   const auto& nBinsY = h2->GetNbinsY();
0121
0122   float xmin = h2->GetXaxis()->GetXmin();
0123   float xmax = h2->GetXaxis()->GetXmax();
0124
0125   std::vector<float> average;
0126   std::vector<float> RMS;
0127   std::vector<float> errorOnRMS;
0128
0129   for (int i = 1; i <= nBinsX; i++) {
0130     float avgInY{0.};
0131     for (int j = 1; j <= nBinsY; j++) {
0132       avgInY += h2->GetBinContent(i, j);
0133     }
0134     avgInY /= nBinsY;
0135     average.push_back(avgInY);
0136   }
0137
0138   for (int i = 1; i <= nBinsX; i++) {
0139     float RMSInY{0.};
0140     for (int j = 1; j <= nBinsY; j++) {
0141       RMSInY += std::pow((h2->GetBinContent(i, j) - average[i - 1]), 2);
0142     }
0143     RMSInY /= nBinsY;
0144     RMS.push_back(std::sqrt(RMSInY));
0145   }
0146
0147   TH1F* h_prof = new TH1F("avgInEta", "profile vs #eta", nBinsX, xmin, xmax);
0148   h_prof->GetXaxis()->SetTitle(h2->GetXaxis()->GetTitle());
0149   h_prof->GetYaxis()->SetTitle(updateTitle(h2->GetZaxis()->GetTitle(), TitleFormatType::MEAN));
0150   for (int i = 1; i <= nBinsX; i++) {
0151     //std::cout << i << " = " << average[i - 1] << std::endl;
0152     h_prof->SetBinContent(i, average[i - 1]);
0153     h_prof->SetBinError(i, RMS[i - 1] / nBinsY);
0154   }
0155
0156   TH1F* h_RMS = new TH1F("RMSInEta", "RMS vs #eta", nBinsX, xmin, xmax);
0157   h_RMS->GetXaxis()->SetTitle(h2->GetXaxis()->GetTitle());
0158   h_RMS->GetYaxis()->SetTitle(updateTitle(h2->GetZaxis()->GetTitle(), TitleFormatType::RMS));
0159   for (int i = 1; i <= nBinsX; i++) {
0160     //std::cout << i << " = " << average[i - 1] << std::endl;
0161     h_RMS->SetBinContent(i, RMS[i - 1]);
0162     h_RMS->SetBinError(i, RMS[i - 1] / std::sqrt(2 * (nBinsY - 1)));
0163   }
0164
0165   return std::make_pair(h_prof, h_RMS);
0166 }
0167
0168 //________________________________________________________________
0169 std::pair<TH1F*, TH1F*> makeProfileVsPhi(TH2F* h2) {
0170   const auto& nBinsX = h2->GetNbinsX();
0171   const auto& nBinsY = h2->GetNbinsY();
0172
0173   float ymin = h2->GetYaxis()->GetXmin();
0174   float ymax = h2->GetYaxis()->GetXmax();
0175
0176   std::vector<float> average;
0177   std::vector<float> RMS;
0178
0179   for (int j = 1; j <= nBinsY; j++) {
0180     float avgInX{0.};
0181     for (int i = 1; i <= nBinsX; i++) {
0182       avgInX += h2->GetBinContent(i, j);
0183     }
0184     avgInX /= nBinsX;
0185     average.push_back(avgInX);
0186   }
0187
0188   for (int j = 1; j <= nBinsY; j++) {
0189     float RMSInX{0.};
0190     for (int i = 1; i <= nBinsX; i++) {
0191       RMSInX += std::pow((h2->GetBinContent(i, j) - average[j - 1]), 2);
0192     }
0193     RMSInX /= nBinsX;
0194     RMS.push_back(std::sqrt(RMSInX));
0195   }
0196
0197   TH1F* h_prof = new TH1F("avgInPhi", "profile vs #phi", nBinsY, ymin, ymax);
0198   h_prof->GetXaxis()->SetTitle(h2->GetYaxis()->GetTitle());
0199   h_prof->GetYaxis()->SetTitle(updateTitle(h2->GetZaxis()->GetTitle(), TitleFormatType::MEAN));
0200   for (int i = 1; i <= nBinsY; i++) {
0201     //std::cout << i << " = " << average[i - 1] << std::endl;
0202     h_prof->SetBinContent(i, average[i - 1]);
0203     h_prof->SetBinError(i, RMS[i - 1] / nBinsX);
0204   }
0205
0206   TH1F* h_RMS = new TH1F("RMSInPhi", "RMS vs #phi", nBinsY, ymin, ymax);
0207   h_RMS->GetXaxis()->SetTitle(h2->GetYaxis()->GetTitle());
0208   h_RMS->GetYaxis()->SetTitle(updateTitle(h2->GetZaxis()->GetTitle(), TitleFormatType::RMS));
0209   for (int i = 1; i <= nBinsY; i++) {
0210     //std::cout << i << " = " << average[i - 1] << std::endl;
0211     h_RMS->SetBinContent(i, RMS[i - 1]);
0212     h_RMS->SetBinError(i, RMS[i - 1] / std::sqrt(2 * (nBinsX - 1)));
0213   }
0214
0215   return std::make_pair(h_prof, h_RMS);
0216 }
0217
0218 //________________________________________________________________
0219 void analyzeDiMuonBiases(const char* inputFile, anaKind type) {
0220   TStopwatch timer;
0221   timer.Start();
0222
0223   std::cout << "I AM GOING TO PERFORM THE " << analysis_types[type] << " analysis!" << std::endl;
0224
0225   TFile* file = TFile::Open(inputFile);
0226   //TTree* tree = dynamic_cast<TTree*>(file->Get("myanalysis/tree"));
0227   TTree* tree = dynamic_cast<TTree*>(file->Get("ZtoMMNtuple/tree"));
0228
0229   if (tree) {
0230     tree->ls();
0231   } else {
0232     std::cout << "could not open the input tree from file.\n exiting";
0233     exit(EXIT_FAILURE);
0234   }
0235   // prepare the output file name
0236   std::string outputString(inputFile);
0237   std::string oldSubstring = "ZmmNtuple_";
0238   std::string newSubstring = "histos_";
0239   outputString.replace(outputString.find(oldSubstring), oldSubstring.length(), newSubstring);
0240   std::size_t pos = outputString.find(".root");
0241
0242   // If ".root" is found, inject the given const char*
0243   TString toInject = TString::Format("__%s", analysis_types[type]);
0244   if (pos != std::string::npos) {
0245     outputString.insert(pos, toInject.Data());
0246   }
0247
0248   // one for iteration
0249
0250   TFile* outputFile = TFile::Open(outputString.c_str(), "RECREATE");
0251   TH1F* hMass = new TH1F("hMass", "mass;m_{#mu#mu} [GeV];events", 80, 70., 110.);
0252   TH1F* hPt = new TH1F("hPt", ";muon p_{T} [GeV];events", 100, 0., 200.);
0253   TH1F* hZPt = new TH1F("hZPt", ";di-muon system p_{T} [GeV];events", 100, 0., 200.);
0254
0255   TH1F* hPtPlus = new TH1F("hPtPlus", ";positive muon p_{T} [GeV];events", 100, 0., 200.);
0256   TH1F* hPtMinus = new TH1F("hPtMinus", ";negative muon p_{T} [GeV];events", 100, 0., 200.);
0257
0258   TH1F* hPtCorr = new TH1F("hPtCorr", ";corrected muon p_{T} [GeV];events", 100, 0., 200.);
0259   TH1F* hPtPlusCorr = new TH1F("hPtPlusCorr", ";corrected positive muon p_{T} [GeV];events", 100, 0., 200.);
0260   TH1F* hPtMinusCorr = new TH1F("hPtMinusCorr", ";corrected negative muon p_{T} [GeV];events", 100, 0., 200.);
0261
0262   TH1F* hMassCorr = new TH1F("hMassCorr", "corrected mass;corrected m_{#mu#mu} [GeV];events", 80, 70., 110.);
0263   TH1F* hDeltaMassCorr = new TH1F(
0264       "hDeltaMassCorr",
0265       "#Delta Mass (corrected - corrected) ;#Delta m_{#mu#mu} (m^{corr}_{#mu#mu} - m^{uncorr}_{#mu#mu}) [GeV];events",
0266       100,
0267       -10.,
0268       10.);
0269   TH1F* hDeltaPtCorr =
0270       new TH1F("hDeltaPtCorr", "#Delta p_{T};#Delta p_{T} (p^{corr}_{T}-p^{uncorr}_{T}) [GeV];events", 100, -10., 10.);
0271
0272   TH1F* hEta = new TH1F("hEta", ";muon #eta;events", 100, -3.0, 3.0);
0273
0274   TH1F* hDeltaD0 = new TH1F("hDeltaD0", ";#Deltad_{0} [cm];events", 100, -0.01, 0.01);
0275   TH1F* hDeltaDz = new TH1F("hDeltaDz", ";#Deltad_{z} [cm];events", 100, -0.1, 0.1);
0276
0277   TH1F* hDeltaD0Corr = new TH1F("hDeltaD0Corr", ";corrected #Deltad_{0} [cm];events", 100, -0.01, 0.01);
0278   TH1F* hDeltaDzCorr = new TH1F("hDeltaDzCorr", ";corrected #Deltad_{z} [cm];events", 100, -0.1, 0.1);
0279
0280   TH1F* hAbsDeltaD0 = new TH1F("hAbsDeltaD0", ";|#Deltad_{0}| [cm];events", 100, 0., 0.01);
0281   TH1F* hAbsDeltaDz = new TH1F("hAbsDeltaDz", ";|#Deltad_{z}| [cm];events", 100, 0., 0.1);
0282
0283   TH1F* hAbsDeltaD0Corr = new TH1F("hAbsDeltaD0Corr", ";corrected |#Deltad_{0}| [cm];events", 100, 0., 0.01);
0284   TH1F* hAbsDeltaDzCorr = new TH1F("hAbsDeltaDzCorr", ";corrected |#Deltad_{z}| [cm];events", 100, 0., 0.1);
0285
0286   TH1F* hD0 = new TH1F("hD0", ";d_{0}(PV) [cm];events", 100, -0.01, 0.01);
0287   TH1F* hDz = new TH1F("hDz", ";d_{z}(PV) [cm];events", 100, -0.1, 0.1);
0288
0289   TH1F* hD0Corr = new TH1F("hD0Corr", ";corrected d_{0}(PV) [cm];events", 100, -0.01, 0.01);
0290   TH1F* hDzCorr = new TH1F("hDzCorr", ";corrected d_{z}(PV) [cm];events", 100, -0.1, 0.1);
0291
0292   TH1F* hPhi = new TH1F("hPhi", ";muon #phi;events", 100, -TMath::Pi(), TMath::Pi());
0293   TH1F* hCorrection = new TH1F(
0294       "hCorrection",
0295       TString::Format("largest correction vs iteration;iteration number; #delta_{%s} correction ", analysis_types[type])
0296           .Data(),
0297       k_MAXITERATIONS,
0298       0.,
0299       k_MAXITERATIONS);
0300
0301   TH2F* hSagitta = new TH2F("hSagitta",
0302                             "#delta_{sagitta};muon #eta;muon #phi;#delta_{sagitta} [TeV^{-1}]",
0303                             k_NBINS,
0304                             -2.5,
0305                             2.5,
0306                             k_NBINS,
0307                             -TMath::Pi(),
0308                             TMath::Pi());
0309
0310   const char* IPtype = (type == anaKind::d0_t) ? "d_{0}" : "d_{z}";
0311   TH2F* hIP = new TH2F("hIP",
0312                        TString::Format("#delta_{%s};muon #eta;muon #phi;#delta_{%s} [#mum]", IPtype, IPtype).Data(),
0313                        k_NBINS,
0314                        -2.47,
0315                        2.47,
0316                        k_NBINS,
0317                        -TMath::Pi(),
0318                        TMath::Pi());
0319
0320   std::map<std::pair<int, int>, float> sagittaCorrections;
0321   std::map<std::pair<int, int>, float> IPCorrections;
0322   //sagittaCorrections.reserve(k_NBINS * k_NBINS);
0323
0324   // initialize the map
0325   for (unsigned int i = 0; i < k_NBINS; i++) {
0326     for (unsigned int j = 0; j < k_NBINS; j++) {
0327       const auto& index = std::make_pair(i, j);
0328       //sagittaCorrections[index] = 3.435/10e3 ; // from GEN-SIM
0329       sagittaCorrections[index] = 0.f;
0330       IPCorrections[index] = 0.f;
0331     }
0332   }
0333
0334   float maxCorrection{1.f};
0335   int iteration = 0;
0336   float threshold = (type == sagitta_t) ? 1e-10 : 1e-10;
0337
0338   while ((std::abs(maxCorrection) > threshold) && iteration < k_MAXITERATIONS) {
0339     float oldMaxCorrection = maxCorrection;
0340     if (type == sagitta_t) {
0341       maxCorrection = updateSagittaMap(tree, sagittaCorrections, hSagitta, iteration, oldMaxCorrection);
0342     } else {
0343       maxCorrection = updateIPMap(tree, IPCorrections, hIP, iteration, type);
0344     }
0345
0346     std::cout << "iteration: " << iteration << " maxCorrection: " << maxCorrection << std::endl;
0347     hCorrection->SetBinContent(iteration, maxCorrection);
0348     iteration++;
0349   }
0350
0351   Float_t posTrackEta, negTrackEta, posTrackPhi, negTrackPhi, posTrackPt, negTrackPt;
0352   Float_t posTrackDz, negTrackDz, posTrackD0, negTrackD0;
0353
0354   tree->SetBranchAddress("posTrackEta", &posTrackEta);
0355   tree->SetBranchAddress("negTrackEta", &negTrackEta);
0356   tree->SetBranchAddress("posTrackPhi", &posTrackPhi);
0357   tree->SetBranchAddress("negTrackPhi", &negTrackPhi);
0358   tree->SetBranchAddress("posTrackPt", &posTrackPt);
0359   tree->SetBranchAddress("negTrackPt", &negTrackPt);
0360   tree->SetBranchAddress("posTrackDz", &posTrackDz);
0361   tree->SetBranchAddress("negTrackDz", &negTrackDz);
0362   tree->SetBranchAddress("posTrackD0", &posTrackD0);
0363   tree->SetBranchAddress("negTrackD0", &negTrackD0);
0364
0365   Float_t invMass;
0366
0367   // Fill control histograms
0368   for (Long64_t i = 0; i < tree->GetEntries(); i++) {
0369     tree->GetEntry(i);
0370
0371     TLorentzVector posTrack, negTrack, mother;
0372     posTrack.SetPtEtaPhiM(posTrackPt, posTrackEta, posTrackPhi, k_muMass);  // assume muon mass for tracks
0373     negTrack.SetPtEtaPhiM(negTrackPt, negTrackEta, negTrackPhi, k_muMass);
0374     mother = posTrack + negTrack;
0375
0376     hDeltaD0->Fill(posTrackD0 - negTrackD0);
0377     hDeltaDz->Fill(posTrackDz - negTrackDz);
0378
0379     hAbsDeltaD0->Fill(std::abs(posTrackD0 - negTrackD0));
0380     hAbsDeltaDz->Fill(std::abs(posTrackDz - negTrackDz));
0381
0382     hD0->Fill(posTrackD0);
0383     hD0->Fill(negTrackD0);
0384
0385     hDz->Fill(posTrackDz);
0386     hDz->Fill(negTrackDz);
0387
0388     hPt->Fill(posTrackPt);
0389     hPt->Fill(negTrackPt);
0390
0391     hPtPlus->Fill(posTrackPt);
0392     hPtMinus->Fill(negTrackPt);
0393
0394     hEta->Fill(posTrackEta);
0395     hEta->Fill(negTrackEta);
0396
0397     hPhi->Fill(posTrackPhi);
0398     hPhi->Fill(negTrackPhi);
0399
0400     invMass = mother.M();
0401     hMass->Fill(invMass);
0402     hZPt->Fill(mother.Pt());
0403
0404     const auto& indexPlus = findEtaPhiBin(hSagitta, posTrackEta, posTrackPhi);
0405     float deltaPlus = sagittaCorrections[indexPlus];
0406
0407     const auto& indexMinus = findEtaPhiBin(hSagitta, negTrackEta, negTrackPhi);
0408     float deltaMinus = sagittaCorrections[indexMinus];
0409
0410     float deltaIPplus = IPCorrections[indexPlus];
0411     float deltaIPminus = IPCorrections[indexMinus];
0412
0413     if (type == anaKind::d0_t) {
0414       hDeltaD0Corr->Fill((posTrackD0 + deltaIPplus) - (negTrackD0 + deltaIPminus));
0415       hDeltaDzCorr->Fill(posTrackDz - negTrackDz);
0416
0417       hAbsDeltaD0Corr->Fill(std::abs((posTrackD0 + deltaIPplus) - (negTrackD0 + deltaIPminus)));
0418       hAbsDeltaDzCorr->Fill(std::abs(posTrackDz - negTrackDz));
0419
0420       hD0Corr->Fill(posTrackD0 - deltaIPplus);
0421       hD0Corr->Fill(negTrackD0 - deltaIPminus);
0422
0423       hDzCorr->Fill(posTrackDz);
0424       hDzCorr->Fill(negTrackDz);
0425     } else {
0426       hDeltaD0Corr->Fill(posTrackD0 - negTrackD0);
0427       hDeltaDzCorr->Fill((posTrackDz + deltaIPplus) - (negTrackDz + deltaIPminus));
0428
0429       hAbsDeltaD0Corr->Fill(std::abs(posTrackD0 - negTrackD0));
0430       hAbsDeltaDzCorr->Fill(std::abs((posTrackDz + deltaIPplus) - (negTrackDz + deltaIPminus)));
0431
0432       hD0Corr->Fill(posTrackD0);
0433       hD0Corr->Fill(negTrackD0);
0434
0435       hDzCorr->Fill(posTrackDz + deltaIPplus);
0436       hDzCorr->Fill(negTrackDz + deltaIPminus);
0437     }
0438
0439     TLorentzVector posTrackCorr, negTrackCorr, motherCorr;
0440     //posTrackCorr.SetPtEtaPhiM(posTrackPt/(1+posTrackPt*deltaPlus), posTrackEta, posTrackPhi, k_muMass);  // assume muon mass for tracks
0441     //negTrackCorr.SetPtEtaPhiM(negTrackPt/(1-negTrackPt*deltaMinus), negTrackEta, negTrackPhi, k_muMass);
0442
0443     posTrackCorr.SetPtEtaPhiM(
0444         pTcorrector(posTrackPt, deltaPlus, 1.), posTrackEta, posTrackPhi, k_muMass);  // assume muon mass for tracks
0445     negTrackCorr.SetPtEtaPhiM(pTcorrector(negTrackPt, deltaMinus, -1.), negTrackEta, negTrackPhi, k_muMass);
0446
0447     hPtPlusCorr->Fill(posTrackCorr.Pt());
0448     hPtMinusCorr->Fill(negTrackCorr.Pt());
0449
0450     hDeltaPtCorr->Fill(posTrackCorr.Pt() - posTrackPt);
0451     hDeltaPtCorr->Fill(negTrackCorr.Pt() - negTrackPt);
0452
0453     //std::cout << "original pT: " << posTrackPt << " corrected pT: " << posTrackPt / (1+posTrackPt*deltaPlus)  << std::endl;
0454     //std::cout << "original pT: " << negTrackPt << " corrected pT: " << negTrackPt / (1-negTrackPt*deltaMinus) << std::endl;
0455
0456     motherCorr = posTrackCorr + negTrackCorr;
0457
0458     hMassCorr->Fill(motherCorr.M());
0459     hDeltaMassCorr->Fill(mother.M() - motherCorr.M());
0460     hPtCorr->Fill(posTrackCorr.Pt());
0461     hPtCorr->Fill(negTrackCorr.Pt());
0462   }
0463
0464   //for (unsigned int iteration = 0; iteration < 1; iteration++) {
0465   //float max = updateSagittaMap(tree, sagittaCorrections, hSagitta, 0);  // zero-th iteration
0466   // std::cout << "maximal correction update is: " << max << std::endl;
0467   // }
0468   //updateSagittaMap(tree,sagittaCorrections,hSagitta,1); // first iteration
0469   //updateSagittaMap(tree,sagittaCorrections,hSagitta,2); // second iteration
0470   //updateSagittaMap(tree,sagittaCorrections,hSagitta,3); // third iteration
0471
0472   for (unsigned int i = 0; i < k_NBINS; i++) {
0473     for (unsigned int j = 0; j < k_NBINS; j++) {
0474       const auto& index = std::make_pair(i, j);
0475       if (type == sagitta_t) {
0476         hSagitta->SetBinContent(i + 1, j + 1, sagittaCorrections[index] * 10e3);  // 1/GeV = 1000/TeV
0477       } else {
0478         hIP->SetBinContent(i + 1, j + 1, IPCorrections[index] * 10e4);  // 1cm = 10e4um
0479       }
0480     }
0481   }
0482
0483   // The first argument is the name of the new TProfile,
0484   // the second argument is the first bin to include (1 in this case, to skip underflow),
0485   // the third argument is the last bin to include (-1 in this case, to skip overflow),
0486   // and the fourth argument is an option "s" to compute the profile using weights from the bin contents.
0487
0488   //TH1D* projVsEta = hSagitta->ProjectionX("projVsEta", 1, k_NBINS);
0489   //projVsEta->GetYaxis()->SetTitle("#delta_{sagitta} [TeV^{-1}]");
0490   //TH1D* projVsPhi = hSagitta->ProjectionY("projVsPhi", 1, k_NBINS);
0491   //projVsPhi->GetYaxis()->SetTitle("#delta_{sagitta} [TeV^{-1}]");
0492   //TProfile* profVsEta = hSagitta->ProfileX("_pfx",1,47,"s");
0493   //TProfile* profVsPhi = hSagitta->ProfileY("_pfy",1,47,"s");
0494
0495   //const auto& momVsEta = makeProfileVsEta(hSagitta);
0496   //const auto& momVsPhi = makeProfileVsPhi(hSagitta);
0497
0498   std::pair<TH1F*, TH1F*> momVsEta = std::make_pair(nullptr, nullptr);
0499   std::pair<TH1F*, TH1F*> momVsPhi = std::make_pair(nullptr, nullptr);
0500   if (type == sagitta_t) {
0501     momVsEta = makeProfileVsEta(hSagitta);
0502     momVsPhi = makeProfileVsPhi(hSagitta);
0503   } else {
0504     momVsEta = makeProfileVsEta(hIP);
0505     momVsPhi = makeProfileVsPhi(hIP);
0506   }
0507
0508   hDeltaD0->Write();
0509   hDeltaDz->Write();
0510
0511   hMass->Write();
0512   hPt->Write();
0513   hPtPlus->Write();
0514   hPtMinus->Write();
0515
0516   hZPt->Write();
0517
0518   hEta->Write();
0519   hPhi->Write();
0520
0521   if (type == anaKind::sagitta_t) {
0522     hMassCorr->Write();
0523     hPtCorr->Write();
0524     hPtPlusCorr->Write();
0525     hPtMinusCorr->Write();
0526     hDeltaMassCorr->Write();
0527     hDeltaPtCorr->Write();
0528   } else if (type == anaKind::d0_t) {
0529     hD0Corr->Write();
0530     hDeltaD0Corr->Write();
0531     hAbsDeltaD0Corr->Write();
0532   } else if (type == anaKind::dz_t) {
0533     hDzCorr->Write();
0534     hDeltaDzCorr->Write();
0535     hAbsDeltaDzCorr->Write();
0536   }
0537
0538   if (type == sagitta_t) {
0539     hMassCorr->Write();
0540     hPtCorr->Write();
0541     hPtPlusCorr->Write();
0542     hPtMinusCorr->Write();
0543     hDeltaMassCorr->Write();
0544     hDeltaPtCorr->Write();
0545   }
0546
0547   hCorrection->Write();
0548
0549   if (type == sagitta_t) {
0550     hSagitta->SetOption("colz");
0551     hSagitta->Write();
0552   } else {
0553     hIP->SetOption("colz");
0554     hIP->Write();
0555   }
0556
0557   momVsEta.first->Write();
0558   momVsPhi.first->Write();
0559   momVsEta.second->Write();
0560   momVsPhi.second->Write();
0561
0562   //projVsEta->Write();
0563   //projVsPhi->Write();
0564   //profVsEta->Write();
0565   //profVsPhi->Write();
0566
0567   outputFile->Close();
0568   file->Close();
0569
0570   timer.Stop();
0571   timer.Print();
0572 }
0573
0574 //________________________________________________________________
0575 float fitResiduals(TH1* histogram) {
0576   TF1* gaussianFunc = new TF1("gaussianFunc", "gaus", 0, 100);
0577   // Perform an initial fit to get the mean (μ) and width (σ) of the Gaussian
0578   histogram->Fit(gaussianFunc, "Q");
0579
0580   double prevMean = gaussianFunc->GetParameter(1);   // Mean of the previous iteration
0581   double prevSigma = gaussianFunc->GetParameter(2);  // Sigma of the previous iteration
0582
0583   int maxIterations = 10;    // Set a maximum number of iterations
0584   double tolerance = 0.001;  // Define a tolerance to check for convergence
0585
0586   for (int iteration = 0; iteration < maxIterations; ++iteration) {
0587     // Set the fit range to 1.5 σ of the previous iteration
0588     double lowerLimit = prevMean - 1.5 * prevSigma;
0589     double upperLimit = prevMean + 1.5 * prevSigma;
0590     gaussianFunc->SetRange(lowerLimit, upperLimit);
0591
0592     // Perform the fit within the restricted range
0593     histogram->Fit(gaussianFunc, "Q");
0594
0595     // Get the fit results for this iteration
0596     double currentMean = gaussianFunc->GetParameter(1);
0597     double currentSigma = gaussianFunc->GetParameter(2);
0598
0599     // Check for convergence
0600     if (std::abs(currentMean - prevMean) < tolerance && std::abs(currentSigma - prevSigma) < tolerance) {
0601       //cout << "Converged after " << iteration + 1 << " iterations." << endl;
0602       break;
0603     }
0604
0605     prevMean = currentMean;
0606     prevSigma = currentSigma;
0607   }
0608
0609   return prevMean;
0610 }
0611
0612 //________________________________________________________________
0613 float updateSagittaMap(TTree* tree,
0614                        std::map<std::pair<int, int>, float>& theMap,
0615                        TH2F*& hSagitta,
0616                        const int iteration,
0617                        float oldMaxCorrection) {
0618   std::cout << "calling the updateSagittaMap" << std::endl;
0619
0620   std::map<std::pair<int, int>, float> deltaCorrection;
0621   //deltaCorrection.reserve(k_NBINS * k_NBINS);
0622   std::map<std::pair<int, int>, float> countsPerBin;
0623   //countsPerBin.reserve(k_NBINS * k_NBINS);
0624   std::map<std::pair<int, int>, TH1F*> deltaInBin;
0625
0626   //float maxRange = (iteration == 0) ? 1. : 1./std::pow(2,iteration);
0627   float maxRange = 0.01;
0628
0629   for (unsigned int i = 0; i < k_NBINS; i++) {
0630     for (unsigned int j = 0; j < k_NBINS; j++) {
0631       const auto& index = std::make_pair(i, j);
0632       deltaCorrection[index] = 0.f;
0633       countsPerBin[index] = 0.f;
0634       deltaInBin[index] = new TH1F(
0635           Form("delta_iter_%i_%i_%i", iteration, i, j),
0636           Form("iteration %i : #delta_%i_%i;estimated #delta_{s} -true #delta_{s} [GeV^{-1}];n.muons", iteration, i, j),
0637           100,
0638           -maxRange,
0639           maxRange);
0640       //deltaInBin[index]->GetXaxis()->SetCanExtend(kTRUE);
0641     }
0642   }
0643
0644   Float_t posTrackEta, negTrackEta, posTrackPhi, negTrackPhi, posTrackPt, negTrackPt, invMass;
0645   Float_t genPosMuonEta, genNegMuonEta, genPosMuonPhi, genNegMuonPhi, genPosMuonPt, genNegMuonPt;  // gen info
0646
0647   tree->SetBranchAddress("posTrackEta", &posTrackEta);
0648   tree->SetBranchAddress("negTrackEta", &negTrackEta);
0649   tree->SetBranchAddress("posTrackPhi", &posTrackPhi);
0650   tree->SetBranchAddress("negTrackPhi", &negTrackPhi);
0651   tree->SetBranchAddress("posTrackPt", &posTrackPt);
0652   tree->SetBranchAddress("negTrackPt", &negTrackPt);
0653
0654   // gen info
0655   tree->SetBranchAddress("genPosMuonEta", &genPosMuonEta);
0656   tree->SetBranchAddress("genNegMuonEta", &genNegMuonEta);
0657   tree->SetBranchAddress("genPosMuonPhi", &genPosMuonPhi);
0658   tree->SetBranchAddress("genNegMuonPhi", &genNegMuonPhi);
0659   tree->SetBranchAddress("genPosMuonPt", &genPosMuonPt);
0660   tree->SetBranchAddress("genNegMuonPt", &genNegMuonPt);
0661
0662   for (Long64_t i = 0; i < tree->GetEntries(); i++) {
0663     tree->GetEntry(i);
0664     TLorentzVector posTrack, negTrack, mother;
0665     posTrack.SetPtEtaPhiM(posTrackPt, posTrackEta, posTrackPhi, k_muMass);  // assume muon mass for tracks
0666     negTrack.SetPtEtaPhiM(negTrackPt, negTrackEta, negTrackPhi, k_muMass);
0667     mother = posTrack + negTrack;
0668     invMass = mother.M();
0669
0670     //std::cout << "invmass:" << invMass << std::endl;
0671
0672     //now deal with positive muon
0673     const auto& indexPlus = findEtaPhiBin(hSagitta, posTrackEta, posTrackPhi);
0674     float deltaSagittaPlus = theMap[indexPlus];
0675
0676     //now deal with negative muon
0677     const auto& indexMinus = findEtaPhiBin(hSagitta, negTrackEta, negTrackPhi);
0678     float deltaSagittaMinus = theMap[indexMinus];
0679
0680     if (indexPlus == std::make_pair(-1, -1) || indexMinus == std::make_pair(-1, -1)) {
0681       continue;
0682     }
0683
0684     /*
0685     float deltaMass{0.f};
0686     if (iteration == 0) {
0687       deltaMass = invMass * invMass - (MZ_PDG * MZ_PDG);
0688     } else {
0689       deltaMass = MZ_PDG * (posTrackPt * deltaSagittaPlus - negTrackPt * deltaSagittaMinus);
0690     }
0691
0692     float deltaDeltaSagittaPlus =
0693         -(deltaMass / (2 * MZ_PDG)) * ((1 + posTrackPt * deltaSagittaPlus) / posTrackPt);
0694     float deltaDeltaSagittaMinus =
0695         (deltaMass / (2 * MZ_PDG)) * ((1 - negTrackPt * deltaSagittaMinus) / negTrackPt);
0696
0697     */
0698
0699     float frac;
0700     if (iteration != 0) {
0701       frac = (posTrackPt * deltaSagittaPlus - negTrackPt * deltaSagittaMinus);
0702     } else {
0703       frac = (invMass - MZ_PDG) / MZ_PDG;
0704     }
0705
0706     /*
0707     TLorentzVector posTrackCorr, negTrackCorr, motherCorr;
0708     posTrackCorr.SetPtEtaPhiM(pTcorrector(posTrackPt,deltaSagittaPlus,1.), posTrackEta, posTrackPhi, k_muMass);  // assume muon mass for tracks
0709     negTrackCorr.SetPtEtaPhiM(pTcorrector(negTrackPt,deltaSagittaMinus,-1.), negTrackEta, negTrackPhi, k_muMass);
0710     motherCorr = posTrackCorr + negTrackCorr;
0711
0712     float frac = (motherCorr.M() - MZ_PDG) / MZ_PDG;
0713     */
0714
0715     float deltaDeltaSagittaPlus = calcDeltaSagitta(1., frac, deltaSagittaPlus, posTrackPt);
0716     float deltaDeltaSagittaMinus = calcDeltaSagitta(-1., frac, deltaSagittaMinus, negTrackPt);
0717
0718     // inverting equation 5 of https://arxiv.org/pdf/2212.07338.pdf
0719     double trueDeltaSagittaPlus = (genPosMuonPt - posTrackPt) / (genPosMuonPt * posTrackPt);
0720     double trueDeltaSagittaMinus = (negTrackPt - genNegMuonPt) / (genNegMuonPt * negTrackPt);
0721
0722     //std::cout << "deltaMass: " << std::setw(10) << frac
0723     //<< " DeltaDeltaSag(+)-true: " << (deltaSagittaPlus + deltaDeltaSagittaPlus)   - trueDeltaSagittaPlus
0724     //        << " DeltaDeltaSag(-)-true: " << (deltaSagittaMinus + deltaDeltaSagittaMinus) - trueDeltaSagittaMinus << std::endl;
0725
0726     deltaCorrection[indexPlus] += deltaDeltaSagittaPlus;
0727     deltaCorrection[indexMinus] += deltaDeltaSagittaMinus;
0728
0729     deltaInBin[indexPlus]->Fill((deltaSagittaPlus + deltaDeltaSagittaPlus) - trueDeltaSagittaPlus);
0730     deltaInBin[indexMinus]->Fill((deltaSagittaMinus + deltaDeltaSagittaMinus) - trueDeltaSagittaMinus);
0731
0732     countsPerBin[indexPlus] += 1.;
0733     countsPerBin[indexMinus] += 1.;
0734   }
0735
0736   //TFile* iterFile = TFile::Open(Form("iteration_%i.root",iteration), "RECREATE");
0737
0738   for (unsigned int i = 0; i < k_NBINS; i++) {
0739     for (unsigned int j = 0; j < k_NBINS; j++) {
0740       const auto& index = std::make_pair(i, j);
0741       //std::cout << index.first << ", " << index.second << " value: " << deltaCorrection[index] << " / " << countsPerBin[index];
0742       deltaCorrection[index] /= countsPerBin[index];
0743       //deltaCorrection[index] = fitResiduals(deltaInBin[index]);
0744       deltaInBin[index]->Write();
0745       //std::cout << " =  " << deltaCorrection[index] << std::endl;
0746     }
0747   }
0748
0749   //iterFile->Close();
0750
0751   std::cout << " ================================ iteration: " << iteration << std::endl;
0752   for (unsigned int i = 0; i < k_NBINS; i++) {
0753     for (unsigned int j = 0; j < k_NBINS; j++) {
0754       const auto& index = std::make_pair(i, j);
0755       //std::cout << i << ", " << j << " initial: " << theMap[index] << " correction: " << deltaCorrection[index]
0756       //          << std::endl;
0757       theMap[index] += deltaCorrection[index];
0758     }
0759   }
0760
0761   // find the largest element of the correction of this iteration
0762   auto maxIter = std::max_element(deltaCorrection.begin(), deltaCorrection.end(), [](const auto& a, const auto& b) {
0763     return std::abs(a.second) < std::abs(b.second);
0764   });
0765
0766   return maxIter->second;
0767 }
0768
0769 //________________________________________________________________
0770 float updateIPMap(
0771     TTree* tree, std::map<std::pair<int, int>, float>& theMap, TH2F*& hIP, const int iteration, anaKind type) {
0772   std::cout << "calling the updateIPMap" << std::endl;
0773
0774   std::map<std::pair<int, int>, float> deltaCorrection;
0775   //deltaCorrection.reserve(k_NBINS * k_NBINS);
0776   std::map<std::pair<int, int>, float> countsPerBin;
0777   //countsPerBin.reserve(k_NBINS * k_NBINS);
0778
0779   for (unsigned int i = 0; i < k_NBINS; i++) {
0780     for (unsigned int j = 0; j < k_NBINS; j++) {
0781       const auto& index = std::make_pair(i, j);
0782       deltaCorrection[index] = 0.f;
0783       countsPerBin[index] = 0.f;
0784     }
0785   }
0786
0787   float posTrackEta, negTrackEta, posTrackPhi, negTrackPhi, posTrackIP, negTrackIP;
0788   tree->SetBranchAddress("posTrackEta", &posTrackEta);
0789   tree->SetBranchAddress("negTrackEta", &negTrackEta);
0790   tree->SetBranchAddress("posTrackPhi", &posTrackPhi);
0791   tree->SetBranchAddress("negTrackPhi", &negTrackPhi);
0792   if (type == d0_t) {
0793     tree->SetBranchAddress("posTrackD0", &posTrackIP);
0794     tree->SetBranchAddress("negTrackD0", &negTrackIP);
0795   } else if (type == dz_t) {
0796     tree->SetBranchAddress("posTrackDz", &posTrackIP);
0797     tree->SetBranchAddress("negTrackDz", &negTrackIP);
0798   } else {
0799     std::cout << "error, you have been calling this script with the wrong settings!";
0800     exit(EXIT_FAILURE);
0801   }
0802
0803   for (Long64_t i = 0; i < tree->GetEntries(); i++) {
0804     tree->GetEntry(i);
0805
0806     // deal with the positive muon
0807     const auto& indexPlus = findEtaPhiBin(hIP, posTrackEta, posTrackPhi);
0808     float deltaIPPlus = theMap[indexPlus];
0809
0810     //now deal with negative muon
0811     const auto& indexMinus = findEtaPhiBin(hIP, negTrackEta, negTrackPhi);
0812     float deltaIPMinus = theMap[indexMinus];
0813
0814     float IPdistance = std::abs(((posTrackIP + deltaIPPlus) - (negTrackIP + deltaIPMinus)));
0815
0816     //this converges
0817     float deltaDeltaIPPlus =
0818         (posTrackIP + deltaIPPlus) < (negTrackIP + deltaIPMinus) ? IPdistance / 2. : -IPdistance / 2.;
0819     float deltaDeltaIPMinus =
0820         (posTrackIP + deltaIPPlus) < (negTrackIP + deltaIPMinus) ? -IPdistance / 2. : IPdistance / 2.;
0821
0822     //std::cout << "deltaMass: " << deltaMass << " DeltaDeltaSagPlus: " <<   deltaDeltaIPPlus << " DeltaDeltaSagMinus: " << deltaDeltaIPMinus << std::endl;
0823
0824     deltaCorrection[indexPlus] += deltaDeltaIPPlus;
0825     deltaCorrection[indexMinus] += deltaDeltaIPMinus;
0826
0827     countsPerBin[indexPlus] += 1.;
0828     countsPerBin[indexMinus] += 1.;
0829   }
0830
0831   for (unsigned int i = 0; i < k_NBINS; i++) {
0832     for (unsigned int j = 0; j < k_NBINS; j++) {
0833       const auto& index = std::make_pair(i, j);
0834       //std::cout << index.first << ", " << index.second << " value: " << deltaCorrection[index] << " / " << countsPerBin[index];
0835       deltaCorrection[index] /= countsPerBin[index];
0836       //std::cout << " =  " << deltaCorrection[index] << std::endl;
0837     }
0838   }
0839
0840   std::cout << " ================================ iteration: " << iteration << std::endl;
0841   for (unsigned int i = 0; i < k_NBINS; i++) {
0842     for (unsigned int j = 0; j < k_NBINS; j++) {
0843       const auto& index = std::make_pair(i, j);
0844       //std::cout << i << ", " << j << " initial: " << theMap[index] << " correction: " << deltaCorrection[index]
0845       //          << std::endl;
0846       theMap[index] += deltaCorrection[index];
0847     }
0848   }
0849
0850   // find the largest element of the correction of this iteration
0851   auto maxIter = std::max_element(deltaCorrection.begin(), deltaCorrection.end(), [](const auto& a, const auto& b) {
0852     return std::abs(a.second) < std::abs(b.second);
0853   });
0854
0855   return maxIter->second;
0856 }