FlowANA/urqmd2mcpico.C

#include <stdlib.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <map>

#include <Rtypes.h>
#include <TFile.h>
#include <TTree.h>
#include <TMath.h>
#include <TStopwatch.h>

void urqmd2mcpico(TString inFileName = "", TString outFileName = "")
{
  //   Create a timer object to benchmark this loop
  TStopwatch timer;
  timer.Start();

  // Init event/particle parameters
  //    Fields below marked with "*" in their comments are of importance as they will be used later in FlowANA for the flow measurements
  float d_bimp; // impact parameter *
  int d_npart;  // number of participants
  float d_phi2; // EP/RP plane angle for v2
  float d_phi3; // EP/RP plane angle for v3
  float d_ecc2; // eccentricity for v2
  float d_ecc3; // eccentricity for v3

  static const int max_nh = 25000;

  int d_nh;               // number of all particles in the event *
  float d_momx[max_nh];   //[momentum x] *
  float d_momy[max_nh];   //[momentum y] *
  float d_momz[max_nh];   //[momentum z] *
  float d_ene[max_nh];    //[energy] *
  int d_pdg[max_nh];      //[particle data group code] *
  int d_hid[max_nh];      //[history id]
  short d_charge[max_nh]; //[electric charge]

  //  Init output file
  TFile *hfile;
  hfile = new TFile(outFileName.Data(), "RECREATE");

  //  Init output file
  TTree *tree;
  tree = new TTree("mctree", "An example of a ROOT tree for mcpico format");

  tree->Branch("bimp", &d_bimp, "bimp/F"); // impact parametr
  tree->Branch("phi2", &d_phi2, "phi2/F"); // phiRP2
  tree->Branch("phi3", &d_phi3, "phi3/F"); // phiRP3
  tree->Branch("ecc2", &d_ecc2, "ecc2/F"); // eccentricity2
  tree->Branch("ecc3", &d_ecc3, "ecc3/F"); // eccentricity3

  tree->Branch("npart", &d_npart, "npart/I");        // number of participants
  tree->Branch("nh", &d_nh, "nh/I");                 // number of particles
  tree->Branch("momx", &d_momx, "momx[nh]/F");       //[x projection of the particle's momentum]
  tree->Branch("momy", &d_momy, "momy[nh]/F");       //[y projection of the particle's momentum]
  tree->Branch("momz", &d_momz, "momz[nh]/F");       //[z projection of the particle's momentum]
  tree->Branch("ene", &d_ene, "ene[nh]/F");          //[energy]
  tree->Branch("hid", &d_hid, "hid[nh]/I");          //[histrory id]
  tree->Branch("pdg", &d_pdg, "pdg[nh]/I");          //[particle data group code]
  tree->Branch("charge", &d_charge, "charge[nh]/S"); //[electric charge]

  // Init objects for input file reading
  std::ifstream is;
  std::stringstream ss;
  std::string str;

  // Init temporary parameters during file reading
  Double_t B, PsiRP, Time,
      r0, rx, ry, rz, p0, px, py, pz,
      m, ityp, i3, ichg, ncl, lcl, orr, dump;
  Int_t Nev, Npart, count,
      pid;
  Int_t ev;
  Int_t printev;
  Double_t told = 0;
  Double_t tnew = 0;
  const Int_t count1 = 3, count2 = 12;

  // Table of the urqmd particle types and corresponding pdg codes
  // Dictionary: {itype, pdg}
  const std::map<Int_t, Int_t> particleURQMD = {{100, 22},
                                                {101, 211},
                                                {101, 111},
                                                {102, 221},
                                                {103, 223},
                                                {104, 213},
                                                {104, 113},
                                                {105, 9000221},
                                                {106, 321},
                                                {106, 311},
                                                {107, 331},
                                                {108, 323},
                                                {108, 313},
                                                {109, 333},
                                                {110, 9000311},
                                                {110, 9000321},
                                                {111, 9000321},
                                                {1, 2212},
                                                {1, 2112},
                                                {27, 3122}};

  // Open input output.f14 file and exit if we fail to do so
  is.open(inFileName.Data());
  if (is.fail())
  {
    std::cerr << "urqmd2mcpico: Open input file: " << inFileName.Data() << std::endl;
    return;
  }

  // Main loop where we read the file. We print out the event count every 100 events.
  ev = 0;
  printev = 100;
  while (!is.eof())
  {
    // Print out information
    if (ev % printev == 0)
    {
      tnew = timer.RealTime();
      printf("\tevent:%d,\trtime=%f s\n", ev, tnew - told);
      told = tnew;
      timer.Continue();
    }

    ss.str("");
    ss.clear();
    getline(is, str);

    if (str[0] == ' ')
      continue;
    for (Int_t i = 0; i < count1 - 1; i++)
    {
      getline(is, str);
    }

    // Read impact parameter
    ss.str("");
    ss.clear();
    getline(is, str);
    ss << str;
    ss >> str >> B;
    getline(is, str);

    // Read number of event
    ss.str("");
    ss.clear();
    getline(is, str);
    ss << str;
    ss >> str >> Nev;
    for (int j = 0; j < count2; j++)
    {
      getline(is, str);
    }
    PsiRP = 0.;
    if (is.eof())
    {
      break;
    }

    // Read number of particles and time
    ss.str("");
    ss.clear();
    ss << str;
    ss >> Npart >> Time;
    getline(is, str);

    // Skip the event if it has no produced particles
    if (Npart <= 0)
      continue;

    // Fill tree info
    d_bimp = B;
    d_nh = Npart;
    d_npart = 0;
    d_phi2 = PsiRP;
    d_phi3 = PsiRP;
    d_ecc2 = 0;
    d_ecc3 = 0;

    // Loop over particles in the event
    for (int j = 0; j < Npart; j++)
    {
      ss.str("");
      ss.clear();
      getline(is, str);
      ss << str;
      ss >> r0 >> rx >> ry >> rz >> p0 >> px >> py >> pz >> m >> ityp >> i3 >> ichg >> dump >> dump >> dump >> dump >> dump >> dump;
      if (particleURQMD.find(TMath::Abs(ityp)) != particleURQMD.end())
      {
        pid = TMath::Sign(particleURQMD.at(TMath::Abs(ityp)), ichg);

        // UrQMD has the same itype for differently charged particles - here's a quick fix for some particle species
        if (ichg == 0 && ityp == 1) pid = 2112;
        if (ichg != 0 && ityp == 1) pid = 2212;
        if (ichg == 0 && ityp == 101) pid = 111;
        if (ichg != 0 && ityp == 101) pid = 211;
        if (ichg == 0 && ityp == 104) pid = 113;
        if (ichg != 0 && ityp == 104) pid = 213;
        if (ichg == 0 && ityp == 106) pid = 311;
        if (ichg != 0 && ityp == 106) pid = 321;
        if (ichg == 0 && ityp == 108) pid = 323;
        if (ichg != 0 && ityp == 108) pid = 313;
        if (ichg == 0 && ityp == 110) pid = 9000311;
        if (ichg != 0 && ityp == 110) pid = 9000321;
      }
      else
      {
        pid = -999.;
      }

      // Fill tree info
      d_momx[j] = px;
      d_momy[j] = py;
      d_momz[j] = pz;
      d_ene[j] = p0;
      d_pdg[j] = pid;
      d_hid[j] = -1;
      d_charge[j] = ichg;
    }

    // Fill the event information in the TTree
    tree->Fill();

    ev++;
  } // end of the while loop

  // Write TTree into ROOT file
  hfile->cd();
  tree->Write();

  //  Stop timer and print results
  timer.Stop();

  Double_t rtime = timer.RealTime();
  Double_t ctime = timer.CpuTime();

  printf("\n%d events has been processed.\n", ev);
  printf("RealTime=%f seconds, CpuTime=%f seconds\n", rtime, ctime);

  hfile->Close();
}