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

#include <iostream>
#include <utility>
#include <deque>
#include <sstream>
#include <fstream>

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

#include "Utilities/ExRootAnalysis/interface/BlockClasses.h"
#include "Utilities/ExRootAnalysis/interface/ExRootTreeReader.h"
#include "Utilities/ExRootAnalysis/interface/ExRootTreeWriter.h"
#include "Utilities/ExRootAnalysis/interface/ExRootTreeBranch.h"

#include "interface/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)
    {
      cout << "** WARNING: cannot get branch '" << it_deque->first << "'" << endl;
    }
  }
}



HEPEVTConverter::HEPEVTConverter(const string& inputFileList, const string& outputFileName)
{
  string buffer;

  gROOT->SetBatch(); 

  ExRootTreeWriter *treeWriter = new ExRootTreeWriter(outputFileName, "GEN");
  ExRootTreeBranch *branchGen = treeWriter->NewBranch("Particle", TRootGenParticle::Class());
  
  ifstream infile(inputFileList.c_str());
  if(!infile.is_open())
    {
      cerr << "** ERROR: Can't open '" << inputFileList << "' for input" << endl;
      exit(1);
    }
  while(1)
    {
      
      TChain *chain = new TChain("h101");
      infile >> buffer;
      if(!infile.good()) break;
      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;
      
      TRootGenParticle *element;
      
      // Loop over all events
      for(entry = 0; entry < allEntries; ++entry)
        {
          // Load selected branches with data from specified event
          treeReader->ReadEntry(entry);
          treeWriter->Clear();
          
          for(particle = 0; particle < event.Nhep; ++particle)
            {
              element = (TRootGenParticle*) 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];
              
              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();
        }
      delete chain;
      delete treeReader;
    }
  treeWriter->Write();
  cout << "** Exiting..." << endl;
  
  delete treeWriter;

}