StFlowMephi/FemtoDstAnalyzer_PID.C

/**
 * \brief Example of how to read a file (list of files) using StFemtoEvent classes
 *
 * RunFemtoDstAnalyzer.C is an example of reading FemtoDst format.
 * One can use either FemtoDst file or a list of femtoDst files (inFile.lis or
 * inFile.list) as an input, and preform physics analysis
 *
 * \author Grigory Nigmatkulov
 * \date May 29, 2018
 */

// This is needed for calling standalone classes
#define _VANILLA_ROOT_

// C++ headers
#include <iostream>
#include <vector>
#include <map>

// ROOT headers
#include "TROOT.h"
#include "TFile.h"
#include "TChain.h"
#include "TVector2.h"
#include "TVector3.h"
#include "TTree.h"
#include "TSystem.h"
#include "TH1.h"
#include "TH2.h"
#include "TMath.h"
#include "TProfile2D.h"
#include "TStopwatch.h"

// FemtoDst headers
#include "StFemtoDstReader.h"
#include "StFemtoDst.h"
#include "StFemtoEvent.h"
#include "StFemtoTrack.h"

// Load libraries (for ROOT_VERSTION_CODE >= 393215)
#if ROOT_VERSION_CODE >= ROOT_VERSION(6, 0, 0)
R__LOAD_LIBRARY(/mnt/pool/rhic/4/parfenovpeter/STAR/build/libStFemtoDst.so)
#endif

#include "Constants.h"

// inFile - is a name of name.FemtoDst.root file or a name
//          of a name.lis(t) files, that contains a list of
//          name1.FemtoDst.root, name2.FemtoDst.root, ... files

//Used functions (see them below)
Bool_t isGoodEvent(StFemtoEvent *const &event);
Bool_t isGoodTrackFlow(StFemtoTrack *const &track, Float_t _energy, TVector3 pVtx);

Double_t GetMeanNsig(StFemtoTrack *const &track, const Int_t PID);
Double_t GetWidthM2(StFemtoTrack *const &track, const Int_t PID);
Double_t GetRotAngle(StFemtoTrack *const &track);
std::pair<Double_t,Double_t> GetNewXY(StFemtoTrack *const &track);
Bool_t isMesonXY(Double_t x, Double_t y);

