BeamBkgGeneratorModule.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.                  *
 **************************************************************************/
 
// Own header.
#include <background/modules/BeamBkgGenerator/BeamBkgGeneratorModule.h>
 
// framework - DataStore
#include <framework/datastore/StoreArray.h>
#include <framework/datastore/StoreObjPtr.h>
 
// framework aux
#include <framework/gearbox/Unit.h>
#include <framework/gearbox/Const.h>
#include <framework/logging/Logger.h>
#include <framework/gearbox/GearDir.h>
 
// data objects
#include <framework/dataobjects/EventMetaData.h>
#include <mdst/dataobjects/MCParticle.h>
#include <generators/SAD/dataobjects/SADMetaHit.h>
 
// random generator
#include <TRandom.h>
 
// coordinates translation
#include <TGeoMatrix.h>
#include <generators/SAD/ReaderSAD.h>
 
using namespace std;
using namespace Belle2;
 
//-----------------------------------------------------------------
//                 Register module
//-----------------------------------------------------------------
 
REG_MODULE(BeamBkgGenerator);
 
//-----------------------------------------------------------------
//                 Implementation
//-----------------------------------------------------------------
 
BeamBkgGeneratorModule::BeamBkgGeneratorModule() : Module()
{
  // set module description
  setDescription("Beam background generator based on SAD files. "
                 "The generator picks up particles from the SAD file randomly "
                 "according to their rates. "
                 "Number of events is determined from 'realTime' and overall rate, and "
                 "the generator terminates the execution when this number is reached.");
 
  // Add parameters
  addParam("fileName", m_fileName, "name of the SAD file converted to root");
  addParam("treeName", m_treeName, "name of the TTree in the SAD file", string("sad"));
  addParam("ringName", m_ringName, "name of the superKEKB ring (LER or HER)");
  addParam("realTime", m_realTime,
           "equivalent superKEKB running time to generate sample [ns].");
}
 
BeamBkgGeneratorModule::~BeamBkgGeneratorModule()
{
}
 
void BeamBkgGeneratorModule::initialize()
{
  // register MCParticles
 
  StoreArray<MCParticle> mcParticles;
  mcParticles.registerInDataStore();
 
  StoreArray<SADMetaHit> sadHits;
  sadHits.registerInDataStore();
 
 
  // check steering parameters
 
  if (m_ringName != "LER" and m_ringName != "HER") {
    B2ERROR("ringName can only be LER or HER");
  }
  if (m_realTime <= 0) {
    B2ERROR("realTime must be positive");
  }
 
  // open SAD file and get TTree
 
  m_file = TFile::Open(m_fileName.c_str());
  if (!m_file) {
    B2ERROR(m_fileName << ": can't open file");
    return;
  }
  m_tree = (TTree*) m_file->Get(m_treeName.c_str());
  if (!m_tree) {
    B2ERROR(m_treeName << ": no such TTree in the SAD file");
    return;
  }
 
  m_tree->SetBranchAddress("s", &m_sad.s);
  m_tree->SetBranchAddress("x", &m_sad.x);
  m_tree->SetBranchAddress("px", &m_sad.px);
  m_tree->SetBranchAddress("y", &m_sad.y);
  m_tree->SetBranchAddress("py", &m_sad.py);
  m_tree->SetBranchAddress("E", &m_sad.E);
  m_tree->SetBranchAddress("rate", &m_sad.rate);
 
  // for SADMetaHit
  m_tree->SetBranchAddress("ss", &m_sad.ss);
  m_tree->SetBranchAddress("sraw", &m_sad.sraw);
  m_tree->SetBranchAddress("ssraw", &m_sad.ssraw);
  m_tree->SetBranchAddress("nturn", &m_sad.nturn);
  m_tree->SetBranchAddress("xraw", &m_sad.xraw);
  m_tree->SetBranchAddress("yraw", &m_sad.yraw);
  m_tree->SetBranchAddress("r", &m_sad.r);
  m_tree->SetBranchAddress("rr", &m_sad.rr);
  m_tree->SetBranchAddress("dp_over_p0", &m_sad.dp_over_p0);
  m_tree->SetBranchAddress("watt", &m_sad.watt);
 
  int numEntries = m_tree->GetEntries();
  if (numEntries <= 0) {
    B2ERROR("SAD tree is empty");
    return;
  }
 
  // construct vector of cumulative rates and determine number of events to generate
 
  m_rates.push_back(0);
  for (int i = 0; i < numEntries; i++) {
    m_tree->GetEntry(i);
    m_rates.push_back(m_sad.rate + m_rates.back());
    m_counters.push_back(0);
  }
 
  m_numEvents = gRandom->Poisson(m_rates.back() * m_realTime / Unit::s);
 
  // set rotation from SAD to Belle II
 
  if (m_ringName == "LER") {
    GearDir ring("/Detector/SuperKEKB/LER/");
    m_rotation.SetAngle(ring.getAngle("angle"));
    m_ring = 2;
  } else {
    GearDir ring("/Detector/SuperKEKB/HER/");
    m_rotation.SetAngle(ring.getAngle("angle"));
    m_ring = 1;
  }
 
  m_sectionOrdering.insert(m_sectionOrdering.end(), {1, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2});
 
  B2RESULT("BG rate: " << m_rates.back() / 1e6 << " MHz, events to generate: "
           << m_numEvents);
 
}
 
void BeamBkgGeneratorModule::beginRun()
{
}
 
