OfflineValidation/macros/analyzeDiMuonBiases_fast.C

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