MpdFlowScripts/get-dca/get_dca.c

#include <Rtypes.h>
#include <TChain.h>
#include <TClonesArray.h>
#include <TStopwatch.h>

#include <TROOT.h>
#include <TSystem.h>
#include <TMath.h>
// #include <TTree.h>
#include <TFile.h>
#include <TH1.h>
#include <TF1.h>

#include <FairMCEventHeader.h>
#include <FairMCTrack.h>
#include <MpdEvent.h>
#include <MpdKalmanTrack.h>
#include <MpdVertex.h>
#include <MpdTrack.h>
#include <MpdHelix.h>

#include <iostream>

#include "get_fit.c"
#include "MakeFitDCA.c"

using namespace std;

int GetPtBin(Float_t pt);
int GetEtaBin(Float_t eta);

void get_fit(TString inFileName, TString outFileName);

// const float ptBins[] = {0.,0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.5, 3.};
// const int NptBins = 12;

// const float etaBins[] = {-1.5,-1.2,-1.,-0.8,-0.6,-0.4,-0.2,0.,0.2,0.4,0.6,0.8,1.,1.2,1.5};
// const int NetaBins = 14;

// const int Ndim = 3;


const Double_t F_CUR0 = 0.3 * 0.01 * 5 / 10; // 5kG

MpdHelix MakeHelix(const MpdKalmanTrack *tr)
{
  Double_t r = tr->GetPosNew();
  Double_t phi = tr->GetParam(0) / r;
  Double_t x = r * TMath::Cos(phi);
  Double_t y = r * TMath::Sin(phi);
  Double_t dip = tr->GetParam(3);
  Double_t cur = F_CUR0 * TMath::Abs (tr->GetParam(4));
  TVector3 o(x, y, tr->GetParam(1));
  Int_t h = (Int_t) TMath::Sign(1.1,tr->GetParam(4));
  MpdHelix helix(cur, dip, tr->GetParam(2)-TMath::PiOver2()*h, o, h);
  return helix;
}

void get_dca(TString inFileName , TString outFileName)
{
  if (inFileName.IsNull() || outFileName.IsNull())
  {
    std::cout << "Provide an input/output information!" << std::endl;
    return;
  }

  TStopwatch timer;
  timer.Start();

  // Init output histograms
  TH1F* h_dca[Ndim][NptBins][NetaBins];

  char name[200];
  char title[200];

  for (int dim = 0; dim < Ndim; ++dim)
  {
    for (int ptbin = 0; ptbin < NptBins; ++ptbin)
    {
      for (int etabin = 0; etabin < NetaBins; ++etabin)
      {
        sprintf(name,"h_dca_%i_%i_%i",dim,ptbin,etabin);
        sprintf(title,"DCA distribution for %.2f < p_{t} < %.2f and %.2f < #eta < %.2f", ptBins[ptbin], ptBins[ptbin+1],etaBins[etabin],etaBins[etabin+1]);
        h_dca[dim][ptbin][etabin] = new TH1F(name,title,4000,-50.,50);
      }
    }
  }

  TChain *dstTree = new TChain("mpdsim");
  dstTree->Add(inFileName.Data());

  TFile  *outFile = new TFile(outFileName.Data(),"RECREATE");

  MpdEvent *MPDEvent = nullptr;
  dstTree->SetBranchAddress("MPDEvent.", &MPDEvent);
  // TClonesArray *MpdGlobalTracks = nullptr;
  // TClonesArray *mpdKalmanTracks = nullptr; //(TClonesArray*) inFile->FindObjectAny("TpcKalmanTrack");
  // dstTree->SetBranchAddress("TpcKalmanTrack", &mpdKalmanTracks);
  TClonesArray *vertexes = nullptr; //(TClonesArray*) inFile->FindObjectAny("Vertex");
  dstTree->SetBranchAddress("Vertex", &vertexes);

  Int_t n_events = dstTree->GetEntries();
  //Int_t n_events = 50;
  for (int event = 0; event < n_events; ++event)
  {
    cout << "EVENT N "<< event <<endl;
    dstTree->GetEntry(event);
    // MpdGlobalTracks = MPDEvent->GetGlobalTracks();

    Int_t Ntracks = MPDEvent->GetGlobalTracks()->GetEntriesFast();

    // MpdVertex *vertex = (MpdVertex*) vertexes->First();
    // TVector3 primaryVertex;
    // vertex->Position(primaryVertex);
    for (int track = 0; track < Ntracks; ++track)
    {
      // MpdKalmanTrack *kalmanTrack = (MpdKalmanTrack*) mpdKalmanTracks->UncheckedAt(track);
      // DCA stuff
      // MpdHelix helix = MakeHelix(kalmanTrack);
      // Double_t pathLength = helix.pathLength(primaryVertex);
      // TVector3 pca;
      // pca = helix.at(pathLength);
      // pca -= primaryVertex;

      // Double_t dcaX = pca.X();
      // Double_t dcaY = pca.Y();
      // Double_t dcaZ = pca.Z();

      MpdTrack *mpdtrack = (MpdTrack*) MPDEvent->GetGlobalTracks()->UncheckedAt(track); 
      Int_t pt_bin = GetPtBin(TMath::Abs(mpdtrack->GetPt()));
      Int_t eta_bin = GetEtaBin(mpdtrack->GetEta());
      if ( (eta_bin == -1) || (pt_bin == -1 )) continue;
      h_dca[0][pt_bin][eta_bin]->Fill(mpdtrack->GetDCAX());
      h_dca[1][pt_bin][eta_bin]->Fill(mpdtrack->GetDCAY());
      h_dca[2][pt_bin][eta_bin]->Fill(mpdtrack->GetDCAZ());
    }
  }

  outFile->cd();
  for (int dim = 0; dim < Ndim; ++dim)
  {
    for (int ptbin = 0; ptbin < NptBins; ++ptbin)
    {
      for (int etabin = 0; etabin < NetaBins; ++etabin)
      {
        std::cout << h_dca[dim][ptbin][etabin]->GetName() << " ("
          << h_dca[dim][ptbin][etabin]->GetEntries() << " entries)"
          << std::endl;
        h_dca[dim][ptbin][etabin]->Write();
      }
    }
  }
  std::cout << "Closing output file." << std::endl;
  outFile->Close();
}