void BeamBkgGeneratorModule::event()
{
 
  // check event counter
 
  if (m_eventCounter >= m_numEvents) {
    StoreObjPtr<EventMetaData> eventMetaData;
    eventMetaData->setEndOfData(); // stop event processing
    return;
  }
  m_eventCounter++;
 
  // get SAD entry
 
  int i = generateEntry();
  m_tree->GetEntry(i);
  m_counters[i]++;
 
  // make SADMetaHit
  int ring_section = -1;
  double ssraw = m_sad.ss;
  if (m_sad.ss < 0) ssraw += 3016.;
  int section = (int)(ssraw * 12 / 3016.);
  if ((unsigned)section < m_sectionOrdering.size()) ring_section = m_sectionOrdering[section];
 
  StoreArray<SADMetaHit> SADMetaHits;
  SADMetaHits.appendNew(SADMetaHit(m_sad.ssraw, m_sad.sraw, m_sad.ss, m_sad.s,
                                   0., m_sad.nturn,
                                   m_sad.x, m_sad.y, m_sad.px, m_sad.py, m_sad.xraw, m_sad.yraw,
                                   m_sad.r, m_sad.rr, m_sad.dp_over_p0, m_sad.E, m_sad.rate,
                                   m_sad.watt, m_ring, ring_section));
 
 
  // transform to Belle II (flip sign of y and s, rotate)
 
  ROOT::Math::XYZVector position(m_sad.x * Unit::m, -m_sad.y * Unit::m, -m_sad.s * Unit::m);
  position = m_rotation * position;
 
  double pz = sqrt(m_sad.E * m_sad.E - m_sad.px * m_sad.px - m_sad.py * m_sad.py);
  if (m_ringName == "LER") pz = -pz;
  ROOT::Math::XYZVector momentum(m_sad.px, -m_sad.py, pz);
  momentum = m_rotation * momentum;
 
  // PDG code and mass
 
  int pdgCode = Const::electron.getPDGCode();
  double mass = Const::electron.getMass();
  if (m_ringName == "LER") pdgCode = -pdgCode;
 
  // append SAD particle to MCParticles
 
  StoreArray<MCParticle> MCParticles;
  MCParticle* part = MCParticles.appendNew();
  part->setPDG(pdgCode);
  part->setMass(mass);
  part->setProductionVertex(position);
  part->setProductionTime(0);
  part->setMomentum(momentum);
  part->setEnergy(sqrt(momentum.Mag2() + mass * mass));
  part->setValidVertex(true);
  part->setStatus(MCParticle::c_PrimaryParticle);
  part->addStatus(MCParticle::c_StableInGenerator);
 
  // FarBeamLine region transformation
  if (abs(m_sad.s * Unit::m) > 4.0 * Unit::m) {
    // initial coordinates in SAD space
    double particlePosSADfar[] = {m_sad.x* Unit::m, -m_sad.y* Unit::m, 0.0 * Unit::m};
    double particleMomSADfar[] = {m_sad.px* Unit::GeV, -m_sad.py* Unit::GeV, pz* Unit::GeV};
    // final coordinates in Geant4 space
    double particlePosGeant4[] = {0.0, 0.0, 0.0};
    double particleMomGeant4[] = {0.0, 0.0, 0.0};
 
    // create a transformation matrix for a given ring
    TGeoHMatrix transMatrix; 
    if (m_ringName == "LER") {
      transMatrix = m_readerSAD.SADtoGeant(ReaderSAD::c_LER, m_sad.s * Unit::m);
    } else {
      transMatrix = m_readerSAD.SADtoGeant(ReaderSAD::c_HER, m_sad.s * Unit::m);
    }
 
    // calculate a new set of coordinates in Geant4 space
    transMatrix.LocalToMaster(particlePosSADfar, particlePosGeant4); // position
    transMatrix.LocalToMasterVect(particleMomSADfar, particleMomGeant4); // momentum
    // apply a new set of coordinates
    part->setMomentum(ROOT::Math::XYZVector(particleMomGeant4[0], particleMomGeant4[1], particleMomGeant4[2]));
    part->setProductionVertex(ROOT::Math::XYZVector(particlePosGeant4[0], particlePosGeant4[1], particlePosGeant4[2]));
  }
}
 
void BeamBkgGeneratorModule::endRun()
{
}
 
void BeamBkgGeneratorModule::terminate()
{
  // close SAD file
  m_file->Close();
 
  B2RESULT("BG rate: " << m_rates.back() / 1e6 << " MHz, equivalent time: "
           << m_realTime / Unit::us << " us, "
           << m_eventCounter << " events generated");
  if (m_eventCounter != m_numEvents)
    B2ERROR("Number of generated events does not match the equivalent running time: "
            << m_numEvents << " events needed, but "
            << m_eventCounter << " generated");
 
  int imax = 0;
  double average = 0;
  for (unsigned i = 0; i < m_counters.size(); i++) {
    if (m_counters[i] > m_counters[imax]) imax = i;
    average += m_counters[i];
  }
  average /= m_counters.size();
  B2RESULT("SAD particle usage: on average " << average << " times, max "
           << m_counters[imax] << " times (entry " << imax << ")");
 
}
 
 
int BeamBkgGeneratorModule::generateEntry() const
{
  double rate = gRandom->Uniform(m_rates.back());
  int i1 = 0;
  int i2 = m_rates.size() - 1;
  while (i2 - i1 > 1) {
    int i = (i1 + i2) / 2;
    if (rate > m_rates[i]) {
      i1 = i;
    } else {
      i2 = i;
    }
  }
  return i1;
}