Geant4Output2Podio.cxx

#include "Geant4Output2Podio.h"
#include "DD4hep/DD4hepUnits.h"
#include "DD4hep/InstanceCount.h"
#include "DD4hep/Primitives.h"
#include "DD4hep/Printout.h"
#include "DDG4/Geant4Data.h"
#include "DDG4/Geant4HitCollection.h"
#include "DDG4/Geant4Particle.h"

#include "TBuffer.h"

#include "CLHEP/Units/SystemOfUnits.h"
#include "G4ParticleDefinition.hh"
#include "G4ParticleTable.hh"
#include "G4IonTable.hh"
#include "G4Event.hh"

#include "dd4pod/TrackerHitCollection.h"
#include "dd4pod/CalorimeterHitCollection.h"
#include "dd4pod/Geant4ParticleCollection.h"

#include "dd4pod/PhotoMultiplierHitCollection.h"
#include "PMTHit.h"

// Geant4 include files
#include "G4HCofThisEvent.hh"

#include "G4Threading.hh"
#include "G4AutoLock.hh"

#include <atomic>

using std::string;

namespace {
  G4Mutex action_mutex = G4MUTEX_INITIALIZER;
}

namespace dd4hep::sim {

  /// Standard constructor
  Geant4Output2Podio::Geant4Output2Podio(Geant4Context* ctxt, const string& nam)
      : Geant4OutputAction(ctxt, nam) {
    declareProperty("RunHeader", m_runHeader);
    declareProperty("RunNumberOffset", m_runNumberOffset);
    declareProperty("EventNumberOffset", m_eventNumberOffset);
    declareProperty("HandleMCTruth", m_handleMCTruth);
    declareProperty("DisabledCollections", m_disabledCollections);
    declareProperty("DisableParticles", m_disableParticles);
    declareProperty("StoreCalorimeterContributions", m_storeCalorimeterContributions);
    declareProperty("EventParametersInt",    m_eventParametersInt);
    declareProperty("EventParametersFloat",  m_eventParametersFloat);
    declareProperty("EventParametersString", m_eventParametersString);
    printout(DEBUG, "Geant4Output2Podio", "Constructed");
    InstanceCount::increment(this);
  }

  /// Default destructor
  Geant4Output2Podio::~Geant4Output2Podio() {
    // @TODO: do we really have to mutex-lock the destructor?
    G4AutoLock protection_lock(&action_mutex);

    printout(INFO, "Geant4Output2Podio", "Closing output stream...");

    if (m_writer) {
      printout(DEBUG, "Geant4Output2Podio", "Closing and removing reader...");
      m_writer->finish();
      // not sure if we can't just let RAII take care of this @TODO
      m_writer.reset(); 
    }
    if (m_store) {
      printout(DEBUG, "Geant4Output2Podio", "Removing Event Store...");
      // not sure if we can't just let RAII take care of this @TODO
      m_store.reset();
    }
    InstanceCount::decrement(this);
  }

  /// Callback to store the Geant4 run information
  void Geant4Output2Podio::beginRun(const G4Run* run) {
    G4AutoLock protection_lock(&action_mutex);
    if (!m_writer && !m_output.empty()) {
      m_store  = std::make_unique<podio::EventStore>();
      m_writer = std::make_unique<podio::ROOTWriter>(m_output, m_store.get());
      printout(INFO, "Geant4Output2Podio", "opened %s for output", m_output.c_str());
      // initialize our enable/disable list
      for (const auto& name : m_disabledCollections) {
        m_disableList.insert(name);
      }
    }
    Geant4OutputAction::beginRun(run);
    m_particleTable = G4ParticleTable::GetParticleTable();
  }
  /// Callback to store the Geant4 run information
  void Geant4Output2Podio::endRun(const G4Run*  /*run*/) {
    // Do nothing
  }

  void Geant4Output2Podio::saveRun(const G4Run* /* run */) {
    // Do nothing
  }

