Root/src/RootVertexWriter.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/Root/RootVertexWriter.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Utilities/Helpers.hpp"
0013 #include "Acts/Utilities/VectorHelpers.hpp"
0014 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0015 #include "ActsFatras/EventData/Barcode.hpp"
0016 #include "ActsFatras/EventData/Particle.hpp"
0017
0018 #include <algorithm>
0019 #include <cstdint>
0020 #include <ios>
0021 #include <limits>
0022 #include <ostream>
0023 #include <stdexcept>
0024 #include <unordered_map>
0025
0026 #include <TFile.h>
0027 #include <TTree.h>
0028
0029 namespace ActsExamples {
0030
0031 RootVertexWriter::RootVertexWriter(const RootVertexWriter::Config& cfg,
0032                                    Acts::Logging::Level lvl)
0033     : WriterT(cfg.inputVertices, "RootVertexWriter", lvl), m_cfg(cfg) {
0034   // inputParticles is already checked by base constructor
0035   if (m_cfg.filePath.empty()) {
0036     throw std::invalid_argument("Missing file path");
0037   }
0038   if (m_cfg.treeName.empty()) {
0039     throw std::invalid_argument("Missing tree name");
0040   }
0041
0042   // open root file and create the tree
0043   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0044   if (m_outputFile == nullptr) {
0045     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0046   }
0047   m_outputFile->cd();
0048   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0049   if (m_outputTree == nullptr) {
0050     throw std::bad_alloc();
0051   }
0052
0053   // setup the branches
0054   m_outputTree->Branch("event_id", &m_eventId);
0055   m_outputTree->Branch("vertex_id", &m_vertexId);
0056   m_outputTree->Branch("process", &m_process);
0057   m_outputTree->Branch("vx", &m_vx);
0058   m_outputTree->Branch("vy", &m_vy);
0059   m_outputTree->Branch("vz", &m_vz);
0060   m_outputTree->Branch("vt", &m_vt);
0061   m_outputTree->Branch("outgoing_particles", &m_outgoingParticles);
0062   m_outputTree->Branch("vertex_primary", &m_vertexPrimary);
0063   m_outputTree->Branch("vertex_secondary", &m_vertexSecondary);
0064   m_outputTree->Branch("generation", &m_generation);
0065 }
0066
0067 RootVertexWriter::~RootVertexWriter() {
0068   if (m_outputFile != nullptr) {
0069     m_outputFile->Close();
0070   }
0071 }
0072
0073 ProcessCode RootVertexWriter::finalize() {
0074   m_outputFile->cd();
0075   m_outputTree->Write();
0076   m_outputFile->Close();
0077
0078   ACTS_INFO("Wrote vertices to tree '" << m_cfg.treeName << "' in '"
0079                                        << m_cfg.filePath << "'");
0080
0081   return ProcessCode::SUCCESS;
0082 }
0083
0084 ProcessCode RootVertexWriter::writeT(const AlgorithmContext& ctx,
0085                                      const SimVertexContainer& vertices) {
0086   // ensure exclusive access to tree/file while writing
0087   std::lock_guard<std::mutex> lock(m_writeMutex);
0088
0089   m_eventId = ctx.eventNumber;
0090   for (const auto& vertex : vertices) {
0091     m_vertexId.push_back(vertex.vertexId().value());
0092     m_process.push_back(static_cast<std::uint32_t>(vertex.process));
0093     // position
0094     m_vx.push_back(Acts::clampValue<float>(vertex.position4.x() /
0095                                            Acts::UnitConstants::mm));
0096     m_vy.push_back(Acts::clampValue<float>(vertex.position4.y() /
0097                                            Acts::UnitConstants::mm));
0098     m_vz.push_back(Acts::clampValue<float>(vertex.position4.z() /
0099                                            Acts::UnitConstants::mm));
0100     m_vt.push_back(Acts::clampValue<float>(vertex.position4.w() /
0101                                            Acts::UnitConstants::mm));
0102     // TODO ingoing particles
0103     // outgoing particles
0104     std::vector<std::uint64_t> outgoing;
0105     for (const auto& particle : vertex.outgoing) {
0106       outgoing.push_back(particle.value());
0107     }
0108     m_outgoingParticles.push_back(std::move(outgoing));
0109     // decoded barcode components
0110     m_vertexPrimary.push_back(vertex.vertexId().vertexPrimary());
0111     m_vertexSecondary.push_back(vertex.vertexId().vertexSecondary());
0112     m_generation.push_back(vertex.vertexId().generation());
0113   }
0114
0115   m_outputTree->Fill();
0116
0117   m_vertexId.clear();
0118   m_process.clear();
0119   m_vx.clear();
0120   m_vy.clear();
0121   m_vz.clear();
0122   m_vt.clear();
0123   m_outgoingParticles.clear();
0124   m_vertexPrimary.clear();
0125   m_vertexSecondary.clear();
0126   m_generation.clear();
0127
0128   return ProcessCode::SUCCESS;
0129 }
0130
0131 }  // namespace ActsExamples