source: svn/trunk/src/HEPEVTConverter.cc@ 606

/***********************************************************************
**                                                                    **
**          /----------------------------------------------\          **
**         |  Delphes, a framework for the fast simulation  |         **
**         |       of a  generic collider experiment        |         **
**          \-------------  arXiv:0903.2225v1  ------------/          **
**                                                                    **
**                                                                    **
**   This package uses:                                               **
**   ------------------                                               **
**   ROOT: Nucl. Inst. & Meth. in Phys. Res. A389 (1997) 81-86        **
**   FastJet algorithm: Phys. Lett. B641 (2006) [hep-ph/0512210]      **
**   Hector: JINST 2:P09005 (2007) [physics.acc-ph:0707.1198v2]       **
**   FROG: [hep-ex/0901.2718v1]                                       **
**   HepMC: Comput. Phys. Commun.134 (2001) 41                        **
**                                                                    **
** ------------------------------------------------------------------ **
**                                                                    **
**   Main authors:                                                    **
**   -------------                                                    **
**                                                                    **
**                Severine Ovyn                Xavier Rouby           **
**          severine.ovyn@uclouvain.be      xavier.rouby@cern         **
**                                                                    **
**         Center for Particle Physics and Phenomenology (CP3)        **
**                Universite catholique de Louvain (UCL)              **
**                     Louvain-la-Neuve, Belgium                      **
**                                                                    **
**                      Copyright (C) 2008-2009,                      **
**                        All rights reserved.                        **
**                                                                    **
***********************************************************************/

#include <iostream>
#include <utility>
#include <deque>
#include <sstream>
#include <fstream>
#include <cmath>
#include <cstdlib> // for exit()

#include "TROOT.h"
#include "TTree.h"
#include "TBranch.h"
#include "TLeaf.h"
#include "TString.h"
#include "TLorentzVector.h"
#include "TChain.h"

#include "BlockClasses.h"
#include "ExRootTreeWriter.h"
#include "ExRootTreeBranch.h"

#include "HEPEVTConverter.h"

using namespace std;

//------------------------------------------------------------------------------

HEPTreeReader::HEPTreeReader(TTree *tree, HEPEvent *event) : fChain(tree), fCurrentTree(-1), fEvent(event)
{
  if(!fChain) return;

  TBranch *branch;
  TLeaf *leaf;
  TString name;
  Int_t i;

  name = "Nhep";
  branch = fChain->GetBranch(name);
  branch->SetAddress(&(event->Nhep));
  fBranches.push_back(make_pair(name, branch));

  TString intNames[3] = {"Idhep", "Jsmhep", "Jsdhep"};
  Int_t **intData[3] = {&event->Idhep, &event->Jsmhep, &event->Jsdhep};

  for(i = 0; i < 3; ++i)
  {
    name = intNames[i];
    branch = fChain->GetBranch(name);
    leaf = branch->GetLeaf(name);
    *intData[i] = new Int_t[leaf->GetNdata()];
    branch->SetAddress(*intData[i]);
    fBranches.push_back(make_pair(name, branch));
  }

  TString floatNames[2] = {"Phep", "Vhep"};
  Float_t **floatData[2] = {&event->Phep, &event->Vhep};

  for(i = 0; i < 2; ++i)
  {
    name = floatNames[i];
    branch = fChain->GetBranch(name);
    leaf = branch->GetLeaf(name);
    *floatData[i] = new Float_t[leaf->GetNdata()];
    branch->SetAddress(*floatData[i]);
    fBranches.push_back(make_pair(name, branch));
  }
}

//------------------------------------------------------------------------------

HEPTreeReader::~HEPTreeReader()
{
  if(!fChain) return;

  delete[] fEvent->Idhep;
  delete[] fEvent->Jsmhep;
  delete[] fEvent->Jsdhep;
  delete[] fEvent->Phep;
  delete[] fEvent->Vhep;
}

//------------------------------------------------------------------------------

Bool_t HEPTreeReader::ReadEntry(Long64_t element)
{
  if(!fChain) return kFALSE;

  Int_t treeEntry = fChain->LoadTree(element);
  if(treeEntry < 0) return kFALSE;

  if(fChain->IsA() == TChain::Class())
  {
    TChain *chain = static_cast<TChain*>(fChain);
    if(chain->GetTreeNumber() != fCurrentTree)
    {
      fCurrentTree = chain->GetTreeNumber();
      Notify();
    }
  }

  deque< pair<TString, TBranch*> >::iterator it_deque;
  TBranch *branch;

  for(it_deque = fBranches.begin(); it_deque != fBranches.end(); ++it_deque)
  {
    branch = it_deque->second;
    if(branch)
    {
      branch->GetEntry(treeEntry);
    }
  }

  return kTRUE;
}