Double_t GetSigmNewX(StFemtoTrack *const &track, Int_t PID);
Double_t GetSigmNewY(StFemtoTrack *const &track, Int_t PID);
//_________________
void FemtoDstAnalyzer_PID(const Char_t *inFile = "st_physics_12150008_raw_4030001.femtoDst.root",
                                         const Char_t *outFile = "./oPIDTest.root")
{

  std::cout << "Hi! Lets do some physics, Master!" << std::endl;
  TStopwatch timer;
  timer.Start();

#if ROOT_VERSION_CODE < ROOT_VERSION(6, 0, 0)
  gSystem->Load("../libStFemtoDst.so");
#endif

  StFemtoDstReader *femtoReader = new StFemtoDstReader(inFile);
  femtoReader->Init();

  // This is a way if you want to spead up IO
  std::cout << "Explicit read status for some branches" << std::endl;
  femtoReader->SetStatus("*", 0);
  femtoReader->SetStatus("Event", 1);
  femtoReader->SetStatus("Track", 1);
  std::cout << "Status has been set" << std::endl;

  std::cout << "Now I know what to read, Master!" << std::endl;

  if (!femtoReader->chain())
  {
    std::cout << "No chain has been found." << std::endl;
  }
  Long64_t eventsInTree = femtoReader->tree()->GetEntries();
  std::cout << "eventsInTree: " << eventsInTree << std::endl;
  Long64_t events2read = femtoReader->chain()->GetEntries();

  std::cout << "Number of events to read: " << events2read << std::endl;


  const int NptBins = 18;
  const double ptBin [NptBins+1] = {0.2,0.4,0.6,0.8,1.,1.2,1.4,1.6,1.8,2.,2.2,2.4,2.6,2.8,3.,3.2,3.4,3.6,3.8};
  TH2D *hNsigPionMSqr[NptBins];
  TH2D *hPionNewXY[NptBins];
  TH2D *hPionNewXYCut[NptBins];
  TH2D *hPionNewXYPion[NptBins];
  TH2D *hPionNewXYKaon[NptBins];
  TH2D *hPionNewXYProton[NptBins];
  TH2D *hNsigKaonMSqr[NptBins];
  TH2D *hNsigProtonMSqr[NptBins];

  TH1D *hNsigPion = new TH1D(Form("hNsigPion") ,Form("n#sigma_{#pi} for %.2f < p < %.2f GeV/c;n#sigma_{#pi};N_{counts}",0.6,0.7),1400,-35.,35.);
  TH1D *hNsigKaon = new TH1D(Form("hNsigKaon") ,Form("n#sigma_{K} for %.2f < p_{T} < %.2f GeV/c;n#sigma_{K};N_{counts}",0.6,0.7),1400,-35.,35.);
  TH1D *hNsigProton = new TH1D(Form("hNsigProton") ,Form("n#sigma_{p} for %.2f < p_{T} < %.2f GeV/c;n#sigma_{p};N_{counts}",0.9,1.2),1400,-35.,35.);

  TH1D *hNsigPionCut = new TH1D(Form("hNsigPionCut") ,Form("n#sigma_{#pi} w/ Cut for %.2f < p < %.2f GeV/c;n#sigma_{#pi};N_{counts}",0.6,0.7),1400,-35.,35.);
  TH1D *hNsigKaonCut = new TH1D(Form("hNsigKaonCut") ,Form("n#sigma_{K} w/ Cut for %.2f < p_{T} < %.2f GeV/c;n#sigma_{K};N_{counts}",0.6,0.7),1400,-35.,35.);
  TH1D *hNsigProtonCut = new TH1D(Form("hNsigProtonCut") ,Form("n#sigma_{p} w/ Cut for %.2f < p_{T} < %.2f GeV/c;n#sigma_{p};N_{counts}",0.9,1.2),1400,-35.,35.);

  //Initialization
  for (int i=0; i<NptBins; i++)
  {
    hNsigPionMSqr[i] = new TH2D(Form("hNsigPionMSqr%i",i) ,Form("n#sigma_{#pi} vs M^{2} for %.2f < p_{T} < %.2f GeV/c;n#sigma_{#pi};m^{2} [GeV/c^{2}]^{2}",ptBin[i],ptBin[i+1]),1400,-35.,35.,1400,-2.,5.);
    hNsigKaonMSqr[i] = new TH2D(Form("hNsigKaonMSqr%i",i) ,Form("n#sigma_{K} vs M^{2} for %.2f < p_{T} < %.2f GeV/c;n#sigma_{K};m^{2} [GeV/c^{2}]^{2}",ptBin[i],ptBin[i+1]),1400,-35.,35.,1400,-2.,5.);
    hNsigProtonMSqr[i] = new TH2D(Form("hNsigProtonMSqr%i",i) ,Form("n#sigma_{p} vs M^{2} for %.2f < p_{T} < %.2f GeV/c;n#sigma_{p};m^{2} [GeV/c^{2}]^{2}",ptBin[i],ptBin[i+1]),1400,-35.,35.,1400,-2.,5.);
    hPionNewXY[i] = new TH2D(Form("hPionNewXY%i",i) ,Form("x(n#sigma_{#pi},M^{2}) vs y(n#sigma_{#pi},M^{2}) for %.2f < p_{T} < %.2f GeV/c;X;Y",ptBin[i],ptBin[i+1]),600,-0.5,1.,600,-1.,0.5);
    hPionNewXYCut[i] = new TH2D(Form("hPionNewXYCut%i",i) ,Form("x(n#sigma_{#pi},M^{2}) vs y(n#sigma_{#pi},M^{2}) with cut for %.2f < p_{T} < %.2f GeV/c;X;Y",ptBin[i],ptBin[i+1]),600,-0.5,1.,600,-1.,0.5);
    hPionNewXYPion[i] = new TH2D(Form("hPionNewXYPion%i",i) ,Form("x(n#sigma_{#pi},M^{2}) vs y(n#sigma_{#pi},M^{2}) for #pi^{#pm} for %.2f < p_{T} < %.2f GeV/c;X;Y",ptBin[i],ptBin[i+1]),600,-0.5,1.,600,-1.,0.5);
    hPionNewXYKaon[i] = new TH2D(Form("hPionNewXYKaon%i",i) ,Form("x(n#sigma_{#pi},M^{2}) vs y(n#sigma_{#pi},M^{2}) for K^{#pm} for %.2f < p_{T} < %.2f GeV/c;X;Y",ptBin[i],ptBin[i+1]),600,-0.5,1.,600,-1.,0.5);
    hPionNewXYProton[i] = new TH2D(Form("hPionNewXYProton%i",i) ,Form("x(n#sigma_{#pi},M^{2}) vs y(n#sigma_{#pi},M^{2}) for (p#bar{p}) for %.2f < p_{T} < %.2f GeV/c;X;Y",ptBin[i],ptBin[i+1]),600,-0.5,1.,600,-1.,0.5);
  }

	std::pair<Double_t,Double_t> XY_pid;

  // Loop over events
  for (Long64_t iEvent = 0; iEvent < events2read; iEvent++)
  {

    if ((iEvent+1) % 1000 == 0){
      std::cout << "Working on event #[" << (iEvent + 1)
                << "/" << events2read << "]" << std::endl;
    }

    Bool_t readEvent = femtoReader->readFemtoEvent(iEvent);
    if (!readEvent)
    {
      std::cout << "Something went wrong, Master! Nothing to analyze..." << std::endl;
      break;
    }

    // Retrieve femtoDst
    StFemtoDst *dst = femtoReader->femtoDst();

    // Retrieve event information
    StFemtoEvent *event = dst->event();
    if (!event)
    {
      std::cout << "Something went wrong, Master! Event is hiding from me..." << std::endl;
      break;
    }

    TVector3 pVtx = event->primaryVertex();

    // Event selection
    if ( !isGoodEvent(event) ) continue;

    // Track analysis
    Int_t nTracks = dst->numberOfTracks();
    

    // Track loop
    for (Int_t iTrk = 0; iTrk < nTracks; iTrk++)
    {

      // Retrieve i-th femto track
      StFemtoTrack *femtoTrack = dst->track(iTrk);

      if (!femtoTrack)
        continue;

      // Must be a primary track
      if (!femtoTrack->isPrimary())
        continue;

      //Track selection
      if (!isGoodTrackFlow(femtoTrack, energy, pVtx)) continue;

      // Determine pt bin for PID plots
      int i_pt = -1;
      for (int i=0;i<NptBins;i++)
      {
        if (femtoTrack->pt() >= (ptBin[i]) && femtoTrack->pt() < (ptBin[i+1]))
	{
	  i_pt = i;
	}
      }
      if (i_pt == -1) continue;

      if (!femtoTrack->isTofTrack())
      {
        if (femtoTrack->ptot() >= 0.6 && femtoTrack->ptot() < 0.7)
	{
	  hNsigPion->Fill(femtoTrack->nSigmaPion());
	  if (femtoTrack->nSigmaKaon() > 2 && femtoTrack->nSigmaPion() < 2) hNsigPionCut->Fill(femtoTrack->nSigmaPion());

	  hNsigKaon->Fill(femtoTrack->nSigmaKaon());
	  if (femtoTrack->nSigmaPion() > 3 && femtoTrack->nSigmaKaon() < 2) hNsigKaonCut->Fill(femtoTrack->nSigmaKaon());
	}
        if (femtoTrack->ptot() >= 0.9 && femtoTrack->ptot() < 1.2)
	{
	  hNsigProton->Fill(femtoTrack->nSigmaProton());
	  if (femtoTrack->nSigmaPion() > 3 && femtoTrack->nSigmaProton() < 2) hNsigProtonCut->Fill(femtoTrack->nSigmaProton());
	}
      }

      // Check if track has TOF signal
      if (femtoTrack->isTofTrack())
      {
        if (femtoTrack->gDCA(pVtx).Mag() >= 1.) continue;
	XY_pid = GetNewXY(femtoTrack);
	hPionNewXY[i_pt]->Fill(XY_pid.first,XY_pid.second);
	if (isMesonXY(XY_pid.first, XY_pid.second))
	{
	  hPionNewXYCut[i_pt]->Fill(XY_pid.first,XY_pid.second);
				
	  if ( TMath::Power(XY_pid.first / (2*GetSigmNewX(femtoTrack,0)),2) + TMath::Power(XY_pid.second / (2*GetSigmNewY(femtoTrack,0)),2) < 1. && 
	     TMath::Power((XY_pid.first - 0.23) / (2*GetSigmNewX(femtoTrack,1)),2) + TMath::Power(XY_pid.second / (2*GetSigmNewY(femtoTrack,1)),2) > 1. )
	  {
	    hPionNewXYPion[i_pt]->Fill(XY_pid.first,XY_pid.second);
	  }
	  if ( TMath::Power((XY_pid.first - 0.23) / (2*GetSigmNewX(femtoTrack,1)),2) + TMath::Power(XY_pid.second / (2*GetSigmNewY(femtoTrack,1)),2) < 1. &&
	       TMath::Power(XY_pid.first / (2*GetSigmNewX(femtoTrack,0)),2) + TMath::Power(XY_pid.second / (2*GetSigmNewY(femtoTrack,0)),2) > 1. )
	  {
	    hPionNewXYKaon[i_pt]->Fill(XY_pid.first,XY_pid.second);
	  }
	}
	else
	{
	  hPionNewXYProton[i_pt]->Fill(XY_pid.first,XY_pid.second);
	}

        hNsigPionMSqr[i_pt]->Fill(femtoTrack->nSigmaPion(),femtoTrack->massSqr());
        hNsigKaonMSqr[i_pt]->Fill(femtoTrack->nSigmaKaon(),femtoTrack->massSqr());
        hNsigProtonMSqr[i_pt]->Fill(femtoTrack->nSigmaProton(),femtoTrack->massSqr());

      } //if( isTofTrack() )

    } //for(Int_t iTrk=0; iTrk<nTracks; iTrk++)

  } //for(Long64_t iEvent=0; iEvent<events2read; iEvent++)

  femtoReader->Finish();

  TFile *output = new TFile(outFile,"recreate");

  for (int i=0; i<NptBins; i++)
  {
    hNsigPionMSqr[i] -> Write();
    hNsigKaonMSqr[i] -> Write();
    hNsigProtonMSqr[i] -> Write();
  }
  
	for (int i=0; i<NptBins; i++)
  {
    hPionNewXY[i] -> Write();
    hPionNewXYCut[i] -> Write();
	}

	for (int i=0; i<NptBins; i++)
  {
    hPionNewXYPion[i] -> Write();
    hPionNewXYKaon[i] -> Write();
    hPionNewXYProton[i] -> Write();
	}

  hNsigPion->Write();
  hNsigKaon->Write();
  hNsigProton->Write();
  hNsigPionCut->Write();
  hNsigKaonCut->Write();
  hNsigProtonCut->Write();

  output->Close();

  std::cout << "I'm done with analysis. We'll have a Nobel Prize, Master!"
            << std::endl;
  timer.Stop();
  timer.Print();
}