  /// Commit data at end of filling procedure
  void Geant4Output2Podio::commit(OutputContext<G4Event>& /*ctxt*/) {
    if (m_writer)
      [[likely]] {
        G4AutoLock protection_lock(&action_mutex);
        m_writer->writeEvent();
        m_store->clearCollections();
      }
    else
      [[unlikely]] { except("+++ Failed to write output file / output stream not accessible"); }
  }

  /// Initialize the event store
  /// It is save if this function gets called mutliple times, as nothing will happen
  /// if the collections already exist
  void Geant4Output2Podio::initEventStore(OutputContext<G4Event>& ctxt) {
    printout(INFO, "Geant4Output2Podio", "Initializing the Event Store");

    // mcparticles
    if (!m_disableParticles && !m_store->getCollectionIDTable()->present("mcparticles")) {
      createCollection<dd4pod::Geant4ParticleCollection>("mcparticles");
    }

    // Now init the hit collections
    const G4Event* evt  = ctxt.context;
    const auto     hce  = evt->GetHCofThisEvent();
    const size_t   nCol = hce->GetNumberOfCollections();

    for (size_t i = 0; i < nCol; ++i) {
      const auto        hc   = hce->GetHC(i);
      const std::string name = hc->GetName();
      const auto        collection = dynamic_cast<Geant4HitCollection*>(hc);
      if (!collection) {
        printout(WARNING, "Geant4Output2Podio", "Skipping unknown GEANT4 hit collection: %s", name.c_str());
        continue;
      }
      if (m_disableList.count(name)) {
        printout(DEBUG, "Geant4Output2Podio", "Skipping hit collection (by user request): %s", name.c_str());
        continue;
      } 
      // Ensure we don't create the same collection twice
      if (m_store->getCollectionIDTable()->present(name)) {
        printout(DEBUG, "Geant4Output2Podio", "Skipping collection already on the data store: %s", name.c_str());
        continue;
      }
      // Tracker hit
      if (typeid(Geant4Tracker::Hit) == collection->type().type()) {
        createCollection<dd4pod::TrackerHitCollection>(name);
        // Calorimeter hit
      } else if (typeid(Geant4Calorimeter::Hit) == collection->type().type()) {
        createCollection<dd4pod::CalorimeterHitCollection>(name);
        // PMT hit
      } else if (typeid(npdet::PMTHit) == collection->type().type()) {
        createCollection<dd4pod::PhotoMultiplierHitCollection>(name);
      } else {
        printout(
            WARNING, "Geant4Output2Podio",
            "Encountered unknown hit type %s, collectionection will not be stored: %s",
            collection->type().type().name(), name.c_str());
      }
    }
  }

  // When this function is called we know we want to store mcparticles
  void Geant4Output2Podio::saveParticles(const Geant4ParticleMap* parts) {

    // Only store particles if not empty
    if (parts->particleMap.empty()) {
      return;
    }

    // create the output collection
    dd4pod::Geant4ParticleCollection* podcol =
        const_cast<dd4pod::Geant4ParticleCollection*>(
            &m_store->get<dd4pod::Geant4ParticleCollection>("mcparticles"));

    const auto& pm = parts->particles();

    // loop over particle and copy them to podio data model
    for (const auto& [key, part] : pm) {
      int charge = 0;
      // particles
      const auto part_def = m_particleTable->FindParticle(part->pdgID);
      // ions
      int Z{0}, A{0}, L{0}, J{0};
      double E{0.};
      const auto is_ion = G4IonTable::GetNucleusByEncoding(part->pdgID, Z, A, L, E, J);
      if (part_def != nullptr) {
        charge = static_cast<int>(std::lround(part_def->GetPDGCharge() / CLHEP::eplus));
      } else if (is_ion) {
        charge = Z;
      } else {
        printout((part->genStatus == 1 ? ERROR : INFO), "Geant4Output2Podio",
                 "++ saveEvent: pdgID %d unknown, charge set to 0", part->pdgID);
      }
      dd4pod::MutableGeant4Particle podpart{
          part->id,
          part->g4Parent,
          part->reason,
          part->mask,
          part->steps,
          part->secondaries,
          part->pdgID,
          part->status,
          {part->colorFlow[0], part->colorFlow[1]},
          part->genStatus,
          charge, // part->charge,
          part->_spare[0],
          {part->spin[0], part->spin[1], part->spin[2]},
          {part->vsx, part->vsy, part->vsz},
          {part->vex, part->vey, part->vez},
          {part->psx / 1000., part->psy / 1000., part->psz / 1000.},
          {part->pex / 1000., part->pey / 1000., part->pez / 1000.},
          part->mass / 1000.,
          part->time,
          part->properTime};
      for (const auto& ip : part->parents) {
        podpart.addparents(ip);
      }
      for (const auto& ip : part->daughters) {
        podpart.adddaughters(ip);
      }
      podcol->push_back(podpart);
    }
  }

