RecoilCorrector.h

#ifndef HTT_RecoilCorrector_h
#define HTT_RecoilCorrector_h

#include <TFile.h>
#include <TH1.h>
#include <TF1.h>
#include <TString.h>
#include <TRandom.h>
#include <TMath.h>
#include <assert.h>

class RecoilCorrector {
  
 public:
 RecoilCorrector(TString fileName) {

  TString cmsswBase = TString( getenv ("CMSSW_BASE") );
  TString baseDir = cmsswBase + "/src";

  _fileName = baseDir+"/"+fileName;
  TFile * file = new TFile(_fileName);
  if (file->IsZombie()) {
    std::cout << "file " << _fileName << " is not found...   quitting " << std::endl;
    exit(-1);
  }

  TH1D * projH = (TH1D*)file->Get("projH");
  if (projH==NULL) {
    std::cout << "File should contain histogram with the name projH " << std::endl;
    std::cout << "Check content of the file " << _fileName << std::endl;
  }

  TString firstBinStr  = projH->GetXaxis()->GetBinLabel(1);
  TString secondBinStr = projH->GetXaxis()->GetBinLabel(2);

  TString paralZStr = firstBinStr;
  TString perpZStr  = secondBinStr;
  if (firstBinStr.Contains("Perp")) {
    paralZStr = secondBinStr;
    perpZStr  = firstBinStr;
  }
  std::cout << "Parallel component      (U1) : " << paralZStr << std::endl;
  std::cout << "Perpendicular component (U2) : " << perpZStr << std::endl;
  TH1D * ZPtBinsH = (TH1D*)file->Get("ZPtBinsH");
  if (ZPtBinsH==NULL) {
    std::cout << "File should contain histogram with the name ZPtBinsH " << std::endl;
    std::cout << "Check content of the file " << _fileName << std::endl;
    exit(-1);
  }
  int nZPtBins = ZPtBinsH->GetNbinsX();
  float ZPtBins[10];
  TString ZPtStr[10];
  for (int i=0; i<=nZPtBins; ++i) {
    ZPtBins[i] = ZPtBinsH->GetXaxis()->GetBinLowEdge(i+1);
    if (i<nZPtBins)
      ZPtStr[i] = ZPtBinsH->GetXaxis()->GetBinLabel(i+1);
  }

  TH1D * nJetBinsH = (TH1D*)file->Get("nJetBinsH");
  if (nJetBinsH==NULL) {
    std::cout << "File should contain histogram with the name nJetBinsH" << std::endl;
    std::cout << "Check content of the file " << _fileName << std::endl;
    exit(-1);
  }
  int nJetsBins = nJetBinsH->GetNbinsX();
  TString nJetsStr[5];
  for (int i=0; i<nJetsBins; ++i) {
    nJetsStr[i] = nJetBinsH->GetXaxis()->GetBinLabel(i+1);
  }

  InitMEtWeights(file,
                 perpZStr,
                 paralZStr,
                 nZPtBins,
                 ZPtBins,
                 ZPtStr,
                 nJetsBins,
                 nJetsStr);

  _epsrel = 5e-4;
  _epsabs = 5e-4;
  _range = 0.95;

  }


