source: HiSusy/trunk/PythiaDelphes/main41.cc @ 1

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

#include <map>

#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;

// main41.cc is a part of the PYTHIA event generator.
// Copyright (C) 2012 Mikhail Kirsanov, Torbjorn Sjostrand.
// PYTHIA is licenced under the GNU GPL version 2, see COPYING for details.
// Please respect the MCnet Guidelines, see GUIDELINES for details.

// This is a simple test program.
// It illustrates how HepMC can be interfaced to Pythia8.
// It studies the charged multiplicity distribution at the LHC.
// HepMC events are output to the hepmcout41.dat file.
// Written by Mikhail Kirsanov based on main01.cc.

// WARNING: typically one needs 25 MB/100 events at the LHC.
// Therefore large event samples may be impractical.

#include "Pythia.h"
#include "HepMCInterface.h"

#include "HepMC/GenEvent.h"
#include "HepMC/IO_GenEvent.h"

// Following line is a deprecated alternative, removed in recent versions
//#include "HepMC/IO_Ascii.h"
//#include "HepMC/IO_AsciiParticles.h"

// Following line to be used with HepMC 2.04 onwards.
#ifdef HEPMC_HAS_UNITS
#include "HepMC/Units.h"
#endif

using namespace Pythia8; 

class HepMCConverter
{
public:

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

  void Clear();
  Bool_t EventReady();

  Bool_t ReadLine(FILE *inputFile);

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

  Int_t fEventNumber, fMPI, fProcessID, fSignalCode, fVertexCounter;
  Double_t fScale, fAlphaQCD, fAlphaQED;

  Int_t fID1, fID2;
  Double_t fX1, fX2, fScalePDF, fPDF1, fPDF2;

  Int_t fOutVertexCode, fVertexID, fInCounter, fOutCounter;
  Double_t fX, fY, fZ, fT;
  
  Int_t fParticleCode, fPID, fStatus, fInVertexCode;
  Double_t fPx, fPy, fPz, fE, fMass, fTheta, fPhi;
  
  Int_t fParticleCounter;

  char *fBuffer;

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

  map< Int_t, pair < Int_t, Int_t > > fMotherMap;    
  map< Int_t, pair < Int_t, Int_t > > fDaughterMap;   
};

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

void HepMCConverter::Clear()
{
  fVertexCounter = -1;
  fInCounter = -1;
  fOutCounter = -1;
  fMotherMap.clear();
  fParticleCounter = 0;
}

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

Bool_t HepMCConverter::EventReady()
{
  return (fVertexCounter == 0) && (fInCounter == 0) && (fOutCounter == 0);
}

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