Bool_t isGoodEvent(StFemtoEvent *const &event)
{
  if (!event) return false;
  if (event == nullptr) return false;

  if (event->primaryVertex().Perp() > cutVtxR) return false;
  if (TMath::Abs(event->primaryVertex().Z()) > cutVtxZEnergy.at(energy)) return false;

  if ((energy == 200.) && TMath::Abs(event->primaryVertex().Z() - event->vpdVz()) > cutVpdVz) return false;

  return true;
}

Bool_t isGoodTrackFlow(StFemtoTrack *const &track, Float_t _energy, TVector3 pVtx)
{
  if (!track) return false;
  // if (!track->isPrimary()) return false;
  if (track->nHitsFit() < cutNhits) return false;
  if (track->nHitsPoss() <= cutNhitsPoss) return false;
  if ((Double_t) track->nHitsFit()/track->nHitsPoss() < cutNhitsRatio) return false;
  if (TMath::Abs(track->eta()) >= cutEta) return false;
  
  if (track->pt() <= cutPtMin.at(_energy)) return false;
  //if (track->pt() > cutPtMax) return false;
  if (track->ptot() > cutPMax) return false;
  if (track->gDCA(pVtx).Mag() >= cutDCA.at(_energy)) return false;
  return true;
}

Double_t GetMeanNsig(StFemtoTrack *const &track, const Int_t PID)
{
  Double_t pt, par0, par1, par2, par3;
  if (PID == 0) return 0.;
  if (PID == 1)
  {
    par0  = 5.84263e+00;
    par1  = -6.65912e+00;
    par2  = 8.32204e-01;
    par3  = 8.28440e-01;
  }
  if (PID == 2)
  {
    par0  = 1.09685e+01;
    par1  = -8.61499e+00;
    par2  = 8.92938e-01;
    par3  = 8.08397e-01;
  }
  pt = track->pt();
  if (pt != 0)
	{
    return (par0*TMath::Power(1./pt,par2) + par1 + par3*pt);
	}
	else
  {
		return 0.;
  }
}