int GetPtBin(Float_t pt)
{
  int pt_bin = -1;
  for (int i = 0; i < NptBins; ++i)
    if ((pt > ptBins[i]) && (pt <= ptBins[i + 1])) 
      pt_bin = i;
  return pt_bin;
}

int GetEtaBin(Float_t eta)
{
  int eta_bin = -1;
  for (int i = 0; i < NetaBins; ++i)
    if ((eta > etaBins[i]) && (eta <= etaBins[i + 1])) 
      eta_bin = i;
  return eta_bin;
}

int main(int argc, char** argv)
{
  TString iFileName, oFileName;
  Bool_t isFirstFit, isSecondFit;

  isFirstFit = kFALSE; isSecondFit = kFALSE;

  if (argc < 3){
    std::cerr << "./get_dca -i inputfile -o outputfile [--firstfit --secondfit]" << std::endl;
    return 1;
  }

  for (int i=1; i<argc; i++){
    if(std::string(argv[i]) != "-i" &&
       std::string(argv[i]) != "-o" &&
       std::string(argv[i]) != "--firstfit" &&
       std::string(argv[i]) != "--secondfit"){
      std::cerr << "\n[ERROR]: Unknown parameter " << i << ": " <<  argv[i] << std::endl;
      return 1;
    } else {
      if(std::string(argv[i]) == "-i" && i!=argc-1) {
        iFileName = argv[++i];
        continue;
      }
      if(std::string(argv[i]) == "-i" && i==argc-1) {
        std::cerr << "\n[ERROR]: Input file name was not specified " << std::endl;
        return 1;
      }
      if(std::string(argv[i]) == "-o" && i!=argc-1) {
        oFileName = argv[++i];
        continue;
      }
      if(std::string(argv[i]) == "-o" && i==argc-1){
        std::cerr << "\n[ERROR]: Output file name was not specified " << std::endl;
        return 1;
      }
      if(std::string(argv[i]) == "--firstfit") {
        isFirstFit = kTRUE;
        continue;
      }
      if(std::string(argv[i]) == "--secondfit"){
        isSecondFit = kTRUE;
        continue;
      }
    }
  }
  if (isFirstFit == kTRUE && isSecondFit == kTRUE){
    std::cerr << "\n[ERROR]: Both --firstfit and --secondfit cannot be used at the same time " << std::endl;
    return 1;
  }

  if (isFirstFit == kFALSE && isSecondFit == kFALSE) get_dca(iFileName,oFileName);
  if (isFirstFit) get_fit(iFileName,oFileName);
  if (isSecondFit) MakeFitDCA(iFileName, oFileName);
  return 0;
}