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("vertex_id", &m_vertexId);
0051   m_outputTree->Branch("process", &m_process);
0052   m_outputTree->Branch("vx", &m_vx);
0053   m_outputTree->Branch("vy", &m_vy);
0054   m_outputTree->Branch("vz", &m_vz);
0055   m_outputTree->Branch("vt", &m_vt);
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_vertexId.push_back(vertex.vertexId().value());
0087     m_process.push_back(static_cast<std::uint32_t>(vertex.process));
0088     // position
0089     m_vx.push_back(Acts::clampValue<float>(vertex.position4.x() /
0090                                            Acts::UnitConstants::mm));
0091     m_vy.push_back(Acts::clampValue<float>(vertex.position4.y() /
0092                                            Acts::UnitConstants::mm));
0093     m_vz.push_back(Acts::clampValue<float>(vertex.position4.z() /
0094                                            Acts::UnitConstants::mm));
0095     m_vt.push_back(Acts::clampValue<float>(vertex.position4.w() /
0096                                            Acts::UnitConstants::mm));
0097     // TODO ingoing particles
0098     // outgoing particles
0099     std::vector<std::uint64_t> outgoing;
0100     for (const auto& particle : vertex.outgoing) {
0101       outgoing.push_back(particle.value());
0102     }
0103     m_outgoingParticles.push_back(std::move(outgoing));
0104     // decoded barcode components
0105     m_vertexPrimary.push_back(vertex.vertexId().vertexPrimary());
0106     m_vertexSecondary.push_back(vertex.vertexId().vertexSecondary());
0107     m_generation.push_back(vertex.vertexId().generation());
0108   }
0109
0110   m_outputTree->Fill();
0111
0112   m_vertexId.clear();
0113   m_process.clear();
0114   m_vx.clear();
0115   m_vy.clear();
0116   m_vz.clear();
0117   m_vt.clear();
0118   m_outgoingParticles.clear();
0119   m_vertexPrimary.clear();
0120   m_vertexSecondary.clear();
0121   m_generation.clear();
0122
0123   return ProcessCode::SUCCESS;
0124 }
0125
0126 }  // namespace ActsExamples