HepMC/src/HepMCProcessExtractor.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/HepMC/HepMCProcessExtractor.hpp"
0010
0011 #include "ActsExamples/Framework/WhiteBoard.hpp"
0012 #include "ActsExamples/Io/HepMC3/HepMC3Particle.hpp"
0013
0014 #include <stdexcept>
0015
0016 #include <HepMC3/GenEvent.h>
0017 #include <HepMC3/GenParticle.h>
0018 #include <HepMC3/GenVertex.h>
0019
0020 namespace {
0021
0022 /// @brief This method searches for an outgoing particle from a vertex
0023 ///
0024 /// @param [in] vertex The vertex
0025 /// @param [in] id The track ID of the particle
0026 ///
0027 /// @return The particle pointer if found, else nullptr
0028 HepMC3::ConstGenParticlePtr searchProcessParticleById(
0029     const HepMC3::ConstGenVertexPtr& vertex, const int id) {
0030   // Loop over all outgoing particles
0031   for (const auto& particle : vertex->particles_out()) {
0032     const int trackid =
0033         particle->attribute<HepMC3::IntAttribute>("TrackID")->value();
0034     // Compare ID
0035     if (trackid == id) {
0036       return particle;
0037     }
0038   }
0039   return nullptr;
0040 }
0041
0042 /// @brief This method collects the material in X_0 and L_0 a particle has
0043 /// passed from its creation up to a certain vertex.
0044 ///
0045 /// @param [in] vertex The end vertex of the collection
0046 /// @param [in] id The track ID
0047 /// @param [in, out] particle The particle that get the passed material attached
0048 void setPassedMaterial(const HepMC3::ConstGenVertexPtr& vertex, const int id,
0049                        ActsExamples::SimParticle& particle) {
0050   double x0 = 0.;
0051   double l0 = 0.;
0052   HepMC3::ConstGenParticlePtr currentParticle = nullptr;
0053   HepMC3::ConstGenVertexPtr currentVertex = vertex;
0054   // Loop backwards and test whether the track still exists
0055   while (currentVertex && !currentVertex->particles_in().empty() &&
0056          currentVertex->particles_in()[0]->attribute<HepMC3::IntAttribute>(
0057              "TrackID") &&
0058          currentVertex->particles_in()[0]
0059                  ->attribute<HepMC3::IntAttribute>("TrackID")
0060                  ->value() == id) {
0061     // Get the step length
0062     currentParticle = currentVertex->particles_in()[0];
0063     const double stepLength =
0064         currentParticle->attribute<HepMC3::DoubleAttribute>("StepLength")
0065             ->value();
0066     // Add the passed material
0067     x0 +=
0068         stepLength /
0069         currentParticle->attribute<HepMC3::DoubleAttribute>("NextX0")->value();
0070     l0 +=
0071         stepLength /
0072         currentParticle->attribute<HepMC3::DoubleAttribute>("NextL0")->value();
0073     currentVertex = currentParticle->production_vertex();
0074   }
0075   // Assign the passed material to the particle
0076   particle.setMaterialPassed(x0, l0);
0077 }
0078
0079 /// @brief This function collects outgoing particles from a vertex while keeping
0080 /// track of the future of the ingoing particle.
0081 ///
0082 /// @param [in] vertex The vertex
0083 /// @param [in] trackID The track ID of the ingoing particle
0084 ///
0085 /// @return Vector containing the outgoing particles from a vertex
0086 std::vector<ActsExamples::SimParticle> selectOutgoingParticles(
0087     const HepMC3::ConstGenVertexPtr& vertex, const int trackID) {
0088   std::vector<ActsExamples::SimParticle> finalStateParticles;
0089
0090   // Identify the ingoing particle in the outgoing particles
0091   HepMC3::ConstGenParticlePtr procPart =
0092       searchProcessParticleById(vertex, trackID);
0093
0094   // Test whether this particle survives or dies
0095   HepMC3::ConstGenVertexPtr endVertex = procPart->end_vertex();
0096   if (endVertex
0097           ->attribute<HepMC3::StringAttribute>("NextProcessOf-" +
0098                                                std::to_string(trackID))
0099           ->value() != "Death") {
0100     // Store the particle if it survives
0101     finalStateParticles.push_back(
0102         ActsExamples::HepMC3Particle::particle(procPart));
0103   } else {
0104     // Store the leftovers if it dies
0105     for (const HepMC3::ConstGenParticlePtr& procPartOut :
0106          endVertex->particles_out()) {
0107       if (procPartOut->attribute<HepMC3::IntAttribute>("TrackID")->value() ==
0108               trackID &&
0109           procPartOut->end_vertex()) {
0110         for (const HepMC3::ConstGenParticlePtr& dyingPartOut :
0111              procPartOut->end_vertex()->particles_out()) {
0112           finalStateParticles.push_back(
0113               ActsExamples::HepMC3Particle::particle(dyingPartOut));
0114         }
0115       }
0116     }
0117   }
0118
0119   // Record the particles produced in this process
0120   const std::vector<std::string> attributes = endVertex->attribute_names();
0121   for (const auto& att : attributes) {
0122     // Search for initial parameters
0123     if (att.find("InitialParametersOf") != std::string::npos) {
0124       const std::vector<double> mom4 =
0125           endVertex->attribute<HepMC3::VectorDoubleAttribute>(att)->value();
0126       const HepMC3::FourVector& pos4 = endVertex->position();
0127       const int id = stoi(att.substr(att.find("-") + 1));
0128       HepMC3::ConstGenParticlePtr genParticle =
0129           searchProcessParticleById(endVertex, id);
0130       ActsFatras::Barcode barcode = ActsFatras::Barcode().setParticle(id);
0131       auto pid = static_cast<Acts::PdgParticle>(genParticle->pid());
0132
0133       // Build an Acts particle out of the data
0134       ActsExamples::SimParticle simParticle(barcode, pid);
0135       simParticle.setPosition4(pos4.x(), pos4.y(), pos4.z(), pos4.t());
0136       Acts::Vector3 mom3(mom4[0], mom4[1], mom4[2]);
0137       simParticle.setDirection(mom3.normalized());
0138       simParticle.setAbsoluteMomentum(mom3.norm());
0139
0140       // Store the particle
0141       finalStateParticles.push_back(simParticle);
0142     }
0143   }
0144
0145   return finalStateParticles;
0146 }
0147
0148 /// @brief This method filters and sorts the recorded interactions.
0149 ///
0150 /// @param [in] cfg Configuration of the filtering
0151 /// @param [in, out] interactions The recorded interactions
0152 void filterAndSort(
0153     const ActsExamples::HepMCProcessExtractor::Config& cfg,
0154     ActsExamples::ExtractedSimulationProcessContainer& interactions) {
0155   for (auto& interaction : interactions) {
0156     for (auto cit = interaction.after.cbegin();
0157          cit != interaction.after.cend();) {
0158       // Test whether a particle fulfills the conditions
0159       if (cit->pdg() < cfg.absPdgMin || cit->pdg() > cfg.absPdgMax ||
0160           cit->absoluteMomentum() < cfg.pMin) {
0161         interaction.after.erase(cit);
0162       } else {
0163         cit++;
0164       }
0165     }
0166   }
0167
0168   // Sort the particles based on their momentum
0169   for (auto& interaction : interactions) {
0170     std::sort(interaction.after.begin(), interaction.after.end(),
0171               [](ActsExamples::SimParticle& a, ActsExamples::SimParticle& b) {
0172                 return a.absoluteMomentum() > b.absoluteMomentum();
0173               });
0174   }
0175 }
0176 }  // namespace
0177
0178 ActsExamples::HepMCProcessExtractor::~HepMCProcessExtractor() = default;
0179
0180 ActsExamples::HepMCProcessExtractor::HepMCProcessExtractor(
0181     ActsExamples::HepMCProcessExtractor::Config config,
0182     Acts::Logging::Level level)
0183     : ActsExamples::IAlgorithm("HepMCProcessExtractor", level),
0184       m_cfg(std::move(config)) {
0185   if (m_cfg.inputEvents.empty()) {
0186     throw std::invalid_argument("Missing input event collection");
0187   }
0188   if (m_cfg.outputSimulationProcesses.empty()) {
0189     throw std::invalid_argument("Missing output collection");
0190   }
0191   if (m_cfg.extractionProcess.empty()) {
0192     throw std::invalid_argument("Missing extraction process");
0193   }
0194
0195   m_inputEvents.initialize(m_cfg.inputEvents);
0196   m_outputSimulationProcesses.initialize(m_cfg.outputSimulationProcesses);
0197 }
0198
0199 ActsExamples::ProcessCode ActsExamples::HepMCProcessExtractor::execute(
0200     const ActsExamples::AlgorithmContext& context) const {
0201   // Retrieve the initial particles
0202   const auto events = m_inputEvents(context);
0203
0204   ActsExamples::ExtractedSimulationProcessContainer fractions;
0205   for (const HepMC3::GenEvent& event : events) {
0206     // Fast exit
0207     if (event.particles().empty() || event.vertices().empty()) {
0208       break;
0209     }
0210
0211     // Get the initial particle
0212     HepMC3::ConstGenParticlePtr initialParticle = event.particles()[0];
0213     ActsExamples::SimParticle simParticle =
0214         HepMC3Particle::particle(initialParticle);
0215
0216     // Get the final state particles
0217     ActsExamples::SimParticle particleToInteraction;
0218     std::vector<ActsExamples::SimParticle> finalStateParticles;
0219     // Search the process vertex
0220     bool vertexFound = false;
0221     for (const auto& vertex : event.vertices()) {
0222       const std::vector<std::string> attributes = vertex->attribute_names();
0223       for (const auto& attribute : attributes) {
0224         if (vertex->attribute_as_string(attribute).find(
0225                 m_cfg.extractionProcess) != std::string::npos) {
0226           const int procID = stoi(attribute.substr(attribute.find("-") + 1));
0227           // Get the particle before the interaction
0228           particleToInteraction =
0229               HepMC3Particle::particle(vertex->particles_in()[0]);
0230           // Attach passed material to the particle
0231           setPassedMaterial(vertex, procID, particleToInteraction);
0232           // Record the final state particles
0233           finalStateParticles = selectOutgoingParticles(vertex, procID);
0234           vertexFound = true;
0235           break;
0236         }
0237       }
0238       if (vertexFound) {
0239         break;
0240       }
0241     }
0242     fractions.push_back(ActsExamples::ExtractedSimulationProcess{
0243         simParticle, particleToInteraction, finalStateParticles});
0244   }
0245
0246   // Filter and sort the record
0247   filterAndSort(m_cfg, fractions);
0248
0249   ACTS_INFO(events.size() << " processed");
0250
0251   // Write the recorded material to the event store
0252   m_outputSimulationProcesses(context, std::move(fractions));
0253
0254   return ActsExamples::ProcessCode::SUCCESS;
0255 }