  /// Callback to store the Geant4 event
  void Geant4Output2Podio::saveEvent(OutputContext<G4Event>& ctxt) {
    // @TODO evaluate if we need a lock here
    
    // init the collections on first event
    static bool first_event = true;
    if (first_event) [[unlikely]] {
      initEventStore(ctxt);
      first_event = false;
    }

    EventParameters* parameters = context()->event().extension<EventParameters>(false);
    int runNumber(0), eventNumber(0);
    const int eventNumberOffset(m_eventNumberOffset > 0 ? m_eventNumberOffset : 0);
    const int runNumberOffset(m_runNumberOffset > 0 ? m_runNumberOffset : 0);
    // Get event number, run number and parameters from extension ...
    if ( parameters ) {
      runNumber = parameters->runNumber() + runNumberOffset;
      eventNumber = parameters->eventNumber() + eventNumberOffset;
      parameters->extractParameters(*m_store);
    } else { // ... or from DD4hep framework
      runNumber = m_runNo + runNumberOffset;
      eventNumber = ctxt.context->GetEventID() + eventNumberOffset;
    }
    printout(INFO,"Geant4Output2Podio","+++ Saving Podio event %d run %d.", eventNumber, runNumber);

    saveEventParameters<int>(m_eventParametersInt);
    saveEventParameters<float>(m_eventParametersFloat);
    saveEventParameters<std::string>(m_eventParametersString);

    if (m_disableParticles) {
      return;
    }

    const Geant4ParticleMap* parts = context()->event().extension<Geant4ParticleMap>();

    saveParticles(parts);
  }

