source: HiSusy/trunk/Delphes-3.0.0/readers/DelphesLHEF.cpp @ 1

#include <stdexcept>
#include <iostream>
#include <sstream>

#include <stdlib.h>
#include <signal.h>
#include <stdio.h>

#include "TROOT.h"
#include "TApplication.h"

#include "TFile.h"
#include "TObjArray.h"
#include "TStopwatch.h"
#include "TDatabasePDG.h"
#include "TParticlePDG.h"
#include "TLorentzVector.h"

#include "modules/Delphes.h"
#include "classes/DelphesStream.h"
#include "classes/DelphesClasses.h"
#include "classes/DelphesFactory.h"

#include "ExRootAnalysis/ExRootTreeWriter.h"
#include "ExRootAnalysis/ExRootTreeBranch.h"
#include "ExRootAnalysis/ExRootProgressBar.h"

using namespace std;

static const int kBufferSize  = 1024;

class LHEFConverter
{
public:

  LHEFConverter(DelphesFactory *factory, TObjArray *particleOutputArray,
                TObjArray *candidateOutputArray, TObjArray *partonOutputArray);
  ~LHEFConverter();

  void Clear();
  Bool_t EventReady();

  Bool_t ReadLine(FILE *inputFile);

  void AnalyzeEvent(ExRootTreeBranch *branch, Long64_t eventNumber,
                    TStopwatch *readStopWatch, TStopwatch *procStopWatch);
  void AnalyzeParticle();
  
private:

  Int_t fParticleCounter, fProcessID;
  Double_t fWeight, fScalePDF, fAlphaQCD, fAlphaQED;
  
  Int_t fPID, fStatus, fM1, fM2, fC1, fC2;
  Double_t fPx, fPy, fPz, fE, fMass;
  
  Int_t fEventCounter;

  char *fBuffer;

  TDatabasePDG *fPDG;
  DelphesFactory *fFactory;
  
  TIterator *fItParticleOutputArray;
  
  TObjArray *fParticleOutputArray;
  TObjArray *fCandidateOutputArray;
  TObjArray *fPartonOutputArray;  
};

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

void LHEFConverter::Clear()
{
  fEventCounter = -1;
  fParticleCounter = -1;
}

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

Bool_t LHEFConverter::EventReady()
{
  return (fEventCounter == 0) && (fParticleCounter == 0);
}

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

Bool_t LHEFConverter::ReadLine(FILE *inputFile)
{
  int rc;

  if(!fgets(fBuffer, kBufferSize, inputFile)) return kFALSE;

  DelphesStream bufferStream(fBuffer);

  if(strstr(fBuffer, "<event"))
  {
    Clear();
    fEventCounter = 1;
  }
  else if(fEventCounter > 0)
  {
    rc = bufferStream.ReadInt(fParticleCounter)
      && bufferStream.ReadInt(fProcessID)
      && bufferStream.ReadDbl(fWeight)
      && bufferStream.ReadDbl(fScalePDF)
      && bufferStream.ReadDbl(fAlphaQED)
      && bufferStream.ReadDbl(fAlphaQCD);

    if(!rc)
    {
      cerr << "** ERROR: " << "invalid event format" << endl;
      return kFALSE;
    }

    --fEventCounter;
  }
  else if(fParticleCounter > 0)
  {
    rc = bufferStream.ReadInt(fPID)
      && bufferStream.ReadInt(fStatus)
      && bufferStream.ReadInt(fM1)
      && bufferStream.ReadInt(fM2)
      && bufferStream.ReadInt(fC1)
      && bufferStream.ReadInt(fC2)
      && bufferStream.ReadDbl(fPx)
      && bufferStream.ReadDbl(fPy)
      && bufferStream.ReadDbl(fPz)
      && bufferStream.ReadDbl(fE)
      && bufferStream.ReadDbl(fMass);

    if(!rc)
    {
      cerr << "** ERROR: " << "invalid particle format" << endl;
      return kFALSE;
    }

    AnalyzeParticle();

    --fParticleCounter;
  }
  
  return kTRUE;
}

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

LHEFConverter::LHEFConverter(DelphesFactory *factory, TObjArray *particleOutputArray,
                             TObjArray *candidateOutputArray, TObjArray *partonOutputArray) :
  fBuffer(0), fPDG(0),
  fFactory(factory),
  fItParticleOutputArray(0),
  fParticleOutputArray(particleOutputArray),
  fCandidateOutputArray(candidateOutputArray),
  fPartonOutputArray(partonOutputArray)
{
  fBuffer = new char[kBufferSize];
    
  fPDG = TDatabasePDG::Instance();
  fItParticleOutputArray = fParticleOutputArray->MakeIterator();
}

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

LHEFConverter::~LHEFConverter()
{
  if(fItParticleOutputArray) delete fItParticleOutputArray;
  if(fBuffer) delete[] fBuffer;
}

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

void LHEFConverter::AnalyzeEvent(ExRootTreeBranch *branch, Long64_t eventNumber,
                                 TStopwatch *readStopWatch, TStopwatch *procStopWatch)
{
  LHEFEvent *element;
  
  element = static_cast<LHEFEvent *>(branch->NewEntry());
  element->Number = eventNumber;

  element->ProcessID = fProcessID;
  element->Weight = fWeight;
  element->ScalePDF = fScalePDF;
  element->AlphaQED = fAlphaQED;
  element->AlphaQCD = fAlphaQCD;

  element->ReadTime = readStopWatch->RealTime();
  element->ProcTime = procStopWatch->RealTime();
}

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