Bool_t HepMCConverter::ReadLine(FILE *inputFile)
{
  map< Int_t, pair< Int_t, Int_t > >::iterator itMotherMap;
  map< Int_t, pair< Int_t, Int_t > >::iterator itDaughterMap;
  char key;
  int rc;

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

  DelphesStream bufferStream(fBuffer + 1);

  key = fBuffer[0];

  if(key == 'E')
  {
    Clear();

    rc = bufferStream.ReadInt(fEventNumber)
      && bufferStream.ReadInt(fMPI)
      && bufferStream.ReadDbl(fScale)
      && bufferStream.ReadDbl(fAlphaQCD)
      && bufferStream.ReadDbl(fAlphaQED)
      && bufferStream.ReadInt(fProcessID)
      && bufferStream.ReadInt(fSignalCode)
      && bufferStream.ReadInt(fVertexCounter);

    if(!rc)
    {
      cerr << "** ERROR: " << "invalid event format" << endl;
      return kFALSE;
    }
  }
  else if(key == 'F')
  {
    rc = bufferStream.ReadInt(fID1)
      && bufferStream.ReadInt(fID2)
      && bufferStream.ReadDbl(fX1)
      && bufferStream.ReadDbl(fX2)
      && bufferStream.ReadDbl(fScalePDF)
      && bufferStream.ReadDbl(fPDF1)
      && bufferStream.ReadDbl(fPDF2);

    if(!rc)
    {
      cerr << "** ERROR: " << "invalid PDF format" << endl;
      return kFALSE;
    }
  }
  else if(key == 'V' && fVertexCounter > 0)
  {
    rc = bufferStream.ReadInt(fOutVertexCode)
      && bufferStream.ReadInt(fVertexID)
      && bufferStream.ReadDbl(fX)
      && bufferStream.ReadDbl(fY)
      && bufferStream.ReadDbl(fZ)
      && bufferStream.ReadDbl(fT)
      && bufferStream.ReadInt(fInCounter)
      && bufferStream.ReadInt(fOutCounter);

    if(!rc)
    {
      cerr << "** ERROR: " << "invalid vertex format" << endl;
      return kFALSE;
    }
    --fVertexCounter;
  }
  else if(key == 'P' && fOutCounter > 0)
  {
    rc = bufferStream.ReadInt(fParticleCode)
      && bufferStream.ReadInt(fPID)
      && bufferStream.ReadDbl(fPx)
      && bufferStream.ReadDbl(fPy)
      && bufferStream.ReadDbl(fPz)
      && bufferStream.ReadDbl(fE)
      && bufferStream.ReadDbl(fMass)
      && bufferStream.ReadInt(fStatus)
      && bufferStream.ReadDbl(fTheta)
      && bufferStream.ReadDbl(fPhi)
      && bufferStream.ReadInt(fInVertexCode);

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

    if(fInVertexCode < 0)
    {
      itMotherMap = fMotherMap.find(fInVertexCode);
      if(itMotherMap == fMotherMap.end())
      {
        fMotherMap[fInVertexCode] = make_pair(fParticleCounter, -1);
      }
      else
      {
        itMotherMap->second.second = fParticleCounter;
      }
    }

    if(fInCounter <= 0)
    {
      itDaughterMap = fDaughterMap.find(fOutVertexCode);
      if(itDaughterMap == fDaughterMap.end())
      {
        fDaughterMap[fOutVertexCode] = make_pair(fParticleCounter, fParticleCounter);
      }
      else
      {
        itDaughterMap->second.second = fParticleCounter;
      }
    }

    AnalyzeParticle();

    if(fInCounter > 0)
    {
      --fInCounter;
    }
    else
    {
      --fOutCounter;
    }

    ++fParticleCounter;
  }
  
  return kTRUE;
  
}

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

HepMCConverter::HepMCConverter(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();
}

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

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

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

void HepMCConverter::AnalyzeEvent(ExRootTreeBranch *branch,
                                  TStopwatch *readStopWatch, TStopwatch *procStopWatch)
{
  HepMCEvent *element;
  
  element = static_cast<HepMCEvent *>(branch->NewEntry());
  element->Number = fEventNumber;

  element->ProcessID = fProcessID;
  element->MPI = fMPI;
  element->Scale = fScale;
  element->AlphaQED = fAlphaQED;
  element->AlphaQCD = fAlphaQCD;

  element->ID1 = fID1;
  element->ID2 = fID2;
  element->X1 = fX1;
  element->X2 = fX2;
  element->ScalePDF = fScalePDF;
  element->PDF1 = fPDF1;
  element->PDF2 = fPDF2;
  
  element->ReadTime = readStopWatch->RealTime();
  element->ProcTime = procStopWatch->RealTime();
}

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

void HepMCConverter::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->M2 = 1;
  candidate->D2 = 1;
  if(fInCounter > 0)
  {
    candidate->M1 = 1;
    candidate->Position.SetXYZT(0.0, 0.0, 0.0, 0.0);
  }
  else
  {
    candidate->M1 = fOutVertexCode;
    candidate->Position.SetXYZT(fX, fY, fZ, fT);
  }
  if(fInVertexCode < 0)
  {
    candidate->D1 = fInVertexCode;
  }
  else
  {
    candidate->D1 = 1;
  }

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

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

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