  ~RecoilCorrector(){}

void Correct(float MetPx,
                              float MetPy,
                              float genVPx,
                              float genVPy,
                              float visVPx,
                              float visVPy,
                              int njets,
                              float & MetCorrPx,
                              float & MetCorrPy) {

  float Zpt = TMath::Sqrt(genVPx*genVPx + genVPy*genVPy);
  float U1 = 0;
  float U2 = 0;
  float metU1 = 0;
  float metU2 = 0;
  CalculateU1U2FromMet(MetPx,
                       MetPy,
                       genVPx,
                       genVPy,
                       visVPx,
                       visVPy,
                       U1,
                       U2,
                       metU1,
                       metU2);
  if (Zpt>1000)
    Zpt = 999;

  if (njets>=_nJetsBins)
    njets = _nJetsBins - 1;

  int ZptBin = binNumber(Zpt, _ZPtBins);
  TF1 * metZParalData = _metZParalData[ZptBin][njets];
  TF1 * metZPerpData  = _metZPerpData[ZptBin][njets];

  TF1 * metZParalMC   = _metZParalMC[ZptBin][njets];
  TF1 * metZPerpMC     = _metZPerpMC[ZptBin][njets];

  if (U1>_range*_xminMetZParal[ZptBin][njets]&&U1<_range*_xmaxMetZParal[ZptBin][njets]) {

    int nSumProb = 1;
    double q[1];
    double sumProb[1];

    sumProb[0] = metZParalMC->IntegralOneDim(_xminMetZParalMC[ZptBin][njets],U1,_epsrel,_epsabs,_error);

    if (sumProb[0]<0) {
      sumProb[0] = 1e-5;
    }
    if (sumProb[0]>1) {
      sumProb[0] = 1.0 - 1e-5;
    }


    metZParalData->GetQuantiles(nSumProb,q,sumProb);

    float U1reco = float(q[0]);
    U1 = U1reco;

  }
  else {
    float U1reco = rescale(U1,
                           _meanMetZParalData[ZptBin][njets],
                           _meanMetZParalMC[ZptBin][njets],
                           _rmsMetZParalData[ZptBin][njets],
                           _rmsMetZParalMC[ZptBin][njets]);
    U1 = U1reco;
  }

  if (U2>_range*_xminMetZPerp[ZptBin][njets]&&U2<_range*_xmaxMetZPerp[ZptBin][njets]) {

    int nSumProb = 1;
    double q[1];
    double sumProb[1];

    sumProb[0] = metZPerpMC->IntegralOneDim(_xminMetZPerpMC[ZptBin][njets],U2,_epsrel,_epsabs,_error);

    if (sumProb[0]<0) {
      sumProb[0] = 1e-5;
    }
    if (sumProb[0]>1) {
      sumProb[0] = 1.0 - 1e-5;
    }

    metZPerpData->GetQuantiles(nSumProb,q,sumProb);

    float U2reco = float(q[0]);
    U2 = U2reco;

  }
  else {
    float U2reco = rescale(U2,
                           _meanMetZPerpData[ZptBin][njets],
                           _meanMetZPerpMC[ZptBin][njets],
                           _rmsMetZPerpData[ZptBin][njets],
                           _rmsMetZPerpMC[ZptBin][njets]);
    U2 = U2reco;
  }

  CalculateMetFromU1U2(U1,U2,genVPx,genVPy,visVPx,visVPy,MetCorrPx,MetCorrPy);

}

void CorrectByMeanResolution(float MetPx,
                                              float MetPy,
                                              float genVPx,
                                              float genVPy,
                                              float visVPx,
                                              float visVPy,
                                              int njets,
                                              float & MetCorrPx,
                                              float & MetCorrPy) {

  float Zpt = TMath::Sqrt(genVPx*genVPx + genVPy*genVPy);

  float U1 = 0;
  float U2 = 0;
  float metU1 = 0;
  float metU2 = 0;

  CalculateU1U2FromMet(MetPx,
                       MetPy,
                       genVPx,
                       genVPy,
                       visVPx,
                       visVPy,
                       U1,
                       U2,
                       metU1,
                       metU2);
  if (Zpt>1000)
    Zpt = 999;

  if (njets>=_nJetsBins)
    njets = _nJetsBins - 1;
  int ZptBin = binNumber(Zpt, _ZPtBins);

  U1U2CorrectionsByWidth(U1,
                         U2,
                         ZptBin,
                         njets);

  CalculateMetFromU1U2(U1,U2,genVPx,genVPy,visVPx,visVPy,MetCorrPx,MetCorrPy);

}

  
 private:

  int binNumber(float x, const std::vector<float> bins) const
  {
    for (size_t iB=0; iB<bins.size(); ++iB)
      if (x>=bins[iB]&&x<bins[iB+1])
	return iB;
    return 0;
  }

  int binNumber(float x, int nbins, const float * bins) {

    int binN = 0;

    for (int iB=0; iB<nbins; ++iB) {
      if (x>=bins[iB]&&x<bins[iB+1]) {
	binN = iB;
	break;
      }
    }
    
    return binN;

  }

