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<ActsFatras::Barcode, std::vector<Acts::Vector4>>
0072         particleHits;
0073     // Pre-loop over hits ... write those below threshold
0074     for (const auto& simHit : simHits) {
0075       double momentum = simHit.momentum4Before().head<3>().norm();
0076       if (momentum < m_cfg.momentumThreshold) {
0077         ACTS_VERBOSE("Skipping: Hit below threshold: " << momentum);
0078         continue;
0079       } else if (momentum < m_cfg.momentumThresholdTraj) {
0080         ACTS_VERBOSE(
0081             "Skipping (trajectory): Hit below threshold: " << momentum);
0082         osHits << "v "
0083                << simHit.fourPosition()[Acts::ePos0] / Acts::UnitConstants::mm
0084                << " "
0085                << simHit.fourPosition()[Acts::ePos1] / Acts::UnitConstants::mm
0086                << " "
0087                << simHit.fourPosition()[Acts::ePos2] / Acts::UnitConstants::mm
0088                << std::endl;
0089         continue;
0090       }
0091       ACTS_VERBOSE("Accepting: Hit above threshold: " << momentum);
0092
0093       if (particleHits.find(simHit.particleId()) == particleHits.end()) {
0094         particleHits[simHit.particleId()] = {};
0095       }
0096       particleHits[simHit.particleId()].push_back(simHit.fourPosition());
0097     }
0098     // Draw loop
0099     std::size_t lOffset = 1;
0100     for (auto& [pId, pHits] : particleHits) {
0101       // Draw the particle hits
0102       std::ranges::sort(pHits,
0103                         [](const Acts::Vector4& a, const Acts::Vector4& b) {
0104                           return a[Acts::eTime] < b[Acts::eTime];
0105                         });
0106
0107       osHits << "o particle_" << pId << std::endl;
0108       for (const auto& hit : pHits) {
0109         osHits << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm << " "
0110                << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
0111                << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0112       }
0113       osHits << '\n';
0114
0115       // Interpolate the points, a minimum number of 3 hits is necessary for
0116       // that
0117       std::vector<Acts::Vector4> trajectory;
0118       if (pHits.size() < 3 || m_cfg.nInterpolatedPoints == 0) {
0119         trajectory = pHits;
0120       } else {
0121         // The total number of points is the number of hits times the number of
0122         // interpolated points plus the number of hits
0123         trajectory = Acts::Interpolation3D::spline(
0124             pHits, pHits.size() * (m_cfg.nInterpolatedPoints + 1) - 1,
0125             m_cfg.keepOriginalHits);
0126       }
0127
0128       osTrajectory << "o particle_trajectory_" << pId << std::endl;
0129       for (const auto& hit : trajectory) {
0130         osTrajectory << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm
0131                      << " " << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
0132                      << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0133       }
0134       // Draw the line
0135       for (std::size_t iv = lOffset + 1; iv < lOffset + trajectory.size();
0136            ++iv) {
0137         osTrajectory << "l " << iv - 1 << " " << iv << '\n';
0138       }
0139       osTrajectory << '\n';
0140       // Increase the offset count
0141       lOffset += trajectory.size();
0142     }
0143   }
0144   osHits.close();
0145   osTrajectory.close();
0146
0147   return ActsExamples::ProcessCode::SUCCESS;
0148 }
0149
0150 ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::finalize() {
0151   return ActsExamples::ProcessCode::SUCCESS;
0152 }