PrintMCParticlesModule.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 <analysis/modules/PrintMCParticles/PrintMCParticlesModule.h>
#include <analysis/dataobjects/StringWrapper.h>
 
#include <mdst/dataobjects/MCParticle.h>
 
#include <framework/logging/LogConnectionConsole.h>
 
#include <boost/format.hpp>
#include <boost/algorithm/string.hpp>
 
#include <TDatabasePDG.h>
#include <Math/Vector3D.h>
 
using namespace Belle2;
 
//-----------------------------------------------------------------
//                 Register the Module
//-----------------------------------------------------------------
REG_MODULE(PrintMCParticles);
 
//-----------------------------------------------------------------
//                 Implementation
//-----------------------------------------------------------------
 
namespace {
  std::string statusToString(const MCParticle& mc)
  {
    static std::map<int, std::string> names {
      {MCParticle::c_PrimaryParticle, "PrimaryParticle"},
      {MCParticle::c_StableInGenerator, "StableInGenerator"},
      {MCParticle::c_LeftDetector, "LeftDetector"},
      {MCParticle::c_StoppedInDetector, "StoppedInDetector"},
      {MCParticle::c_IsVirtual, "IsVirtual"},
      {MCParticle::c_Initial, "Initial"},
      {MCParticle::c_IsISRPhoton, "IsISRPhoton"},
      {MCParticle::c_IsFSRPhoton, "IsFSRPhoton"},
      {MCParticle::c_IsPHOTOSPhoton, "IsPHOTOSPhoton"},
    };
    std::vector<std::string> set;
    for (auto&& [status, name] : names) {
      if (mc.hasStatus(status)) set.emplace_back(name);
    }
    return boost::join(set, ", ");
  }

  std::string g4ProcessToName(int process)
  {
    static std::map<int, std::string> translation{
      {1, "CoulombScattering"},
      {2, "Ionisation"},
      {3, "Bremsstrahlung"},
      {4, "PairProdByCharged"},
      {5, "Annihilation"},
      {6, "AnnihilationToMuMu"},
      {7, "AnnihilationToHadrons"},
      {8, "NuclearStopping"},
      {9, "ElectronGeneralProcess"},
      {10, "MultipleScattering"},
      {11, "Rayleigh"},
      {12, "PhotoElectricEffect"},
      {13, "ComptonScattering"},
      {14, "GammaConversion"},
      {15, "GammaConversionToMuMu"},
      {16, "GammaGeneralProcess"},
      {17, "PositronGeneralProcess"},
      {18, "AnnihilationToTauTau"},
      {21, "Cerenkov"},
      {22, "Scintillation"},
      {23, "SynchrotronRadiation"},
      {24, "TransitionRadiation"},
      {25, "SurfaceReflection"},
      {40, "DarkBremsstrahlung"},
      {49, "MuonPairProdByCharged"},
      {111, "HadronElastic"},
      {116, "NeutronGeneral"},
      {121, "HadronInelastic"},
      {131, "Capture"},
      {132, "MuAtomicCapture"},
      {141, "Fission"},
      {151, "HadronAtRest"},
      {152, "LeptonAtRest"},
      {161, "ChargeExchange"},
      // {210, "RadioactiveDecay"},
      {201, "Decay"},
      {202, "Decay_WithSpin"},
      {203, "Decay_PionMakeSpin"},
      {210, "Decay_Radioactive"},
      {211, "Decay_Unknown"},
      {221, "Decay_MuAtom "},
      {231, "Decay_External"},
      {310, "EMDissociation"},
    };
    if (auto it = translation.find(process); it != translation.end()) {
      return it->second;
    }
    return "[Unknown process]";
  }
}
 
PrintMCParticlesModule::PrintMCParticlesModule() : Module()
{
  //Set module properties
  setDescription("Print an MCParticle List");
  setPropertyFlags(c_ParallelProcessingCertified | c_TerminateInAllProcesses);
 
  //Parameter definition
  addParam("storeName", m_particleList, "Name of the StoreArray to print", m_particleList);
  addParam("onlyPrimaries", m_onlyPrimaries, "Show only primary particles", true);
  addParam("maxLevel", m_maxLevel, "Show only up to specified depth level; any value <= 0 means no limit (default: -1)", -1);
  addParam("showVertices", m_showVertices, "Show also the particle production vertices and times", false);
  addParam("showMomenta", m_showMomenta, "Show also the particle momenta", false);
  addParam("showProperties", m_showProperties, "Show the basic particle properties", false);
  addParam("showStatus", m_showStatus, "Show extendend status information of the particle", false);
  addParam("suppressPrint", m_suppressPrint, "Suppress print the information", false);
  addParam("storeCompact", m_storeCompact, "Store a compact single-line string in the ROOT branch instead of the full indented tree",
           false);
}
 
