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_vertex_primary",
0056                        &m_incomingParticlesVertexPrimary);
0057   m_outputTree->Branch("incoming_particles_vertex_secondary",
0058                        &m_incomingParticlesVertexSecondary);
0059   m_outputTree->Branch("incoming_particles_particle",
0060                        &m_incomingParticlesParticle);
0061   m_outputTree->Branch("incoming_particles_generation",
0062                        &m_incomingParticlesGeneration);
0063   m_outputTree->Branch("incoming_particles_sub_particle",
0064                        &m_incomingParticlesSubParticle);
0065   m_outputTree->Branch("outgoing_particles_vertex_primary",
0066                        &m_outgoingParticlesVertexPrimary);
0067   m_outputTree->Branch("outgoing_particles_vertex_secondary",
0068                        &m_outgoingParticlesVertexSecondary);
0069   m_outputTree->Branch("outgoing_particles_particle",
0070                        &m_outgoingParticlesParticle);
0071   m_outputTree->Branch("outgoing_particles_generation",
0072                        &m_outgoingParticlesGeneration);
0073   m_outputTree->Branch("outgoing_particles_sub_particle",
0074                        &m_outgoingParticlesSubParticle);
0075   m_outputTree->Branch("vertex_primary", &m_vertexPrimary);
0076   m_outputTree->Branch("vertex_secondary", &m_vertexSecondary);
0077   m_outputTree->Branch("generation", &m_generation);
0078 }
0079
0080 RootVertexWriter::~RootVertexWriter() {
0081   if (m_outputFile != nullptr) {
0082     m_outputFile->Close();
0083   }
0084 }
0085
0086 ProcessCode RootVertexWriter::finalize() {
0087   m_outputFile->cd();
0088   m_outputTree->Write();
0089   m_outputFile->Close();
0090
0091   ACTS_INFO("Wrote vertices to tree '" << m_cfg.treeName << "' in '"
0092                                        << m_cfg.filePath << "'");
0093
0094   return ProcessCode::SUCCESS;
0095 }
0096
0097 ProcessCode RootVertexWriter::writeT(const AlgorithmContext& ctx,
0098                                      const SimVertexContainer& vertices) {
0099   // ensure exclusive access to tree/file while writing
0100   std::lock_guard<std::mutex> lock(m_writeMutex);
0101
0102   m_eventId = ctx.eventNumber;
0103   for (const auto& vertex : vertices) {
0104     m_process.push_back(static_cast<std::uint32_t>(vertex.process));
0105     // position
0106     m_vx.push_back(Acts::clampValue<float>(vertex.position4.x() /
0107                                            Acts::UnitConstants::mm));
0108     m_vy.push_back(Acts::clampValue<float>(vertex.position4.y() /
0109                                            Acts::UnitConstants::mm));
0110     m_vz.push_back(Acts::clampValue<float>(vertex.position4.z() /
0111                                            Acts::UnitConstants::mm));
0112     m_vt.push_back(Acts::clampValue<float>(vertex.position4.w() /
0113                                            Acts::UnitConstants::mm));
0114
0115     // incoming particles
0116     std::vector<std::uint32_t> incoming_vertex_primary;
0117     std::vector<std::uint32_t> incoming_vertex_secondary;
0118     std::vector<std::uint32_t> incoming_particle_component;
0119     std::vector<std::uint32_t> incoming_generation;
0120     std::vector<std::uint32_t> incoming_sub_particle;
0121     for (const auto& particle : vertex.incoming) {
0122       incoming_vertex_primary.push_back(particle.vertexPrimary());
0123       incoming_vertex_secondary.push_back(particle.vertexSecondary());
0124       incoming_particle_component.push_back(particle.particle());
0125       incoming_generation.push_back(particle.generation());
0126       incoming_sub_particle.push_back(particle.subParticle());
0127     }
0128     m_incomingParticlesVertexPrimary.push_back(
0129         std::move(incoming_vertex_primary));
0130     m_incomingParticlesVertexSecondary.push_back(
0131         std::move(incoming_vertex_secondary));
0132     m_incomingParticlesParticle.push_back(
0133         std::move(incoming_particle_component));
0134     m_incomingParticlesGeneration.push_back(std::move(incoming_generation));
0135     m_incomingParticlesSubParticle.push_back(std::move(incoming_sub_particle));
0136
0137     // outgoing particles
0138     std::vector<std::uint32_t> outgoing_vertex_primary;
0139     std::vector<std::uint32_t> outgoing_vertex_secondary;
0140     std::vector<std::uint32_t> outgoing_particle_component;
0141     std::vector<std::uint32_t> outgoing_generation;
0142     std::vector<std::uint32_t> outgoing_sub_particle;
0143     for (const auto& particle : vertex.outgoing) {
0144       outgoing_vertex_primary.push_back(particle.vertexPrimary());
0145       outgoing_vertex_secondary.push_back(particle.vertexSecondary());
0146       outgoing_particle_component.push_back(particle.particle());
0147       outgoing_generation.push_back(particle.generation());
0148       outgoing_sub_particle.push_back(particle.subParticle());
0149     }
0150     m_outgoingParticlesVertexPrimary.push_back(
0151         std::move(outgoing_vertex_primary));
0152     m_outgoingParticlesVertexSecondary.push_back(
0153         std::move(outgoing_vertex_secondary));
0154     m_outgoingParticlesParticle.push_back(
0155         std::move(outgoing_particle_component));
0156     m_outgoingParticlesGeneration.push_back(std::move(outgoing_generation));
0157     m_outgoingParticlesSubParticle.push_back(std::move(outgoing_sub_particle));
0158
0159     // decoded barcode components
0160     m_vertexPrimary.push_back(vertex.vertexId().vertexPrimary());
0161     m_vertexSecondary.push_back(vertex.vertexId().vertexSecondary());
0162     m_generation.push_back(vertex.vertexId().generation());
0163   }
0164
0165   m_outputTree->Fill();
0166
0167   m_process.clear();
0168   m_vx.clear();
0169   m_vy.clear();
0170   m_vz.clear();
0171   m_vt.clear();
0172   m_incomingParticlesVertexPrimary.clear();
0173   m_incomingParticlesVertexSecondary.clear();
0174   m_incomingParticlesParticle.clear();
0175   m_incomingParticlesGeneration.clear();
0176   m_incomingParticlesSubParticle.clear();
0177   m_outgoingParticlesVertexPrimary.clear();
0178   m_outgoingParticlesVertexSecondary.clear();
0179   m_outgoingParticlesParticle.clear();
0180   m_outgoingParticlesGeneration.clear();
0181   m_outgoingParticlesSubParticle.clear();
0182   m_vertexPrimary.clear();
0183   m_vertexSecondary.clear();
0184   m_generation.clear();
0185
0186   return ProcessCode::SUCCESS;
0187 }
0188
0189 }  // namespace ActsExamples