  TString _fileName;

void InitMEtWeights(TFile * _file,
                                     TString  _perpZStr,
                                     TString  _paralZStr,
                                     int nZPtBins,
                                     float * ZPtBins,
                                     TString * _ZPtStr,
                                     int nJetsBins,
                                     TString * _nJetsStr) {

        std::vector<float> newZPtBins;
        std::vector<std::string> newZPtStr;
        std::vector<std::string> newNJetsStr;

        std::string newPerpZStr  = std::string(_perpZStr);
        std::string newParalZStr = std::string(_paralZStr);

        for (int idx=0; idx<nZPtBins+1; ++idx)
          newZPtBins.push_back(ZPtBins[idx]);
        for (int idx=0; idx<nZPtBins; ++idx )
          newZPtStr.push_back(std::string(_ZPtStr[idx]));

        for (int idx=0; idx<nJetsBins; ++idx)
          newNJetsStr.push_back(std::string(_nJetsStr[idx]));

        InitMEtWeights(_file,
                       newZPtBins,
                       newPerpZStr,
                       newParalZStr,
                       newZPtStr,
                       newNJetsStr);
}


void InitMEtWeights(TFile * _fileMet,
                                     const std::vector<float>& ZPtBins,
                                     const std::string _perpZStr,
                                     const std::string _paralZStr,
                                     const std::vector<std::string>& _ZPtStr,
                                     const std::vector<std::string>& _nJetsStr)
{

  _nZPtBins = ZPtBins.size()-1; // the -1 is on purpose!
  _nJetsBins = _nJetsStr.size();

  _ZPtBins = ZPtBins;

  for (int ZPtBin=0; ZPtBin<_nZPtBins; ++ZPtBin) {
    for (int jetBin=0; jetBin<_nJetsBins; ++jetBin) {

      TString binStrPerpData  = _perpZStr  + "_" + _nJetsStr[jetBin] + _ZPtStr[ZPtBin] + "_data";
      TString binStrParalData = _paralZStr + "_" + _nJetsStr[jetBin] + _ZPtStr[ZPtBin] + "_data";
      TString binStrPerpMC    = _perpZStr  + "_" + _nJetsStr[jetBin] + _ZPtStr[ZPtBin] + "_mc";
      TString binStrParalMC   = _paralZStr + "_" + _nJetsStr[jetBin] + _ZPtStr[ZPtBin] + "_mc";

      _metZParalData[ZPtBin][jetBin] = (TF1*)_fileMet->Get(binStrParalData);
      _metZPerpData[ZPtBin][jetBin]  = (TF1*)_fileMet->Get(binStrPerpData);
      _metZParalMC[ZPtBin][jetBin]   = (TF1*)_fileMet->Get(binStrParalMC);
      _metZPerpMC[ZPtBin][jetBin]    = (TF1*)_fileMet->Get(binStrPerpMC);

      if (_metZParalData[ZPtBin][jetBin]==NULL) {
        std::cout << "Function with name " << binStrParalData
                  << " is not found in file " << _fileName << "... quitting program..." << std::endl;
        exit(-1);

      }
      if (_metZPerpData[ZPtBin][jetBin]==NULL) {
        std::cout << "Function with name " << binStrPerpData
                  << " is not found in file " << _fileName << "... quitting program..." << std::endl;
        exit(-1);

      }

      if (_metZParalMC[ZPtBin][jetBin]==NULL) {
        std::cout << "Function with name " << binStrParalMC
                  << " is not found in file " << _fileName << "... quitting program..." << std::endl;
        exit(-1);

      }
      if (_metZPerpMC[ZPtBin][jetBin]==NULL) {
        std::cout << "Function with name " << binStrPerpMC
                  << " is not found in file " << _fileName << "... quitting program..." << std::endl;
        exit(-1);

      }

      std::cout << _ZPtStr[ZPtBin] << " : " << _nJetsStr[jetBin] << std::endl;

      double xminD,xmaxD;

      _metZParalData[ZPtBin][jetBin]->GetRange(xminD,xmaxD);
      _xminMetZParalData[ZPtBin][jetBin] = float(xminD);
      _xmaxMetZParalData[ZPtBin][jetBin] = float(xmaxD);

      _metZPerpData[ZPtBin][jetBin]->GetRange(xminD,xmaxD);
      _xminMetZPerpData[ZPtBin][jetBin] = float(xminD);
      _xmaxMetZPerpData[ZPtBin][jetBin] = float(xmaxD);

      _metZParalMC[ZPtBin][jetBin]->GetRange(xminD,xmaxD);
      _xminMetZParalMC[ZPtBin][jetBin] = float(xminD);
      _xmaxMetZParalMC[ZPtBin][jetBin] = float(xmaxD);

      _metZPerpMC[ZPtBin][jetBin]->GetRange(xminD,xmaxD);
      _xminMetZPerpMC[ZPtBin][jetBin] = float(xminD);
      _xmaxMetZPerpMC[ZPtBin][jetBin] = float(xmaxD);

      _xminMetZParal[ZPtBin][jetBin] = TMath::Max(_xminMetZParalData[ZPtBin][jetBin],_xminMetZParalMC[ZPtBin][jetBin]);
      _xmaxMetZParal[ZPtBin][jetBin] = TMath::Min(_xmaxMetZParalData[ZPtBin][jetBin],_xmaxMetZParalMC[ZPtBin][jetBin]);

      _xminMetZPerp[ZPtBin][jetBin] = TMath::Max(_xminMetZPerpData[ZPtBin][jetBin],_xminMetZPerpMC[ZPtBin][jetBin]);
      _xmaxMetZPerp[ZPtBin][jetBin] = TMath::Min(_xmaxMetZPerpData[ZPtBin][jetBin],_xmaxMetZPerpMC[ZPtBin][jetBin]);

      _meanMetZParalData[ZPtBin][jetBin] = _metZParalData[ZPtBin][jetBin]->Mean(_xminMetZParalData[ZPtBin][jetBin],_xmaxMetZParalData[ZPtBin][jetBin]);
      _rmsMetZParalData[ZPtBin][jetBin] = TMath::Sqrt(_metZParalData[ZPtBin][jetBin]->CentralMoment(2,_xminMetZParalData[ZPtBin][jetBin],_xmaxMetZParalData[ZPtBin][jetBin]));
      _meanMetZPerpData[ZPtBin][jetBin] = 0;
      _rmsMetZPerpData[ZPtBin][jetBin] = TMath::Sqrt(_metZPerpData[ZPtBin][jetBin]->CentralMoment(2,_xminMetZPerpData[ZPtBin][jetBin],_xmaxMetZPerpData[ZPtBin][jetBin]));

      _meanMetZParalMC[ZPtBin][jetBin] = _metZParalMC[ZPtBin][jetBin]->Mean(_xminMetZParalMC[ZPtBin][jetBin],_xmaxMetZParalMC[ZPtBin][jetBin]);
      _rmsMetZParalMC[ZPtBin][jetBin] = TMath::Sqrt(_metZParalMC[ZPtBin][jetBin]->CentralMoment(2,_xminMetZParalMC[ZPtBin][jetBin],_xmaxMetZParalMC[ZPtBin][jetBin]));
      _meanMetZPerpMC[ZPtBin][jetBin] = 0;
      _rmsMetZPerpMC[ZPtBin][jetBin] = TMath::Sqrt(_metZPerpMC[ZPtBin][jetBin]->CentralMoment(2,_xminMetZPerpMC[ZPtBin][jetBin],_xmaxMetZPerpMC[ZPtBin][jetBin]));


    }
  }

}
void CalculateU1U2FromMet(float metPx,
                                           float metPy,
                                           float genZPx,
                                           float genZPy,
                                           float diLepPx,
                                           float diLepPy,
                                           float & U1,
                                           float & U2,
                                           float & metU1,
                                           float & metU2) {

  float hadRecX = metPx + diLepPx - genZPx;
  float hadRecY = metPy + diLepPy - genZPy;

  float hadRecPt = TMath::Sqrt(hadRecX*hadRecX+hadRecY*hadRecY);

  float phiHadRec = TMath::ATan2(hadRecY,hadRecX);

  float phiDiLep = TMath::ATan2(diLepPy,diLepPx);

  float phiMEt = TMath::ATan2(metPy,metPx);

  float metPt = TMath::Sqrt(metPx*metPx+metPy*metPy);

  float phiZ = TMath::ATan2(genZPy,genZPx);

  float deltaPhiZHadRec  = phiHadRec - phiZ;
  float deltaPhiDiLepMEt = phiMEt - phiDiLep;

  U1 = hadRecPt * TMath::Cos(deltaPhiZHadRec);
  U2 = hadRecPt * TMath::Sin(deltaPhiZHadRec);

  metU1 = metPt * TMath::Cos(deltaPhiDiLepMEt);
  metU2 = metPt * TMath::Sin(deltaPhiDiLepMEt);
}

void CalculateMetFromU1U2(float U1,
                                           float U2,
                                           float genZPx,
                                           float genZPy,
                                           float diLepPx,
                                           float diLepPy,
                                           float & metPx,
                                           float & metPy) {

  float hadRecPt = TMath::Sqrt(U1*U1+U2*U2);

  float deltaPhiZHadRec = TMath::ATan2(U2,U1);

  float phiZ = TMath::ATan2(genZPy,genZPx);

  float phiHadRec = phiZ + deltaPhiZHadRec;

  float hadRecX = hadRecPt*TMath::Cos(phiHadRec);
  float hadRecY = hadRecPt*TMath::Sin(phiHadRec);

  metPx = hadRecX + genZPx - diLepPx;
  metPy = hadRecY + genZPy - diLepPy;
}

void U1U2CorrectionsByWidth(float & U1,
                                             float & U2,
                                             int ZptBin,
                                             int njets) {

  if (njets>=_nJetsBins)
    njets = _nJetsBins - 1;

  float width = U1 - _meanMetZParalMC[ZptBin][njets];
  width *= _rmsMetZParalData[ZptBin][njets]/_rmsMetZParalMC[ZptBin][njets];
  U1 = _meanMetZParalData[ZptBin][njets] + width;

  width = U2;
  width *= _rmsMetZPerpData[ZptBin][njets]/_rmsMetZPerpMC[ZptBin][njets];
  U2 = width;

}

float rescale(float x,
                               float meanData,
                               float meanMC,
                               float resolutionData,
                               float resolutionMC) {

  float width = x - meanMC;
  width *= resolutionData/resolutionMC;
  return meanData + width;

}


float CorrectionsBySampling(float x, TF1 * funcMC, TF1 * funcData) {

  int nSumProb = 1;
  double q[1];
  double sumProb[1];

  double xD = double(x);

  double xminD = 0;
  double xmaxD = 0;

  funcMC->GetRange(xminD,xmaxD);

  float xmin = float(xminD);

  sumProb[0] = funcMC->IntegralOneDim(xmin,xD,_epsrel,_epsabs,_error);

  funcData->GetQuantiles(nSumProb,q,sumProb);

  float output = float(q[0]);

  return output;

}



