Root/src/RootSeedWriter.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/RootSeedWriter.hpp"
0010
0011 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0012 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0013
0014 #include <ios>
0015 #include <stdexcept>
0016
0017 #include <TFile.h>
0018 #include <TTree.h>
0019
0020 namespace ActsExamples {
0021
0022 RootSeedWriter::RootSeedWriter(const RootSeedWriter::Config& config,
0023                                Acts::Logging::Level level)
0024     : WriterT(config.inputSeeds, "RootSeedWriter", level), m_cfg(config) {
0025   // inputParticles is already checked by base constructor
0026   if (m_cfg.filePath.empty()) {
0027     throw std::invalid_argument("Missing file path");
0028   }
0029   if (m_cfg.treeName.empty()) {
0030     throw std::invalid_argument("Missing tree name");
0031   }
0032
0033   // open root file and create the tree
0034   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0035   if (m_outputFile == nullptr) {
0036     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0037   }
0038   m_outputFile->cd();
0039   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0040   if (m_outputTree == nullptr) {
0041     throw std::bad_alloc();
0042   }
0043
0044   // setup the branches
0045   m_outputTree->Branch("event_id", &m_eventId);
0046   m_outputTree->Branch("measurement_id_1", &m_measurementId_1,
0047                        "measurement_id_1/l");
0048   m_outputTree->Branch("measurement_id_2", &m_measurementId_2,
0049                        "measurement_id_2/l");
0050   m_outputTree->Branch("measurement_id_3", &m_measurementId_3,
0051                        "measurement_id_3/l");
0052   if (m_cfg.writingMode != "small") {
0053     m_outputTree->Branch("geometry_id_1", &m_geometryId_1, "geometry_id_1/l");
0054     m_outputTree->Branch("geometry_id_2", &m_geometryId_2, "geometry_id_2/l");
0055     m_outputTree->Branch("geometry_id_3", &m_geometryId_3, "geometry_id_3/l");
0056     m_outputTree->Branch("x_1", &m_x_1);
0057     m_outputTree->Branch("x_2", &m_x_2);
0058     m_outputTree->Branch("x_3", &m_x_3);
0059     m_outputTree->Branch("y_1", &m_y_1);
0060     m_outputTree->Branch("y_2", &m_y_2);
0061     m_outputTree->Branch("y_3", &m_y_3);
0062     m_outputTree->Branch("z_1", &m_z_1);
0063     m_outputTree->Branch("z_2", &m_z_2);
0064     m_outputTree->Branch("z_3", &m_z_3);
0065     m_outputTree->Branch("var_r_1", &m_var_r_1);
0066     m_outputTree->Branch("var_r_2", &m_var_r_2);
0067     m_outputTree->Branch("var_r_3", &m_var_r_3);
0068     m_outputTree->Branch("var_z_1", &m_var_z_1);
0069     m_outputTree->Branch("var_z_2", &m_var_z_2);
0070     m_outputTree->Branch("var_z_3", &m_var_z_3);
0071     m_outputTree->Branch("z_vertex", &m_z_vertex);
0072     m_outputTree->Branch("seed_quality", &m_seed_quality);
0073   }
0074 }
0075
0076 RootSeedWriter::~RootSeedWriter() {
0077   if (m_outputFile != nullptr) {
0078     m_outputFile->Close();
0079   }
0080 }
0081
0082 ProcessCode RootSeedWriter::finalize() {
0083   m_outputFile->cd();
0084   m_outputTree->Write();
0085   m_outputFile->Close();
0086
0087   ACTS_VERBOSE("Wrote seeds to tree '" << m_cfg.treeName << "' in '"
0088                                        << m_cfg.filePath << "'");
0089   return ProcessCode::SUCCESS;
0090 }
0091
0092 ProcessCode RootSeedWriter::writeT(const AlgorithmContext& ctx,
0093                                    const SeedContainer& seeds) {
0094   // ensure exclusive access to tree/file while writing
0095   std::lock_guard<std::mutex> lock(m_writeMutex);
0096
0097   // Get the event number
0098   m_eventId = ctx.eventNumber;
0099   for (const auto& seed : seeds) {
0100     const auto& spacePoints = seed.spacePoints();
0101
0102     const auto slink_1 = spacePoints[0].sourceLinks()[0].get<IndexSourceLink>();
0103     const auto slink_2 = spacePoints[1].sourceLinks()[0].get<IndexSourceLink>();
0104     const auto slink_3 = spacePoints[2].sourceLinks()[0].get<IndexSourceLink>();
0105
0106     m_measurementId_1 = slink_1.index();
0107     if (m_cfg.writingMode != "small") {
0108       m_geometryId_1 = slink_1.geometryId().value();
0109       m_x_1 = spacePoints[0].x();
0110       m_y_1 = spacePoints[0].y();
0111       m_z_1 = spacePoints[0].z();
0112       m_var_r_1 = spacePoints[0].varianceR();
0113       m_var_z_1 = spacePoints[0].varianceZ();
0114     }
0115
0116     m_measurementId_2 = slink_2.index();
0117     if (m_cfg.writingMode != "small") {
0118       m_geometryId_2 = slink_2.geometryId().value();
0119       m_x_2 = spacePoints[1].x();
0120       m_y_2 = spacePoints[1].y();
0121       m_z_2 = spacePoints[1].z();
0122       m_var_r_2 = spacePoints[1].varianceR();
0123       m_var_z_2 = spacePoints[1].varianceZ();
0124     }
0125
0126     m_measurementId_3 = slink_3.index();
0127     if (m_cfg.writingMode != "small") {
0128       m_geometryId_3 = slink_3.geometryId().value();
0129       m_x_3 = spacePoints[2].x();
0130       m_y_3 = spacePoints[2].y();
0131       m_z_3 = spacePoints[2].z();
0132       m_var_r_3 = spacePoints[2].varianceR();
0133       m_var_z_3 = spacePoints[2].varianceZ();
0134     }
0135
0136     if (m_cfg.writingMode != "small") {
0137       m_z_vertex = seed.vertexZ();
0138       m_seed_quality = seed.quality();
0139     }
0140     // Fill the tree
0141     m_outputTree->Fill();
0142   }
0143   return ProcessCode::SUCCESS;
0144 }
0145
0146 }  // namespace ActsExamples