void LHEFConverter::AnalyzeParticle()
{
  Candidate *candidate;
  TParticlePDG *pdgParticle;

  candidate = fFactory->NewCandidate();
  
  candidate->PID = fPID;

  candidate->Status = fStatus;

  pdgParticle = fPDG->GetParticle(fPID);
  candidate->Charge = pdgParticle ? Int_t(pdgParticle->Charge()/3.0) : -999;
  candidate->Mass = pdgParticle ? pdgParticle->Mass() : -999.9;
  
  candidate->Momentum.SetPxPyPzE(fPx, fPy, fPz, fE);
  candidate->Position.SetXYZT(0.0, 0.0, 0.0, 0.0);

  candidate->M1 = fM1 - 1;
  candidate->M2 = fM2 - 1;

  candidate->D1 = -1;
  candidate->D2 = -1;

  fParticleOutputArray->Add(candidate);
  
  if(!pdgParticle) return;

  if(fStatus == 1)
  {
    fCandidateOutputArray->Add(candidate);
  }
  else if(fStatus == 2)
  {
    fPartonOutputArray->Add(candidate);
  }
}

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

static bool interrupted = false;

void SignalHandler(int sig)
{
        interrupted = true;
}

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

int main(int argc, char *argv[])
{
  char appName[] = "DelphesLHEF";
  stringstream message;
  FILE *inputFile = 0;
  TFile *outputFile = 0;
  TStopwatch readStopWatch, procStopWatch;
  ExRootTreeWriter *treeWriter = 0;
  ExRootTreeBranch *branchEvent = 0;
  ExRootConfReader *confReader = 0;
  Delphes *modularDelphes = 0;
  DelphesFactory *factory = 0;
  TObjArray *candidateOutputArray = 0, *particleOutputArray = 0, *partonOutputArray = 0;
  LHEFConverter *converter = 0;
  Int_t i;
  Long64_t length, eventCounter;

  if(argc < 3)
  {
    cout << " Usage: " << appName << " config_file" << " output_file" << " [input_file(s)]" << endl;
    cout << " config_file - configuration file in Tcl format," << endl;
    cout << " output_file - output file in ROOT format," << endl;
    cout << " input_file(s) - input file(s) in LHEF format," << endl;
    cout << " with no input_file, or when input_file is -, read standard input." << endl;
    return 1;
  }

  signal(SIGINT, SignalHandler);

  gROOT->SetBatch();

  int appargc = 1;
  char *appargv[] = {appName};
  TApplication app(appName, &appargc, appargv);
        
  try
  {
    outputFile = TFile::Open(argv[2], "RECREATE");

    if(outputFile == NULL)
    {
      message << "can't open " << argv[2];
      throw runtime_error(message.str());
    }

    treeWriter = new ExRootTreeWriter(outputFile, "Delphes");

    branchEvent = treeWriter->NewBranch("Event", LHEFEvent::Class());

    confReader = new ExRootConfReader;
    confReader->ReadFile(argv[1]);
    
    modularDelphes = new Delphes("Delphes");
    modularDelphes->SetConfReader(confReader);
    modularDelphes->SetTreeWriter(treeWriter);
    
    factory = modularDelphes->GetFactory();
    particleOutputArray = modularDelphes->ExportArray("particles");
    candidateOutputArray = modularDelphes->ExportArray("candidates");
    partonOutputArray = modularDelphes->ExportArray("partons");

    converter = new LHEFConverter(factory, particleOutputArray, candidateOutputArray, partonOutputArray);

    modularDelphes->InitTask();

    i = 3;
    do
    {
      if(interrupted) break;

      if(i == argc || strcmp(argv[i], "-") == 0)
      {
        cout << "** Reading standard input" << endl;
        inputFile = stdin;
        length = -1;
      }
      else
      {
        cout << "** Reading " << argv[i] << endl;
        inputFile = fopen(argv[i], "r");

        if(inputFile == NULL)
        {
          message << "can't open " << argv[i];
          throw runtime_error(message.str());
        }

        fseek(inputFile, 0L, SEEK_END);
        length = ftell(inputFile);
        fseek(inputFile, 0L, SEEK_SET);

        if(length <= 0) continue;
      }

      ExRootProgressBar progressBar(length);
  
      // Loop over all objects
      eventCounter = 0;
      treeWriter->Clear();
      modularDelphes->Clear();
      converter->Clear();
      readStopWatch.Start();
      while(converter->ReadLine(inputFile) && !interrupted)
      {
        if(converter->EventReady())
        {
          ++eventCounter;

          readStopWatch.Stop();
          procStopWatch.Start();
          modularDelphes->ProcessTask();
          procStopWatch.Stop();

          converter->AnalyzeEvent(branchEvent, eventCounter, &readStopWatch, &procStopWatch);
    
          treeWriter->Fill();

          treeWriter->Clear();
          modularDelphes->Clear();
          converter->Clear();      

          readStopWatch.Start();
        }
        progressBar.Update(ftell(inputFile), eventCounter);
      }

      fseek(inputFile, 0L, SEEK_END);
      progressBar.Update(ftell(inputFile), eventCounter, kTRUE);
      progressBar.Finish();

      if(inputFile != stdin) fclose(inputFile);

      ++i;
    }
    while(i < argc);

    modularDelphes->FinishTask();
    treeWriter->Write();

    cout << "** Exiting..." << endl;

    delete converter;
    delete modularDelphes;
    delete confReader;
    delete treeWriter;
    delete outputFile;
    
    return 0;
  }
  catch(runtime_error &e)
  {
    if(treeWriter) delete treeWriter;
    if(outputFile) delete outputFile;
    cerr << "** ERROR: " << e.what() << endl;
    return 1;
  }
}