  /// Callback to store each Geant4 hit collection
  void Geant4Output2Podio::saveCollection(OutputContext<G4Event>& /* ctxt */, G4VHitsCollection* collection) {
    // @TODO evaluate if we need a lock here
    //
    const std::string          name    = collection->GetName();
    // Check if this is a collection we want to skip
    if (m_disableList.count(name)) {
      return;
    }
    const size_t  nhits = collection->GetSize();

    const Geant4HitCollection* coll = dynamic_cast<Geant4HitCollection*>(collection);
    if (!coll) [[unlikely]] {
      printout(ERROR, "Geant4Output2Podio", "No such Geant4 hit collection: %s", name.c_str());
      return;
    }

    const Geant4ParticleMap* parts = context()->event().extension<Geant4ParticleMap>(false);

    if (typeid(Geant4Calorimeter::Hit) == coll->type().type()) {
      auto podcol = const_cast<dd4pod::CalorimeterHitCollection*>(
          &m_store->get<dd4pod::CalorimeterHitCollection>(name));

      // copy all the hits
      for (size_t i = 0; i < nhits; ++i) {
        const Geant4Calorimeter::Hit* hit        = coll->hit(i);
        const auto&                   truth_list = hit->truth;

        auto ohit = podcol->create();
        ohit.cellID(hit->cellID);
        ohit.flag(hit->flag);
        ohit.g4ID(hit->g4ID);
        ohit.position({hit->position.x(), hit->position.y(), hit->position.z()});
        const auto& first_truth = truth_list[0];
        // Translate internal GEANT4 trackID to the related particleID in mcparticles
        const int trackID = parts->particleID(first_truth.trackID);
        ohit.truth(dd4pod::MonteCarloContrib(
            {trackID, first_truth.pdgID, first_truth.deposit / 1000.0, first_truth.time,
             first_truth.length, first_truth.x, first_truth.y, first_truth.z}));
        ohit.energyDeposit(hit->energyDeposit / 1000.0);
        // Optionally store all truth contribitions
        if (m_storeCalorimeterContributions) {
          for (const auto& truth : truth_list) {
            const int ID = parts->particleID(truth.trackID);

            dd4pod::MonteCarloContrib contrib{
                ID,         truth.pdgID,  truth.deposit / 1000.0,
                truth.time, truth.length, truth.x,
                truth.y,    truth.z};
            ohit.addcontributions(contrib);
          }
        }
      }
    } else if (typeid(Geant4Tracker::Hit) == coll->type().type()) {
      auto podcol = const_cast<dd4pod::TrackerHitCollection*>(
          &m_store->get<dd4pod::TrackerHitCollection>(name));
      // copy all the hits
      for (size_t i = 0; i < nhits; ++i) {
        const Geant4Tracker::Hit* hit   = coll->hit(i);
        const auto&               truth = hit->truth;

        auto ohit = podcol->create();
        ohit.cellID(hit->cellID);
        ohit.flag(hit->flag);
        ohit.g4ID(hit->g4ID);
        ohit.position({hit->position.x(), hit->position.y(), hit->position.z()});
        ohit.momentum({hit->momentum.x() / 1000.0, hit->momentum.y() / 1000.0,
                       hit->momentum.z() / 1000.0});
        ohit.length(hit->length);
        // Translate internal GEANT4 trackID to the related particleID in mcparticles
        const int trackID = parts->particleID(truth.trackID);
        ohit.truth({trackID, truth.pdgID, truth.deposit / 1000.0, truth.time,
                    truth.length, truth.x, truth.y, truth.z});
        ohit.energyDeposit(hit->energyDeposit / 1000.0);
      }
    } else if (typeid(npdet::PMTHit) ==coll->type().type()) {
      auto podcol = const_cast<dd4pod::PhotoMultiplierHitCollection*>(
          &m_store->get<dd4pod::PhotoMultiplierHitCollection>(name));
      // copy all the hits
      for (size_t i = 0; i < nhits; ++i) {
        const npdet::PMTHit* hit   = coll->hit(i);
        const auto&          truth = hit->truth;

        auto ohit = podcol->create();
        ohit.cellID(hit->cellID);
        ohit.flag(hit->flag);
        ohit.g4ID(hit->g4ID);
        ohit.position({hit->position.x(), hit->position.y(), hit->position.z()});
        ohit.momentum({hit->momentum.x() / 1000.0, hit->momentum.y() / 1000.0,
                       hit->momentum.z() / 1000.0});
        ohit.length(hit->length);
        // Translate internal GEANT4 trackID to the related particleID in mcparticles
        const int trackID = parts->particleID(truth.trackID);
        ohit.truth({trackID, truth.pdgID, truth.deposit / 1000.0, truth.time,
                    truth.length, truth.x, truth.y, truth.z});
        ohit.energyDeposit(hit->energyDeposit / 1000.0);
      }
    } else {
      printout(WARNING, "Geant4Output2Podio", 
               "Encountered unknown hit type %s, hit will not be stored: %s",
               coll->type().type().name(), name.c_str());
    }
  } 
} // namespace dd4hep::sim

using namespace dd4hep::sim;

#include "DDG4/Factories.h"

// clang-format off
DECLARE_GEANT4ACTION(Geant4Output2Podio)