EDM4hep/src/EDM4hepReader.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/EDM4hep/EDM4hepReader.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Plugins/DD4hep/DD4hepDetectorElement.hpp"
0013 #include "Acts/Plugins/EDM4hep/EDM4hepUtil.hpp"
0014 #include "ActsExamples/DD4hepDetector/DD4hepDetector.hpp"
0015 #include "ActsExamples/EventData/GeometryContainers.hpp"
0016 #include "ActsExamples/EventData/SimHit.hpp"
0017 #include "ActsExamples/EventData/SimParticle.hpp"
0018 #include "ActsExamples/Framework/WhiteBoard.hpp"
0019 #include "ActsExamples/Io/EDM4hep/EDM4hepUtil.hpp"
0020 #include "ActsExamples/Utilities/Paths.hpp"
0021 #include "ActsFatras/EventData/Barcode.hpp"
0022
0023 #include <algorithm>
0024 #include <iomanip>
0025 #include <map>
0026 #include <stdexcept>
0027
0028 #include <DD4hep/Detector.h>
0029 #include <edm4hep/MCParticle.h>
0030 #include <edm4hep/SimTrackerHit.h>
0031 #include <edm4hep/SimTrackerHitCollection.h>
0032 #include <podio/Frame.h>
0033 #include <podio/ObjectID.h>
0034
0035 namespace ActsExamples {
0036
0037 EDM4hepReader::EDM4hepReader(const Config& config, Acts::Logging::Level level)
0038     : m_cfg(config),
0039       m_logger(Acts::getDefaultLogger("EDM4hepParticleReader", level)) {
0040   if (m_cfg.outputParticlesGenerator.empty()) {
0041     throw std::invalid_argument(
0042         "Missing output collection generator particles");
0043   }
0044   if (m_cfg.outputParticlesSimulation.empty()) {
0045     throw std::invalid_argument(
0046         "Missing output collection simulated particles");
0047   }
0048   if (m_cfg.outputSimHits.empty()) {
0049     throw std::invalid_argument("Missing output collection sim hits");
0050   }
0051
0052   m_eventsRange = std::make_pair(0, reader().getEntries("events"));
0053
0054   m_outputParticlesGenerator.initialize(m_cfg.outputParticlesGenerator);
0055   m_outputParticlesSimulation.initialize(m_cfg.outputParticlesSimulation);
0056   m_outputSimHits.initialize(m_cfg.outputSimHits);
0057
0058   m_cfg.trackingGeometry->visitSurfaces([&](const auto* surface) {
0059     const auto* detElement = dynamic_cast<const Acts::DD4hepDetectorElement*>(
0060         surface->associatedDetectorElement());
0061
0062     if (detElement == nullptr) {
0063       ACTS_ERROR("Surface has no associated detector element");
0064       return;
0065     }
0066
0067     const auto translation = detElement->sourceElement()
0068                                  .nominal()
0069                                  .worldTransformation()
0070                                  .GetTranslation();
0071     Acts::Vector3 position;
0072     position << translation[0], translation[1], translation[2];
0073     position *= Acts::UnitConstants::cm;
0074
0075     m_surfaceMap.insert({detElement->sourceElement().key(), surface});
0076   });
0077 }
0078
0079 Acts::PodioUtil::ROOTReader& EDM4hepReader::reader() {
0080   bool exists = false;
0081   auto& reader = m_reader.local(exists);
0082   if (!exists) {
0083     reader.openFile(m_cfg.inputPath);
0084   }
0085
0086   return reader;
0087 }
0088
0089 std::string EDM4hepReader::name() const {
0090   return "EDM4hepReader";
0091 }
0092
0093 std::pair<std::size_t, std::size_t> EDM4hepReader::availableEvents() const {
0094   return m_eventsRange;
0095 }
0096
0097 namespace {
0098 std::string vid(unsigned int vtx) {
0099   return "V" + std::to_string(vtx);
0100 }
0101
0102 std::string pid(const SimParticle& particle) {
0103   return "P" + std::to_string(particle.particleId().value());
0104 }
0105
0106 std::string plabel(const SimParticle& particle) {
0107   using namespace Acts::UnitLiterals;
0108   std::stringstream ss;
0109   ss << particle.pdg() << "\\n(" << particle.particleId() << ")\\n"
0110      << "p=" << std::setprecision(3) << particle.absoluteMomentum() / 1_GeV
0111      << " GeV";
0112   return ss.str();
0113 }
0114
0115 }  // namespace
0116
0117 void EDM4hepReader::graphviz(std::ostream& os,
0118                              const std::vector<SimParticle>& particles,
0119                              const ParentRelationship& parents) const {
0120   os << "digraph Event {\n";
0121
0122   std::set<unsigned int> primaryVertices;
0123
0124   for (const auto& particle : particles) {
0125     if (particle.particleId().generation() == 0) {
0126       primaryVertices.insert(particle.particleId().vertexPrimary());
0127
0128       os << vid(particle.particleId().vertexPrimary()) << " -> "
0129          << pid(particle) << ";\n";
0130     }
0131
0132     os << pid(particle) << " [label=\"" << plabel(particle) << "\"];\n";
0133   }
0134
0135   for (const auto [childIdx, parentIdx] : parents) {
0136     const auto& child = particles[childIdx];
0137     const auto& parent = particles[parentIdx];
0138     os << pid(parent) << " -> " << pid(child);
0139
0140     if (parent.particleId().vertexSecondary() ==
0141         child.particleId().vertexSecondary()) {
0142       os << " [style=dashed]";
0143     }
0144
0145     os << ";\n";
0146   }
0147
0148   for (unsigned int vtx : primaryVertices) {
0149     os << vid(vtx) << " [label=\"PV" << vtx << "\"];\n";
0150   }
0151
0152   os << "}";
0153 }
0154
0155 ProcessCode EDM4hepReader::read(const AlgorithmContext& ctx) {
0156   podio::Frame frame = reader().readEntry("events", ctx.eventNumber);
0157   const auto& mcParticleCollection =
0158       frame.get<edm4hep::MCParticleCollection>(m_cfg.inputParticles);
0159
0160   ACTS_DEBUG("Reading EDM4hep inputs");
0161
0162   std::vector<SimParticle> unorderedParticlesInitial;
0163
0164   // Read particles from the input file
0165   // Find particles without parents and group them by vtx position to find
0166   // primary vertices
0167   std::vector<std::pair<Acts::Vector3, std::vector<edm4hep::MCParticle>>>
0168       primaryVertices;
0169   for (const auto& particle : mcParticleCollection) {
0170     if (particle.parents_size() > 0) {
0171       // not a primary vertex
0172       continue;
0173     }
0174     const auto& vtx = particle.getVertex();
0175     Acts::Vector3 vtxPos = {vtx[0], vtx[1], vtx[2]};
0176     vtxPos /= Acts::UnitConstants::mm;
0177
0178     // linear search for vector
0179     auto it = std::ranges::find_if(primaryVertices, [&vtxPos](const auto& v) {
0180       return v.first == vtxPos;
0181     });
0182
0183     if (it == primaryVertices.end()) {
0184       ACTS_DEBUG("Found primary vertex at " << vtx.x << ", " << vtx.y << ", "
0185                                             << vtx.z);
0186       primaryVertices.push_back({vtxPos, {particle}});
0187     } else {
0188       it->second.push_back(particle);
0189     }
0190   }
0191
0192   ACTS_DEBUG("Found " << primaryVertices.size() << " primary vertices");
0193
0194   // key: child, value: parent
0195   ParentRelationship parentRelationship;
0196
0197   // key: input particle index, value: index in the unordered particle
0198   // container
0199   std::unordered_map<int, std::size_t> edm4hepParticleMap;
0200
0201   std::size_t nPrimaryVertices = 0;
0202   // Walk the particle tree
0203   for (const auto& [vtxPos, particles] : primaryVertices) {
0204     nPrimaryVertices += 1;
0205     ACTS_DEBUG("Walking particle tree for primary vertex at "
0206                << vtxPos.x() << ", " << vtxPos.y() << ", " << vtxPos.z());
0207     std::size_t nParticles = 0;
0208     std::size_t nSecondaryVertices = 0;
0209     std::size_t maxGen = 0;
0210     auto startSize = unorderedParticlesInitial.size();
0211     for (const auto& inParticle : particles) {
0212       nParticles += 1;
0213       SimParticle particle =
0214           EDM4hepUtil::readParticle(inParticle)
0215               .withParticleId(SimBarcode{}
0216                                   .setParticle(nParticles)
0217                                   .setVertexPrimary(nPrimaryVertices));
0218       ACTS_VERBOSE("+ add particle " << particle);
0219       ACTS_VERBOSE("  - at " << particle.position().transpose());
0220       ACTS_VERBOSE("  - createdInSim: " << inParticle.isCreatedInSimulation());
0221       ACTS_VERBOSE("  - vertexIsNotEndpointOfParent: "
0222                    << inParticle.vertexIsNotEndpointOfParent());
0223       ACTS_VERBOSE("  - isStopped: " << inParticle.isStopped());
0224       ACTS_VERBOSE("  - endpoint: " << inParticle.getEndpoint().x << ", "
0225                                     << inParticle.getEndpoint().y << ", "
0226                                     << inParticle.getEndpoint().z);
0227       const auto pid = particle.particleId();
0228       unorderedParticlesInitial.push_back(std::move(particle));
0229       edm4hepParticleMap[inParticle.getObjectID().index] =
0230           unorderedParticlesInitial.size() - 1;
0231       processChildren(inParticle, pid, unorderedParticlesInitial,
0232                       parentRelationship, edm4hepParticleMap,
0233                       nSecondaryVertices, maxGen);
0234     }
0235     ACTS_VERBOSE("Primary vertex complete, produced "
0236                  << (unorderedParticlesInitial.size() - startSize)
0237                  << " particles and " << nSecondaryVertices
0238                  << " secondary vertices in " << maxGen << " generations");
0239     setSubParticleIds(std::next(unorderedParticlesInitial.begin(), startSize),
0240                       unorderedParticlesInitial.end());
0241   }
0242
0243   ACTS_DEBUG("Found " << unorderedParticlesInitial.size() << " particles");
0244
0245   std::vector<SimParticle> particlesGeneratorUnordered;
0246   particlesGeneratorUnordered.reserve(mcParticleCollection.size());
0247   std::vector<SimParticle> particlesSimulatedUnordered;
0248   particlesSimulatedUnordered.reserve(mcParticleCollection.size());
0249
0250   for (const auto& inParticle : mcParticleCollection) {
0251     auto particleIt = edm4hepParticleMap.find(inParticle.getObjectID().index);
0252     if (particleIt == edm4hepParticleMap.end()) {
0253       ACTS_ERROR("Particle " << inParticle.getObjectID().index
0254                              << " not found in particle map");
0255       continue;
0256     }
0257     const std::size_t index = particleIt->second;
0258     const auto& particleInitial = unorderedParticlesInitial.at(index);
0259     if (!inParticle.isCreatedInSimulation()) {
0260       particlesGeneratorUnordered.push_back(particleInitial);
0261     }
0262     SimParticle particleSimulated = particleInitial;
0263
0264     float time = inParticle.getTime() * Acts::UnitConstants::ns;
0265     for (const auto& daughter : inParticle.getDaughters()) {
0266       if (!daughter.vertexIsNotEndpointOfParent()) {
0267         time = daughter.getTime() * Acts::UnitConstants::ns;
0268         break;
0269       }
0270     }
0271
0272     particleSimulated.final().setPosition4(
0273         inParticle.getEndpoint()[0] * Acts::UnitConstants::mm,
0274         inParticle.getEndpoint()[1] * Acts::UnitConstants::mm,
0275         inParticle.getEndpoint()[2] * Acts::UnitConstants::mm, time);
0276
0277     Acts::Vector3 momentumFinal = {inParticle.getMomentumAtEndpoint()[0],
0278                                    inParticle.getMomentumAtEndpoint()[1],
0279                                    inParticle.getMomentumAtEndpoint()[2]};
0280     particleSimulated.final().setDirection(momentumFinal.normalized());
0281     particleSimulated.final().setAbsoluteMomentum(momentumFinal.norm());
0282
0283     ACTS_VERBOSE("- Updated particle initial -> final, position: "
0284                  << particleInitial.fourPosition().transpose() << " -> "
0285                  << particleSimulated.final().fourPosition().transpose());
0286     ACTS_VERBOSE("                                     momentum: "
0287                  << particleInitial.fourMomentum().transpose() << " -> "
0288                  << particleSimulated.final().fourMomentum().transpose());
0289
0290     particlesSimulatedUnordered.push_back(particleSimulated);
0291   }
0292
0293   std::ranges::sort(particlesGeneratorUnordered, detail::CompareParticleId{});
0294   std::ranges::sort(particlesSimulatedUnordered, detail::CompareParticleId{});
0295
0296   SimParticleContainer particlesGenerator{particlesGeneratorUnordered.begin(),
0297                                           particlesGeneratorUnordered.end()};
0298   SimParticleContainer particlesSimulated{particlesSimulatedUnordered.begin(),
0299                                           particlesSimulatedUnordered.end()};
0300
0301   if (!m_cfg.graphvizOutput.empty()) {
0302     std::string path = perEventFilepath(m_cfg.graphvizOutput, "particles.dot",
0303                                         ctx.eventNumber);
0304     std::ofstream dot(path);
0305     graphviz(dot, unorderedParticlesInitial, parentRelationship);
0306   }
0307
0308   std::vector<SimHit> simHitsUnordered;
0309
0310   ACTS_DEBUG("Reading sim hits from " << m_cfg.inputSimHits.size()
0311                                       << " sim hit collections");
0312   for (const auto& name : m_cfg.inputSimHits) {
0313     const auto& inputHits = frame.get<edm4hep::SimTrackerHitCollection>(name);
0314
0315     simHitsUnordered.reserve(simHitsUnordered.size() + inputHits.size());
0316
0317     for (const auto& hit : inputHits) {
0318       auto simHit = EDM4hepUtil::readSimHit(
0319           hit,
0320           [&](const auto& inParticle) {
0321             ACTS_VERBOSE(
0322                 "SimHit has source particle: "
0323                 << Acts::EDM4hepUtil::getParticle(hit).getObjectID().index);
0324             auto it = edm4hepParticleMap.find(inParticle.getObjectID().index);
0325             if (it == edm4hepParticleMap.end()) {
0326               ACTS_ERROR(
0327                   "SimHit has source particle that we did not see before");
0328               return SimBarcode{};
0329             }
0330             const auto& particle = unorderedParticlesInitial.at(it->second);
0331             ACTS_VERBOSE("- " << inParticle.getObjectID().index << " -> "
0332                               << particle.particleId());
0333             return particle.particleId();
0334           },
0335           [&](std::uint64_t cellId) {
0336             ACTS_VERBOSE("CellID: " << cellId);
0337
0338             const auto& vm =
0339                 m_cfg.dd4hepDetector->dd4hepDetector().volumeManager();
0340
0341             const auto detElement = vm.lookupDetElement(cellId);
0342
0343             ACTS_VERBOSE(" -> detElement: " << detElement.name());
0344             ACTS_VERBOSE("   -> id: " << detElement.id());
0345             ACTS_VERBOSE("   -> key: " << detElement.key());
0346
0347             Acts::Vector3 position;
0348             position << detElement.nominal()
0349                             .worldTransformation()
0350                             .GetTranslation()[0],
0351                 detElement.nominal().worldTransformation().GetTranslation()[1],
0352                 detElement.nominal().worldTransformation().GetTranslation()[2];
0353             position *= Acts::UnitConstants::cm;
0354
0355             ACTS_VERBOSE("   -> detElement position: " << position.transpose());
0356
0357             auto it = m_surfaceMap.find(detElement.key());
0358             if (it == m_surfaceMap.end()) {
0359               ACTS_ERROR("Unable to find surface for detElement "
0360                          << detElement.name() << " with cellId " << cellId);
0361               throw std::runtime_error("Unable to find surface for detElement");
0362             }
0363             const auto* surface = it->second;
0364             if (surface == nullptr) {
0365               ACTS_ERROR("Unable to find surface for detElement "
0366                          << detElement.name() << " with cellId " << cellId);
0367               throw std::runtime_error("Unable to find surface for detElement");
0368             }
0369             ACTS_VERBOSE("   -> surface: " << surface->geometryId());
0370             return surface->geometryId();
0371           });
0372
0373       simHitsUnordered.push_back(std::move(simHit));
0374     }
0375   }
0376
0377   std::ranges::sort(simHitsUnordered, detail::CompareGeometryId{});
0378
0379   SimHitContainer simHits{simHitsUnordered.begin(), simHitsUnordered.end()};
0380
0381   if (m_cfg.sortSimHitsInTime) {
0382     ACTS_DEBUG("Sorting sim hits in time");
0383     std::multimap<ActsFatras::Barcode, std::size_t> hitsByParticle;
0384
0385     for (std::size_t i = 0; i < simHits.size(); ++i) {
0386       hitsByParticle.insert({simHits.nth(i)->particleId(), i});
0387     }
0388
0389     for (auto it = hitsByParticle.begin(), end = hitsByParticle.end();
0390          it != end; it = hitsByParticle.upper_bound(it->first)) {
0391       ACTS_DEBUG("Particle " << it->first << " has "
0392                              << hitsByParticle.count(it->first) << " hits");
0393
0394       std::vector<std::size_t> hitIndices;
0395       hitIndices.reserve(hitsByParticle.count(it->first));
0396       for (auto hitIndex : makeRange(hitsByParticle.equal_range(it->first))) {
0397         hitIndices.push_back(hitIndex.second);
0398       }
0399
0400       if (logger().doPrint(Acts::Logging::VERBOSE)) {
0401         ACTS_VERBOSE("Before sorting:");
0402         for (const auto& hitIdx : hitIndices) {
0403           ACTS_VERBOSE(" - " << hitIdx << " / " << simHits.nth(hitIdx)->index()
0404                              << " " << simHits.nth(hitIdx)->time());
0405         }
0406       }
0407
0408       std::ranges::sort(hitIndices, {}, [&simHits](std::size_t h) {
0409         return simHits.nth(h)->time();
0410       });
0411
0412       for (std::size_t i = 0; i < hitIndices.size(); ++i) {
0413         auto& hit = *simHits.nth(hitIndices[i]);
0414         SimHit updatedHit{hit.geometryId(),     hit.particleId(),
0415                           hit.fourPosition(),   hit.momentum4Before(),
0416                           hit.momentum4After(), static_cast<std::int32_t>(i)};
0417         hit = updatedHit;
0418       }
0419
0420       if (logger().doPrint(Acts::Logging::VERBOSE)) {
0421         ACTS_VERBOSE("After sorting:");
0422         for (const auto& hitIdx : hitIndices) {
0423           ACTS_VERBOSE(" - " << hitIdx << " / " << simHits.nth(hitIdx)->index()
0424                              << " " << simHits.nth(hitIdx)->time());
0425         }
0426       }
0427     }
0428   }
0429
0430   m_outputParticlesGenerator(ctx, std::move(particlesGenerator));
0431   m_outputParticlesSimulation(ctx, std::move(particlesSimulated));
0432
0433   m_outputSimHits(ctx, std::move(simHits));
0434
0435   return ProcessCode::SUCCESS;
0436 }
0437
0438 void EDM4hepReader::processChildren(
0439     const edm4hep::MCParticle& inParticle, SimBarcode parentId,
0440     std::vector<SimParticle>& particles, ParentRelationship& parentRelationship,
0441     std::unordered_map<int, std::size_t>& particleMap,
0442     std::size_t& nSecondaryVertices, std::size_t& maxGen) const {
0443   constexpr auto indent = [&](std::size_t n) {
0444     std::string result;
0445     for (std::size_t i = 0; i < n; ++i) {
0446       result += "  ";
0447     }
0448     return result;
0449   };
0450
0451   const std::size_t gen = parentId.generation();
0452   maxGen = std::max(maxGen, gen);
0453
0454   ACTS_VERBOSE(indent(gen) << "  - processing daughters for input particle "
0455                            << inParticle.id());
0456   ACTS_VERBOSE(indent(gen) << "    -> found " << inParticle.daughters_size()
0457                            << " daughter(s)");
0458
0459   bool parentDecayed =
0460       std::any_of(inParticle.daughters_begin(), inParticle.daughters_end(),
0461                   [](const edm4hep::MCParticle& daughter) {
0462                     return !daughter.vertexIsNotEndpointOfParent();
0463                   });
0464   std::size_t secondaryVertex = 0;
0465   if (parentDecayed) {
0466     ACTS_VERBOSE(indent(gen) << "    -> parent decays");
0467     secondaryVertex = ++nSecondaryVertices;
0468   }
0469
0470   std::size_t parentIndex = particles.size() - 1;
0471
0472   std::size_t nParticles = 0;
0473   for (const auto& daughter : inParticle.getDaughters()) {
0474     SimParticle particle = EDM4hepUtil::readParticle(daughter);
0475
0476     auto pid = parentId.makeDescendant(nParticles++);
0477     if (daughter.vertexIsNotEndpointOfParent()) {
0478       // incoming particle survived, interaction via descendant
0479     } else {
0480       // incoming particle decayed
0481       pid = pid.setVertexSecondary(secondaryVertex);
0482     }
0483     particle.setParticleId(pid);
0484
0485     ACTS_VERBOSE(indent(particle.particleId().generation())
0486                  << "+ add particle " << particle);
0487     ACTS_VERBOSE(indent(particle.particleId().generation())
0488                  << "  - generation: " << particle.particleId().generation());
0489     ACTS_VERBOSE(indent(particle.particleId().generation())
0490                  << "  - at " << particle.position().transpose());
0491     ACTS_VERBOSE(indent(particle.particleId().generation())
0492                  << "     - createdInSim: "
0493                  << daughter.isCreatedInSimulation());
0494     ACTS_VERBOSE(indent(particle.particleId().generation())
0495                  << "     - vertexIsNotEndpointOfParent: "
0496                  << daughter.vertexIsNotEndpointOfParent());
0497     ACTS_VERBOSE(indent(particle.particleId().generation())
0498                  << "     - isStopped: " << daughter.isStopped());
0499     ACTS_VERBOSE(indent(particle.particleId().generation())
0500                  << "     - endpoint: " << daughter.getEndpoint().x << ", "
0501                  << daughter.getEndpoint().y << ", "
0502                  << daughter.getEndpoint().z);
0503
0504     particles.push_back(std::move(particle));
0505     particleMap[daughter.getObjectID().index] = particles.size() - 1;
0506     parentRelationship[particles.size() - 1] = parentIndex;
0507     processChildren(daughter, pid, particles, parentRelationship, particleMap,
0508                     nSecondaryVertices, maxGen);
0509   }
0510 }
0511
0512 void EDM4hepReader::setSubParticleIds(std::vector<SimParticle>::iterator begin,
0513                                       std::vector<SimParticle>::iterator end) {
0514   std::vector<std::size_t> numByGeneration;
0515   numByGeneration.reserve(10);
0516
0517   for (auto it = begin; it != end; ++it) {
0518     auto& particle = *it;
0519     const auto pid = particle.particleId();
0520     if (pid.generation() >= numByGeneration.size()) {
0521       numByGeneration.resize(pid.generation() + 1, 0);
0522     }
0523     unsigned int nextSubParticle = numByGeneration[pid.generation()]++;
0524
0525     auto newPid = particle.particleId().setSubParticle(nextSubParticle);
0526     particle.setParticleId(newPid);
0527   }
0528 }
0529
0530 }  // namespace ActsExamples