void PrintMCParticlesModule::initialize()
{
  m_mcparticles.isRequired(m_particleList);
  m_stringWrapper.registerInDataStore("MCDecayString");
}
 
void PrintMCParticlesModule::event()
{
  // clear any previous outputs
  m_output.str(std::string());
 
  if (not m_particleList.empty()) {
    m_output << "Content from MCParticle list '" << m_mcparticles.getName() << "'" << std::endl;
  } else {
    m_output << "Content of MCParticle list" << std::endl;
  }
 
  //Loop over the top level particles (no mother particle exists)
  std::vector<MCParticle*> first_gen;
  for (MCParticle& mc : m_mcparticles) {
    if (mc.getMother() != nullptr) continue;
    first_gen.emplace_back(&mc);
  }
  filterPrimaryOnly(first_gen);
 
  // Build the full tree if needed for log output or default (non-compact) storage
  if (not m_suppressPrint or not m_storeCompact) {
    printTree(first_gen);
  }
 
  if (!m_suppressPrint)
    B2INFO(m_output.str());
 
  if (not m_stringWrapper.isValid())
    m_stringWrapper.create();
 
  if (m_storeCompact) {
    std::stringstream compact;
    for (size_t i = 0; i < first_gen.size(); ++i) {
      if (i > 0) compact << " ";
      compact << formatParticleCompact(first_gen[i]);
    }
    m_stringWrapper->setString(compact.str());
  } else {
    m_stringWrapper->setString(std::string(m_output.str()));
  }
}
 
void PrintMCParticlesModule::filterPrimaryOnly(std::vector<MCParticle*>& particles) const
{
  if (!m_onlyPrimaries or particles.empty()) return;
  particles.erase(std::remove_if(particles.begin(), particles.end(), [](const MCParticle * mc) {
    return not mc->hasStatus(MCParticle::c_PrimaryParticle);
  }), particles.end());
}
 
