HepMC3/src/HepMC3Util.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 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 https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Io/HepMC3/HepMC3Util.hpp"
0010
0011 #include "Acts/Utilities/ScopedTimer.hpp"
0012
0013 #include <stdexcept>
0014
0015 #include <HepMC3/GenEvent.h>
0016 #include <HepMC3/GenParticle.h>
0017 #include <HepMC3/GenVertex.h>
0018
0019 namespace ActsExamples {
0020
0021 std::shared_ptr<HepMC3::GenEvent> HepMC3Util::mergeEvents(
0022     std::span<const HepMC3::GenEvent*> genEvents, const Acts::Logger& logger) {
0023   Acts::AveragingScopedTimer mergeTimer("Merging generator events", logger(),
0024                                         Acts::Logging::DEBUG);
0025
0026   std::vector<std::shared_ptr<HepMC3::GenParticle>> particles;
0027
0028   auto event = std::make_shared<HepMC3::GenEvent>();
0029   event->set_units(HepMC3::Units::GEV, HepMC3::Units::MM);
0030
0031   for (const auto& genEvent : genEvents) {
0032     auto sample = mergeTimer.sample();
0033     particles.clear();
0034     particles.reserve(genEvent->particles_size());
0035
0036     auto copyAttributes = [&](const auto& src, auto& dst) {
0037       for (const auto& attr : src.attribute_names()) {
0038         auto value = src.attribute_as_string(attr);
0039         dst.add_attribute(attr,
0040                           std::make_shared<HepMC3::StringAttribute>(value));
0041       }
0042     };
0043
0044     copyAttributes(*genEvent, *event);
0045
0046     // Add to combined event
0047     for (const auto& srcParticle : genEvent->particles()) {
0048       if (srcParticle->id() - 1 != static_cast<int>(particles.size())) {
0049         throw std::runtime_error("Particle id is not consecutive");
0050       }
0051       auto particle = std::make_shared<HepMC3::GenParticle>();
0052       particle->set_momentum(srcParticle->momentum());
0053       particle->set_generated_mass(srcParticle->generated_mass());
0054       particle->set_pid(srcParticle->pid());
0055       particle->set_status(srcParticle->status());
0056
0057       particles.push_back(particle);
0058       event->add_particle(particle);
0059
0060       copyAttributes(*srcParticle, *particle);
0061     }
0062
0063     for (const auto& srcVertex : genEvent->vertices()) {
0064       auto vertex = std::make_shared<HepMC3::GenVertex>(srcVertex->position());
0065       vertex->set_status(srcVertex->status());
0066
0067       event->add_vertex(vertex);
0068
0069       copyAttributes(*srcVertex, *vertex);
0070
0071       for (const auto& srcParticle : srcVertex->particles_in()) {
0072         const auto& particle = particles.at(srcParticle->id() - 1);
0073         vertex->add_particle_in(particle);
0074       }
0075       for (const auto& srcParticle : srcVertex->particles_out()) {
0076         const auto& particle = particles.at(srcParticle->id() - 1);
0077         vertex->add_particle_out(particle);
0078       }
0079     }
0080   }
0081
0082   return event;
0083 }
0084
0085 std::string_view HepMC3Util::compressionExtension(Compression compression) {
0086   switch (compression) {
0087     using enum Compression;
0088     case none:
0089       return "";
0090     case zlib:
0091       return ".gz";
0092     case lzma:
0093       return ".xz";
0094     case bzip2:
0095       return ".bz2";
0096     case zstd:
0097       return ".zst";
0098     default:
0099       throw std::invalid_argument{"Unknown compression value"};
0100   }
0101 }
0102
0103 std::span<const HepMC3Util::Compression>
0104 HepMC3Util::availableCompressionModes() {
0105   using enum Compression;
0106   static const auto values = []() -> std::vector<HepMC3Util::Compression> {
0107     return {
0108         none,
0109 #ifdef HEPMC3_Z_SUPPORT
0110         zlib,
0111 #endif
0112 #ifdef HEPMC3_LZMA_SUPPORT
0113         lzma,
0114 #endif
0115 #ifdef HEPMC3_BZ2_SUPPORT
0116         bzip2,
0117 #endif
0118 #ifdef HEPMC3_ZSTD_SUPPORT
0119         zstd,
0120 #endif
0121     };
0122   }();
0123   return values;
0124 }
0125
0126 std::ostream& HepMC3Util::operator<<(std::ostream& os,
0127                                      HepMC3Util::Compression compression) {
0128   switch (compression) {
0129     using enum HepMC3Util::Compression;
0130     case none:
0131       return os << "none";
0132     case zlib:
0133       return os << "zlib";
0134     case lzma:
0135       return os << "lzma";
0136     case bzip2:
0137       return os << "bzip2";
0138     case zstd:
0139       return os << "zstd";
0140     default:
0141       throw std::invalid_argument{"Unknown compression value"};
0142   }
0143 }
0144
0145 }  // namespace ActsExamples