Geant4HepMC/src/EventRecording.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2020 CERN for the benefit of the Acts project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Geant4HepMC/EventRecording.hpp"
0010
0011 #include "ActsExamples/EventData/SimParticle.hpp"
0012 #include "ActsExamples/Framework/WhiteBoard.hpp"
0013 #include "ActsExamples/Geant4/DetectorConstructionFactory.hpp"
0014 #include "ActsExamples/Geant4/GdmlDetectorConstruction.hpp"
0015
0016 #include <iostream>
0017 #include <stdexcept>
0018
0019 #include <FTFP_BERT.hh>
0020 #include <G4RunManager.hh>
0021 #include <G4VUserDetectorConstruction.hh>
0022 #include <HepMC3/GenParticle.h>
0023
0024 #include "EventAction.hpp"
0025 #include "PrimaryGeneratorAction.hpp"
0026 #include "RunAction.hpp"
0027 #include "SteppingAction.hpp"
0028
0029 ActsExamples::EventRecording::~EventRecording() {
0030   m_runManager = nullptr;
0031 }
0032
0033 ActsExamples::EventRecording::EventRecording(
0034     const ActsExamples::EventRecording::Config& config,
0035     Acts::Logging::Level level)
0036     : ActsExamples::IAlgorithm("EventRecording", level),
0037       m_cfg(config),
0038       m_runManager(std::make_unique<G4RunManager>()) {
0039   if (m_cfg.inputParticles.empty()) {
0040     throw std::invalid_argument("Missing input particle collection");
0041   }
0042   if (m_cfg.outputHepMcTracks.empty()) {
0043     throw std::invalid_argument("Missing output event collection");
0044   }
0045   if (!m_cfg.detectorConstructionFactory) {
0046     throw std::invalid_argument("Missing detector construction object");
0047   }
0048
0049   m_inputParticles.initialize(m_cfg.inputParticles);
0050   m_outputEvents.initialize(m_cfg.outputHepMcTracks);
0051
0052   // Now set up the Geant4 simulation
0053
0054   // G4RunManager deals with the lifetime of these objects
0055   m_runManager->SetUserInitialization(
0056       m_cfg.detectorConstructionFactory->factorize().release());
0057   m_runManager->SetUserInitialization(new FTFP_BERT);
0058   m_runManager->SetUserAction(new ActsExamples::Geant4::HepMC3::RunAction());
0059   m_runManager->SetUserAction(
0060       new ActsExamples::Geant4::HepMC3::EventAction(m_cfg.processesCombine));
0061   m_runManager->SetUserAction(
0062       new ActsExamples::Geant4::HepMC3::PrimaryGeneratorAction(m_cfg.seed1,
0063                                                                m_cfg.seed2));
0064   m_runManager->SetUserAction(
0065       new ActsExamples::Geant4::HepMC3::SteppingAction(m_cfg.processesReject));
0066   m_runManager->Initialize();
0067 }
0068
0069 ActsExamples::ProcessCode ActsExamples::EventRecording::execute(
0070     const ActsExamples::AlgorithmContext& context) const {
0071   // ensure exclusive access to the geant run manager
0072   std::lock_guard<std::mutex> guard(m_runManagerLock);
0073
0074   // Retrieve the initial particles
0075   const auto initialParticles = m_inputParticles(context);
0076
0077   // Storage of events that will be produced
0078   std::vector<HepMC3::GenEvent> events;
0079   events.reserve(initialParticles.size());
0080
0081   for (const auto& part : initialParticles) {
0082     // Prepare the particle gun
0083     ActsExamples::Geant4::HepMC3::PrimaryGeneratorAction::instance()
0084         ->prepareParticleGun(part);
0085
0086     // Begin with the simulation
0087     m_runManager->BeamOn(1);
0088
0089     // Test if the event was aborted
0090     if (Geant4::HepMC3::SteppingAction::instance()->eventAborted()) {
0091       continue;
0092     }
0093
0094     // Set event start time
0095     HepMC3::GenEvent event =
0096         ActsExamples::Geant4::HepMC3::EventAction::instance()->event();
0097     HepMC3::FourVector shift(0., 0., 0., part.time() / Acts::UnitConstants::mm);
0098     event.shift_position_by(shift);
0099
0100     // Set beam particle properties
0101     const Acts::Vector4 momentum4 =
0102         part.fourMomentum() / Acts::UnitConstants::GeV;
0103     HepMC3::FourVector beamMom4(momentum4[0], momentum4[1], momentum4[2],
0104                                 momentum4[3]);
0105     auto beamParticle = event.particles()[0];
0106     beamParticle->set_momentum(beamMom4);
0107     beamParticle->set_pid(part.pdg());
0108
0109     if (m_cfg.processSelect.empty()) {
0110       // Store the result
0111       events.push_back(std::move(event));
0112     } else {
0113       bool storeEvent = false;
0114       // Test if the event has a process of interest in it
0115       for (const auto& vertex : event.vertices()) {
0116         if (vertex->id() == -1) {
0117           vertex->add_particle_in(beamParticle);
0118         }
0119         const std::vector<std::string> vertexAttributes =
0120             vertex->attribute_names();
0121         for (const auto& att : vertexAttributes) {
0122           if ((vertex->attribute_as_string(att).find(m_cfg.processSelect) !=
0123                std::string::npos) &&
0124               !vertex->particles_in().empty() &&
0125               vertex->particles_in()[0]->attribute<HepMC3::IntAttribute>(
0126                   "TrackID") &&
0127               vertex->particles_in()[0]
0128                       ->attribute<HepMC3::IntAttribute>("TrackID")
0129                       ->value() == 1) {
0130             storeEvent = true;
0131             break;
0132           }
0133         }
0134         if (storeEvent) {
0135           break;
0136         }
0137       }
0138       // Store the result
0139       if (storeEvent) {
0140         // Remove vertices w/o outgoing particles and particles w/o production
0141         // vertices
0142         while (true) {
0143           bool sane = true;
0144           for (const auto& v : event.vertices()) {
0145             if (!v) {
0146               continue;
0147             }
0148             if (v->particles_out().empty()) {
0149               event.remove_vertex(v);
0150               sane = false;
0151             }
0152           }
0153           for (const auto& p : event.particles()) {
0154             if (!p) {
0155               continue;
0156             }
0157             if (!p->production_vertex()) {
0158               event.remove_particle(p);
0159               sane = false;
0160             }
0161           }
0162           if (sane) {
0163             break;
0164           }
0165         }
0166         events.push_back(std::move(event));
0167       }
0168     }
0169   }
0170
0171   ACTS_INFO(initialParticles.size() << " initial particles provided");
0172   ACTS_INFO(events.size() << " tracks generated");
0173
0174   // Write the recorded material to the event store
0175   m_outputEvents(context, std::move(events));
0176
0177   return ActsExamples::ProcessCode::SUCCESS;
0178 }