void HepMCConverter::FinalizeParticles()
{
  Candidate *candidate;
  map< Int_t, pair< Int_t, Int_t > >::iterator itMotherMap;
  map< Int_t, pair< Int_t, Int_t > >::iterator itDaughterMap;
  
  fItParticleOutputArray->Reset();
  while((candidate = static_cast<Candidate *>(fItParticleOutputArray->Next())))
  {
    if(candidate->M1 > 0)
    {
      candidate->M1 = -1;
      candidate->M2 = -1;
    }
    else
    {
      itMotherMap = fMotherMap.find(candidate->M1);
      if(itMotherMap == fMotherMap.end())
      {
        candidate->M1 = -1;
        candidate->M2 = -1;
      }
      else
      {
        candidate->M1 = itMotherMap->second.first;
        candidate->M2 = itMotherMap->second.second;
      }
    }
    if(candidate->D1 > 0)
    {
      candidate->D1 = -1;
      candidate->D2 = -1;
    }
    else
    {
      itDaughterMap = fDaughterMap.find(candidate->D1);
      if(itDaughterMap == fDaughterMap.end())
      {
        candidate->D1 = -1;
        candidate->D2 = -1;
      }
      else
      {
        candidate->D1 = itDaughterMap->second.first;
        candidate->D2 = itDaughterMap->second.second;
      }
    }
  }
}

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

static bool interrupted = false;

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

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

int main(int argc, char *argv[])
{
  char appName[] = "DelphesHepMC";
  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;
  HepMCConverter *converter = 0;
  Int_t i;
  Long64_t length, eventCounter;

  //PYTHIA8 BEGIN
  // Interface for conversion from Pythia8::Event to HepMC one. 
  HepMC::I_Pythia8 ToHepMC;
  //  ToHepMC.set_crash_on_problem();

  // Specify file where HepMC events will be stored.
  HepMC::IO_GenEvent ascii_io("hepmcout41.dat", std::ios::out);
  // Following line is a deprecated alternative, removed in recent versions
  // HepMC::IO_Ascii ascii_io("hepmcout31.dat", std::ios::out);
  // Line below is an eye-readable one-way output, uncomment the include above
  // HepMC::IO_AsciiParticles ascii_io("hepmcout31.dat", std::ios::out);

  // Generator. Process selection. LHC initialization. Histogram.
  Pythia pythia;
  pythia.readString("Beams:eCM = 8000.");    
  pythia.readString("HardQCD:all = on");    
  pythia.readString("PhaseSpace:pTHatMin = 20.");    
  pythia.init();
  //PYTHIA8 END

  //DELPHES
  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 HepMC 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", HepMCEvent::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 HepMCConverter(factory, particleOutputArray, candidateOutputArray, partonOutputArray);

    modularDelphes->InitTask();
  
    i = 3;
    do
    {
      if(interrupted) break;

      if(i == argc || strcmp(argv[i], "-") == 0)
      {
        cout << "** Reading in memory" << endl;
        inputFile = stdin;
        length = -1;
      }

      ExRootProgressBar progressBar(length);
  
      // Loop over all objects
      eventCounter = 0;
      treeWriter->Clear();
      modularDelphes->Clear();
      converter->Clear();
      readStopWatch.Start();

      pythia.next();

#ifdef HEPMC_HAS_UNITS
      // This form with arguments is only meaningful for HepMC 2.04 onwards, 
      // and even then unnecessary if HepMC was built with GeV and mm as units.
      HepMC::GenEvent* hepmcevt = new HepMC::GenEvent(HepMC::Units::GEV, HepMC::Units::MM);
#else
      // This form is needed for backwards compatibility. 
      // In HepMCInterface.cc a conversion from GeV to MeV will be done.
      // To retain units in GeV uncomment line below.
      //ToHepMC.set_convert_to_mev( false);
      HepMC::GenEvent* hepmcevt = new HepMC::GenEvent();
#endif
      
      // Fill HepMC event, including PDF info.
      ToHepMC.fill_next_event( pythia, hepmcevt );
        
      while( converter->ReadLine(inputFile)  && !interrupted)
        {
          if(converter->EventReady())
            {
              ++eventCounter;
              
              converter->FinalizeParticles();
              
              readStopWatch.Stop();
              procStopWatch.Start();
              modularDelphes->ProcessTask();
              procStopWatch.Stop();
              
              converter->AnalyzeEvent(branchEvent, &readStopWatch, &procStopWatch);
              
              treeWriter->Fill();
              
              treeWriter->Clear();
              modularDelphes->Clear();
              converter->Clear();
              
              readStopWatch.Start();
            }
        }
      delete hepmcevt;

      ++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;
  }
}