Double_t GetWidthM2(StFemtoTrack *const &track, const Int_t PID)
{
  Double_t pt, pol0, pol1, pol2;
  if (PID != 0) return 0.;
  pol0  = 8.74975e-04;
  pol1  = -1.62659e-03;
  pol2  = 2.89828e-02;

	pt = track->pt();
  return (pol0 + pol1*pt + pol2*pt*pt);
}

Double_t GetScaleFactor(StFemtoTrack *const &track)
{
  Double_t wM2pion = GetWidthM2(track,0);
  Double_t exScale = 1./0.8273; // = 1. for BES energies
  if (wM2pion == 0) return 1.;
  else return (exScale/wM2pion);
}

Double_t GetRotAngle(StFemtoTrack *const &track)
{
	Double_t scaleFactor = GetScaleFactor(track);
  Double_t MeanKaonNsig = GetMeanNsig(track,1);

	Double_t new_x = (MeanKaonNsig - 0.)/scaleFactor;
	Double_t new_y = (kaonMassSqr - pionMassSqr);
	
  return ( -TMath::ATan2(new_y,new_x) );
}

std::pair<Double_t,Double_t> GetNewXY(StFemtoTrack *const &track)
{ 
  std::pair<Double_t,Double_t> coord;
	Double_t scaleFactor = GetScaleFactor(track);

	Double_t new_x = (track->nSigmaPion() - 0.)/scaleFactor;
	Double_t new_y = (track->massSqr() - pionMassSqr);

	Double_t rotAngle = GetRotAngle(track);
  coord.first  = new_x * TMath::Cos(rotAngle) - new_y * TMath::Sin(rotAngle);
	coord.second = new_x * TMath::Sin(rotAngle) + new_y * TMath::Cos(rotAngle);

  return coord;
}

