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