  //float * _ZPtBins;
  std::vector<float> _ZPtBins;

  Double_t 	_epsrel;
  Double_t 	_epsabs;
  Double_t   	_error;
  float _range;

  int _nZPtBins;
  int _nJetsBins;

  TF1 * _metZParalData[5][3];
  TF1 * _metZPerpData[5][3];
  TF1 * _metZParalMC[5][3];
  TF1 * _metZPerpMC[5][3];

  TH1F * _metZParalDataHist[5][3];
  TH1F * _metZPerpDataHist[5][3];
  TH1F * _metZParalMCHist[5][3];
  TH1F * _metZPerpMCHist[5][3];

  float _meanMetZParalData[5][3];
  float _meanMetZParalMC[5][3];
  float _meanMetZPerpData[5][3];
  float _meanMetZPerpMC[5][3];
  
  float _rmsMetZParalData[5][3];
  float _rmsLeftMetZParalData[5][3];
  float _rmsRightMetZParalData[5][3];
  
  float _rmsMetZParalMC[5][3];
  float _rmsLeftMetZParalMC[5][3];
  float _rmsRightMetZParalMC[5][3];

  float _rmsMetZPerpData[5][3];
  float _rmsLeftMetZPerpData[5][3];
  float _rmsRightMetZPerpData[5][3];
  
  float _rmsMetZPerpMC[5][3];
  float _rmsLeftMetZPerpMC[5][3];
  float _rmsRightMetZPerpMC[5][3];
  
  float _xminMetZPerp[5][3];
  float _xmaxMetZPerp[5][3];

  float _xminMetZPerpData[5][3];
  float _xmaxMetZPerpData[5][3];

  float _xminMetZPerpMC[5][3];
  float _xmaxMetZPerpMC[5][3];

  float _xminMetZParal[5][3];
  float _xmaxMetZParal[5][3];

  float _xminMetZParalData[5][3];
  float _xmaxMetZParalData[5][3];

  float _xminMetZParalMC[5][3];
  float _xmaxMetZParalMC[5][3];

};

#endif