Bool_t isMesonXY(Double_t x, Double_t y)
// Used to exclude proton peak from new XY coord distribution
{
  Double_t pol0, pol1;
  pol0 = -0.92857143;
	pol1 =  1.4285714;

	// Not a proton
	if (y > (pol0 + pol1*x)) return true;
  // Proton
	if (y <= (pol0 + pol1*x)) return false;

	// just in case...
	return false;
}

Double_t GetSigmNewX(StFemtoTrack *const &track, Int_t PID)
{
  if (PID != 0 && PID != 1 && PID != 2) return 0.;
  Double_t pol0, pol1, pol2, pol3;
  Double_t pt = track->pt();
  if (PID == 0) // pion
  {
    pol0 = 0.00061920122;
    pol1 = 0.00064431355;
    pol2 = 0.013672155;
    pol3 = 0.0012772833;
  }
  if (PID == 1) // kaon
  {
    pol0 = 0.0057203659;
    pol1 = 0.0013767144;
    pol2 = 0.0063761902;
    pol3 = 0.0073111853;
  }
  return ( pol0 + pol1*pt + pol2*pt*pt + pol3*pt*pt*pt );
}

Double_t GetSigmNewY(StFemtoTrack *const &track, Int_t PID)
{
  if (PID != 0 && PID != 1 && PID != 2) return 0.;
  Double_t pol0, pol1, pol2, pol3;
  Double_t pt = track->pt();
  if (PID == 0) // pion
  {
    pol0 = 0.0017345792;
    pol1 = -0.0059824003;
    pol2 = 0.029537517;
    pol3 = -0.0037990834;
  }
  if (PID == 1) // kaon
  {
    pol0 = 0.0041058086;
    pol1 = -0.0041708451;
    pol2 = 0.026109827;
    pol3 = -0.0018772891;
  }
  return ( pol0 + pol1*pt + pol2*pt*pt + pol3*pt*pt*pt );
}