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 <fstream>
0022 #include <stdexcept>
0023 #include <unordered_map>
0024 #include <vector>
0025
0026 ActsExamples::ObjSimHitWriter::ObjSimHitWriter(
0027     const ActsExamples::ObjSimHitWriter::Config& config,
0028     Acts::Logging::Level level)
0029     : WriterT(config.inputSimHits, "ObjSimHitWriter", level), m_cfg(config) {
0030   // inputSimHits is already checked by base constructor
0031   if (m_cfg.outputStem.empty()) {
0032     throw std::invalid_argument("Missing output filename stem");
0033   }
0034 }
0035
0036 ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::writeT(
0037     const AlgorithmContext& ctx, const ActsExamples::SimHitContainer& simHits) {
0038   // ensure exclusive access to tree/file while writing
0039   std::scoped_lock lock(m_writeMutex);
0040
0041   // open per-event file for all simhit components
0042   std::string pathSimHit = perEventFilepath(
0043       m_cfg.outputDir, m_cfg.outputStem + ".obj", ctx.eventNumber);
0044   std::string pathSimTrajectory = perEventFilepath(
0045       m_cfg.outputDir, m_cfg.outputStem + "_trajectory.obj", ctx.eventNumber);
0046
0047   std::ofstream osHits(pathSimHit, std::ofstream::out | std::ofstream::trunc);
0048   std::ofstream osTrajectory(pathSimTrajectory,
0049                              std::ofstream::out | std::ofstream::trunc);
0050
0051   if (!osHits || !osTrajectory) {
0052     throw std::ios_base::failure("Could not open '" + pathSimHit + "' or '" +
0053                                  pathSimTrajectory + "' to write");
0054   }
0055
0056   // Only hit plotting
0057   if (!m_cfg.drawConnections) {
0058     // Write data from internal immplementation
0059     for (const auto& simHit : simHits) {
0060       // local simhit information in global coord.
0061       const Acts::Vector4& globalPos4 = simHit.fourPosition();
0062
0063       osHits << "v " << globalPos4[Acts::ePos0] / Acts::UnitConstants::mm << " "
0064              << globalPos4[Acts::ePos1] / Acts::UnitConstants::mm << " "
0065              << globalPos4[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
0066
0067     }  // end simHit loop
0068   } else {
0069     // We need to associate first
0070     std::unordered_map<std::size_t, std::vector<Acts::Vector4>> particleHits;
0071     // Pre-loop over hits ... write those below threshold
0072     for (const auto& simHit : simHits) {
0073       double momentum = simHit.momentum4Before().head<3>().norm();
0074       if (momentum < m_cfg.momentumThreshold) {
0075         ACTS_VERBOSE("Skipping: Hit below threshold: " << momentum);
0076         continue;
0077       } else if (momentum < m_cfg.momentumThresholdTraj) {
0078         ACTS_VERBOSE(
0079             "Skipping (trajectory): Hit below threshold: " << momentum);
0080         osHits << "v "
0081                << simHit.fourPosition()[Acts::ePos0] / Acts::UnitConstants::mm
0082                << " "
0083                << simHit.fourPosition()[Acts::ePos1] / Acts::UnitConstants::mm
0084                << " "
0085                << simHit.fourPosition()[Acts::ePos2] / Acts::UnitConstants::mm
0086                << std::endl;
0087         continue;
0088       }
0089       ACTS_VERBOSE("Accepting: Hit above threshold: " << momentum);
0090
0091       if (particleHits.find(simHit.particleId().value()) ==
0092           particleHits.end()) {
0093         particleHits[simHit.particleId().value()] = {};
0094       }
0095       particleHits[simHit.particleId().value()].push_back(
0096           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::sort(pHits.begin(), pHits.end(),
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 }