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