Obj/src/ObjSimHitWriter.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/Obj/ObjSimHitWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Definitions/Common.hpp"
0013 #include "Acts/Geometry/GeometryIdentifier.hpp"
0014 #include "Acts/Visualization/Interpolation3D.hpp"
0015 #include "ActsExamples/EventData/SimHit.hpp"
0016 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0017 #include "ActsExamples/Utilities/Paths.hpp"
0018 #include "ActsFatras/EventData/Barcode.hpp"
0019 #include "ActsFatras/EventData/Hit.hpp"
0020
0021 #include <algorithm>
0022 #include <fstream>
0023 #include <stdexcept>
0024 #include <unordered_map>
0025 #include <vector>
0026
0027 ActsExamples::ObjSimHitWriter::ObjSimHitWriter(
0028     const ActsExamples::ObjSimHitWriter::Config& config,
0029     Acts::Logging::Level level)
0030     : WriterT(config.inputSimHits, "ObjSimHitWriter", level), m_cfg(config) {
0031   // inputSimHits is already checked by base constructor
0032   if (m_cfg.outputStem.empty()) {
0033     throw std::invalid_argument("Missing output filename stem");
0034   }
0035 }
0036
0037 ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::writeT(
0038     const AlgorithmContext& ctx, const ActsExamples::SimHitContainer& simHits) {
0039   // ensure exclusive access to tree/file while writing
0040   std::scoped_lock lock(m_writeMutex);
0041
0042   // open per-event file for all simhit components
0043   std::string pathSimHit = perEventFilepath(
0044       m_cfg.outputDir, m_cfg.outputStem + ".obj", ctx.eventNumber);
0045   std::string pathSimTrajectory = perEventFilepath(
0046       m_cfg.outputDir, m_cfg.outputStem + "_trajectory.obj", ctx.eventNumber);
0047
0048   std::ofstream osHits(pathSimHit, std::ofstream::out | std::ofstream::trunc);
0049   std::ofstream osTrajectory(pathSimTrajectory,
0050                              std::ofstream::out | std::ofstream::trunc);
0051
0052   if (!osHits || !osTrajectory) {
0053     throw std::ios_base::failure("Could not open '" + pathSimHit + "' or '" +
0054                                  pathSimTrajectory + "' to write");
0055   }
0056
0057   // Only hit plotting
0058   if (!m_cfg.drawConnections) {
0059     // Write data from internal immplementation
0060     for (const auto& simHit : simHits) {
0061       // local simhit information in global coord.
0062       const Acts::Vector4& globalPos4 = simHit.fourPosition();
0063
0064       osHits << "v " << globalPos4[Acts::ePos0] / Acts::UnitConstants::mm << " "
0065              << globalPos4[Acts::ePos1] / Acts::UnitConstants::mm << " "
0066              << globalPos4[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0067
0068     }  // end simHit loop
0069   } else {
0070     // We need to associate first
0071     std::unordered_map<std::size_t, std::vector<Acts::Vector4>> particleHits;
0072     // Pre-loop over hits ... write those below threshold
0073     for (const auto& simHit : simHits) {
0074       double momentum = simHit.momentum4Before().head<3>().norm();
0075       if (momentum < m_cfg.momentumThreshold) {
0076         ACTS_VERBOSE("Skipping: Hit below threshold: " << momentum);
0077         continue;
0078       } else if (momentum < m_cfg.momentumThresholdTraj) {
0079         ACTS_VERBOSE(
0080             "Skipping (trajectory): Hit below threshold: " << momentum);
0081         osHits << "v "
0082                << simHit.fourPosition()[Acts::ePos0] / Acts::UnitConstants::mm
0083                << " "
0084                << simHit.fourPosition()[Acts::ePos1] / Acts::UnitConstants::mm
0085                << " "
0086                << simHit.fourPosition()[Acts::ePos2] / Acts::UnitConstants::mm
0087                << std::endl;
0088         continue;
0089       }
0090       ACTS_VERBOSE("Accepting: Hit above threshold: " << momentum);
0091
0092       if (particleHits.find(simHit.particleId().value()) ==
0093           particleHits.end()) {
0094         particleHits[simHit.particleId().value()] = {};
0095       }
0096       particleHits[simHit.particleId().value()].push_back(
0097           simHit.fourPosition());
0098     }
0099     // Draw loop
0100     std::size_t lOffset = 1;
0101     for (auto& [pId, pHits] : particleHits) {
0102       // Draw the particle hits
0103       std::ranges::sort(pHits,
0104                         [](const Acts::Vector4& a, const Acts::Vector4& b) {
0105                           return a[Acts::eTime] < b[Acts::eTime];
0106                         });
0107
0108       osHits << "o particle_" << pId << std::endl;
0109       for (const auto& hit : pHits) {
0110         osHits << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm << " "
0111                << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
0112                << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0113       }
0114       osHits << '\n';
0115
0116       // Interpolate the points, a minimum number of 3 hits is necessary for
0117       // that
0118       std::vector<Acts::Vector4> trajectory;
0119       if (pHits.size() < 3 || m_cfg.nInterpolatedPoints == 0) {
0120         trajectory = pHits;
0121       } else {
0122         // The total number of points is the number of hits times the number of
0123         // interpolated points plus the number of hits
0124         trajectory = Acts::Interpolation3D::spline(
0125             pHits, pHits.size() * (m_cfg.nInterpolatedPoints + 1) - 1,
0126             m_cfg.keepOriginalHits);
0127       }
0128
0129       osTrajectory << "o particle_trajectory_" << pId << std::endl;
0130       for (const auto& hit : trajectory) {
0131         osTrajectory << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm
0132                      << " " << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
0133                      << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0134       }
0135       // Draw the line
0136       for (std::size_t iv = lOffset + 1; iv < lOffset + trajectory.size();
0137            ++iv) {
0138         osTrajectory << "l " << iv - 1 << " " << iv << '\n';
0139       }
0140       osTrajectory << '\n';
0141       // Increase the offset count
0142       lOffset += trajectory.size();
0143     }
0144   }
0145   osHits.close();
0146   osTrajectory.close();
0147
0148   return ActsExamples::ProcessCode::SUCCESS;
0149 }
0150
0151 ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::finalize() {
0152   return ActsExamples::ProcessCode::SUCCESS;
0153 }