void PrintMCParticlesModule::printTree(const std::vector<MCParticle*>& particles, int level, const std::string& indent)
{
  //If we show extra content make the particle name and pdg code bold if supported
  //And if we also show secondaries make those red to distinguish
  const bool useColor = LogConnectionConsole::terminalSupportsColors(STDOUT_FILENO);
  const bool anyExtraInfo = m_showProperties or m_showMomenta or m_showVertices or m_showStatus;
  std::string colorStart[] = {"", ""};
  std::string colorEnd = "";
  if (useColor) {
    // make it bold if we have extra info, otherwise no need
    colorStart[0] = anyExtraInfo ? "\x1b[31;1m" : "\x1b[31m";
    colorStart[1] = anyExtraInfo ? "\x1b[1m" : "";
    colorEnd = "\x1b[m";
  }
 
  TDatabasePDG* pdb = TDatabasePDG::Instance();
 
  // ok, go over all particles and print them
  const auto num_particles = particles.size();
  size_t particle_index{0};
  for (const auto* mc : particles) {
    // are we the last in the list? if so we need a different box char
    const bool last = ++particle_index == num_particles;
    m_output << indent << (last ? "╰" : "├");
    m_output << (mc->hasStatus(MCParticle::c_PrimaryParticle) ? "── " : "╶╶ ");
    // now print the name and pdg code. Optionally with color
    TParticlePDG* pdef = pdb->GetParticle(mc->getPDG());
    m_output << colorStart[mc->hasStatus(MCParticle::c_PrimaryParticle)] << (pdef ? pdef->GetTitle() : "[UNKNOWN]");
    m_output << " (" << mc->getPDG() << ")" << colorEnd;
 
    // Particle itself is done but we might need to show more things like properties and daughters so ...
    auto daughters = mc->getDaughters();
    filterPrimaryOnly(daughters);
    const bool showDaughters = (not daughters.empty()) and (m_maxLevel <= 0 or level < m_maxLevel);
 
    // new indent level for any information related to this particle.
    const std::string newIndent = indent + (last ? " " : "│") + "   ";
    // and the indent level for any properties associated to this particle
    const std::string propIndent = "\n" + newIndent + (showDaughters ? "│ " : "");
    // limited by level restriction? add an indicator that something is skipped
    if ((not showDaughters) and (not daughters.empty())) {
      m_output << " → …";
    }
    // Now show whatever the user wants
    if (m_showProperties) {
      m_output << propIndent;
      m_output << boost::format("mass=%.3g energy=%.3g charge=%d") % mc->getMass() % mc->getEnergy() % mc->getCharge();
      if (not mc->hasStatus(MCParticle::c_LeftDetector)) m_output << boost::format(" lifetime=%.3g") % mc->getLifetime();
    }
    if (m_showMomenta) {
      const ROOT::Math::XYZVector& p = mc->getMomentum();
      m_output << propIndent;
      m_output << boost::format("p=(%.3g, %.3g, %.3g) |p|=%.3g") % p.X() % p.Y() % p.Z() % p.R();
    }
    if (m_showVertices) {
      const ROOT::Math::XYZVector& v = mc->getVertex();
      m_output << propIndent;
      m_output << boost::format("production vertex=(%.3g, %.3g, %.3g), time=%.3g") % v.X() % v.Y() % v.Z() % mc->getProductionTime();
    }
    if (m_showStatus) {
      m_output << propIndent;
      m_output << "status flags=" << statusToString(*mc);
      if (not mc->hasStatus(MCParticle::c_PrimaryParticle)) {
        m_output << propIndent;
        m_output << "creation process=" << g4ProcessToName(mc->getSecondaryPhysicsProcess());
      }
      m_output << propIndent;
      m_output << "list index=" << mc->getIndex();
    }
    // and if we have shown any extra info and have daughters leave a blank line
    if (showDaughters and anyExtraInfo) m_output << propIndent;
    // now we're done with the particle ...
    m_output << std::endl;
    // but maybe we need to show the daughters
    if (showDaughters) {
      printTree(daughters, level + 1, newIndent);
    }
    // and if we show any extra info also leave blank line to sibling
    if (anyExtraInfo and not last) m_output << newIndent << std::endl;
  }
}
 
std::string PrintMCParticlesModule::formatParticleCompact(const MCParticle* mc, int level) const
{
  TDatabasePDG* pdb = TDatabasePDG::Instance();
  TParticlePDG* pdef = pdb->GetParticle(mc->getPDG());
  std::string name = pdef ? pdef->GetTitle() : "UNKNOWN(" + std::to_string(mc->getPDG()) + ")";
 
  if (name == "gamma") {
    if (mc->hasStatus(MCParticle::c_IsISRPhoton)) {
      name = "gammaI";
    } else if (mc->hasStatus(MCParticle::c_IsFSRPhoton)) {
      name = "gammaF";
    } else if (mc->hasStatus(MCParticle::c_IsPHOTOSPhoton)) {
      name = "gammaP";
    }
  }
 
  if (not mc->hasStatus(MCParticle::c_PrimaryParticle)) {
    name = "~" + name;
  }
 
  auto daughters = mc->getDaughters();
  filterPrimaryOnly(daughters);
 
  const bool showDaughters = (not daughters.empty()) and (m_maxLevel <= 0 or level < m_maxLevel);
 
  if (not showDaughters) {
    if (not daughters.empty()) {
      name += " -> ...";
    }
    return name;
  }
 
  std::string result = name + " -> ";
  for (size_t i = 0; i < daughters.size(); ++i) {
    if (i > 0) result += " ";
    const auto* daughter = daughters[i];
    std::string dstr = formatParticleCompact(daughter, level + 1);
 
    auto granddaughters = daughter->getDaughters();
    filterPrimaryOnly(granddaughters);
    if ((not granddaughters.empty()) and (m_maxLevel <= 0 or level + 1 < m_maxLevel)) {
      result += "[" + dstr + "]";
    } else {
      result += dstr;
    }
  }
 
  return result;
}