LHEReader.cc

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#include <framework/gearbox/Unit.h>
#include <framework/gearbox/Const.h>
#include <framework/logging/Logger.h>
#include <generators/lhe/LHEReader.h>
#include <mdst/dataobjects/MCParticle.h>
 
#include <string>
#include <boost/lexical_cast.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>
 
using namespace std;
using namespace Belle2;
 
const boost::char_separator<char> LHEReader::sep(",; \t");
 
void LHEReader::open(const string& filename)
{
  m_lineNr = 0;
  m_input.open(filename.c_str());
  if (!m_input) throw(LHECouldNotOpenFileError() << filename);
}
 
int LHEReader::getEvent(MCParticleGraph& graph, double& eventWeight)
{
  int nparticles = readEventHeader(eventWeight);
  if (nparticles <= 0) {
    throw (LHEEmptyEventError() << m_lineNr << nparticles);
  }
 
  int first = graph.size();
  //Make list of particles
  for (int i = 0; i < nparticles; i++) {
    graph.addParticle();
  }
 
  //Read particles from file
  for (int i = 0; i < nparticles; ++i) {
    MCParticleGraph::GraphParticle& p = graph[first + i];
    int mother = readParticle(p);
 
    // add the mother
    if (mother > 0) {
      MCParticleGraph::GraphParticle* q = &graph[mother - 1];
      p.comesFrom(*q);
    }
 
    // if positron goes to positive z-direction, we have to rotate along y-axis by 180deg
    if (m_wrongSignPz) {
      ROOT::Math::PxPyPzEVector p4 = p.get4Vector();
      p4.SetPz(-1.0 * p4.Pz());
      p4.SetPx(-1.0 * p4.Px());
      p.set4Vector(p4);
    }
 
    // initial 2 (e+/e-), virtual 3 (Z/gamma*)
    // check if particle should be made virtual according to steering options:
    if (i < m_indexVirtual && i >= m_indexInitial)
      p.addStatus(MCParticle::c_IsVirtual);
 
    if (i < m_indexInitial)
      p.addStatus(MCParticle::c_Initial);
 
    if (m_indexVirtual < m_indexInitial) B2WARNING("IsVirtual particle requested but is overwritten by Initial");
 
  }
  return -1;
}
 
 
bool LHEReader::skipEvents(int n)
{
  double weight;
  for (int i = 0; i < n; i++) {
    int nparticles = readEventHeader(weight);
    if (nparticles < 0) return false;
    for (int j = 0; j < nparticles; j++) getLine();
  }
  return true;
}
 
int LHEReader::countEvents(const std::string& filename)
{
  std::ifstream input(filename.c_str());
  if (!input) throw(LHECouldNotOpenFileError() << filename);
 
  int count = 0;
  int lineNr = 0;
  std::string line;
 
  auto countEvent = [&]() {
    count++;
    while (std::getline(input, line)) {
      lineNr++;
      boost::trim(line);
      if (line == "</event>" || line.empty()) break;
    }
  };
 
  while (std::getline(input, line)) {
    lineNr++;
    boost::trim(line);
    if (line == "<event>") {
      countEvent();
    }
  }
 
  B2INFO("Counted " << count << " events in " << filename << ".");
  return count;
}
 
//===================================================================
//                  Protected methods
//===================================================================
 
std::string LHEReader::getLine()
{
  std::string line;
  do {
    getline(m_input, line);
    m_lineNr++;
    size_t commentPos = line.find_first_of('#');
    if (commentPos != string::npos) {
      line = line.substr(0, commentPos);
    }
    boost::trim(line);
 
  } while (line == "" && !m_input.eof());
 
  return line;
}
 
int LHEReader::readEventHeader(double& eventWeight)
{
 
  // Search for next <event>
  std::string line2;
  do {
    getline(m_input, line2);
    m_lineNr++;
    size_t commentPos = line2.find_first_of('#');
    if (commentPos != string::npos) {
      line2 = line2.substr(0, commentPos);
    }
    boost::trim(line2);
 
  } while (line2 != "<event>" && !m_input.eof());
 
  //Get number of particles from file
  int nparticles = -1;
  string line = getLine();
 
  if (line == "" || m_input.eof()) return -1;
 
  vector<double> fields;
  fields.reserve(15);
 
  tokenizer tokens(line, sep);
  int index(0);
 
  for (const string& tok : tokens) {
    ++index;
    try {
      fields.push_back(boost::lexical_cast<double>(tok));
    } catch (boost::bad_lexical_cast& e) {
      throw (LHEConvertFieldError() << m_lineNr << index << tok);
    }
  }
 
  if (fields.size() < 3) {
    nparticles = static_cast<int>(fields[0]);
    eventWeight = 1.0;
  } else {
    nparticles = static_cast<int>(fields[0]);
    eventWeight = static_cast<double>(fields[2]);
  }
  return nparticles;
}
 
int LHEReader::readParticle(MCParticleGraph::GraphParticle& particle)
{
  int mother = -1;
 
  string line = getLine();
  vector<double> fields;
  fields.reserve(15);
 
  tokenizer tokens(line, sep);
  int index(0);
 
  for (const string& tok : tokens) {
    ++index;
    try {
      fields.push_back(boost::lexical_cast<double>(tok));
    } catch (boost::bad_lexical_cast& e) {
      throw (LHEConvertFieldError() << m_lineNr << index << tok);
    }
  }
 
  switch (fields.size()) {
    case 13:
      particle.addStatus(MCParticle::c_PrimaryParticle);
      particle.setPDG(static_cast<int>(fields[0]));
      mother = static_cast<int>(fields[2]);
      particle.setMomentum(ROOT::Math::XYZVector(fields[6], fields[7], fields[8]));
      particle.setEnergy(fields[9]);
      particle.setMass(fields[10]);
      break;
    default:
      throw (LHEParticleFormatError() << m_lineNr << fields.size());
  }
 
  return mother;
}