//------------------------------------------------------------------------------

void HEPTreeReader::Notify()
{
  // Called when loading a new file.
  // Get branch pointers.
  if(!fChain) return;

  deque< pair<TString, TBranch*> >::iterator it_deque;

  for(it_deque = fBranches.begin(); it_deque != fBranches.end(); ++it_deque)
  {
    it_deque->second = fChain->GetBranch(it_deque->first);
    if(!it_deque->second)
    {
      cerr << left  << setw(40) <<"**        WARNING: cannot get branch: "<<""
           << left  << setw(20) << it_deque->first                        <<""
           << right << setw(9) <<" **"<<""<<endl;
    }
  }
}



HEPEVTConverter::HEPEVTConverter(const string& inputFileList, const string& outputFileName, const PdgTable& pdg, const int& Nevents) : DataConverter(pdg,Nevents)
{
  string buffer;
  int nevt_already_processed=0;

  gROOT->SetBatch(); 

  ExRootTreeWriter *treeWriter = new ExRootTreeWriter(outputFileName, "GEN");
  ExRootTreeBranch *branchGen = treeWriter->NewBranch("Particle", TRootC::GenParticle::Class());
  
  ifstream infile(inputFileList.c_str());
  if(!infile.is_open())
    {
      cerr << left  << setw(30) <<"**   ERROR: Can't open     "<<""
           << left  << setw(20) << inputFileList                        <<""
           << right << setw(19) <<"for input **"<<""<<endl;
      exit(1);
    }
  while(1)
    {
      if (Nevt>0 && nevt_already_processed >=Nevt) break; // enough events already processed 
      TChain *chain = new TChain("h101");
      infile >> buffer;
        ifstream checking_the_file(buffer.c_str());
        if(!checking_the_file.good()) 
          { 
             cerr << left  << setw(30) <<"**   ERROR: Can't find file  "<<""
                  << left  << setw(36) << buffer                        <<""
                  << right << setw(3) <<" **"<<""<<endl;
             continue;
          }
        else checking_the_file.close();

      if(!infile.good()) break;  // end of the list
      chain->Add(buffer.c_str());
      
      HEPEvent event;
      HEPTreeReader *treeReader = new HEPTreeReader(chain, &event);
      
      Long64_t entry, allEntries = treeReader->GetEntries();
      Int_t address, particle;
      Float_t signEta;
      
      TRootC::GenParticle *element;
      allEntries = (Nevt<0)?allEntries:min((int)allEntries,Nevt);
      // Loop over all events
      for(entry = 0; entry < allEntries; ++entry)
        {
          if (Nevt>0 && nevt_already_processed >=Nevt) break; // enough events already processed
          // Load selected branches with data from specified event
          treeReader->ReadEntry(entry);
          treeWriter->Clear();
          
          for(particle = 0; particle < event.Nhep; ++particle)
            {
              element = (TRootC::GenParticle*) branchGen->NewEntry();
              
              element->PID = event.Idhep[particle];     
              element->Status = event.Jsmhep[particle]/16000000 
                + event.Jsdhep[particle]/16000000*100;
              element->M1 = (event.Jsmhep[particle]%16000000)/4000 -1;
              element->M2 = event.Jsmhep[particle]%4000 -1;
              element->D1 = (event.Jsdhep[particle]%16000000)/4000 -1;
              element->D2 = event.Jsdhep[particle]%4000 -1;
              address = particle*5;
              element->E = event.Phep[address + 3];
              element->Px = event.Phep[address + 0];
              element->Py = event.Phep[address + 1];
              element->Pz = event.Phep[address + 2];
              element->M = event.Phep[address + 4];
              
              TLorentzVector vector(element->Px, element->Py, element->Pz, element->E);
              element->PT = vector.Perp();
              signEta = (element->Pz >= 0.0) ? 1.0 : -1.0;
              element->Eta = vector.CosTheta()*vector.CosTheta() == 1.0 ? signEta*999.9 : vector.Eta();
              element->Phi = vector.Phi();
              
              address = particle*4;
              element->T = event.Vhep[address + 3];
              element->X = event.Vhep[address + 0];
              element->Y = event.Vhep[address + 1];
              element->Z = event.Vhep[address + 2];
            }
          treeWriter->Fill();
          ++nevt_already_processed;
        }
      delete chain;
      delete treeReader;
    }
  treeWriter->Write();
  
  delete treeWriter;

}