Root/src/RootSpacepointWriter.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/RootSpacepointWriter.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 <ostream>
0019 #include <stdexcept>
0020 #include <vector>
0021
0022 #include <TFile.h>
0023 #include <TTree.h>
0024
0025 ActsExamples::RootSpacepointWriter::RootSpacepointWriter(
0026     const ActsExamples::RootSpacepointWriter::Config& config,
0027     Acts::Logging::Level level)
0028     : WriterT(config.inputSpacepoints, "RootSpacepointWriter", level),
0029       m_cfg(config) {
0030   // inputParticles is already checked by base constructor
0031   if (m_cfg.filePath.empty()) {
0032     throw std::invalid_argument("Missing file path");
0033   }
0034   if (m_cfg.treeName.empty()) {
0035     throw std::invalid_argument("Missing tree name");
0036   }
0037
0038   // open root file and create the tree
0039   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0040   if (m_outputFile == nullptr) {
0041     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0042   }
0043   m_outputFile->cd();
0044   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0045   if (m_outputTree == nullptr) {
0046     throw std::bad_alloc();
0047   }
0048
0049   m_inputMeasurementParticlesMap.maybeInitialize(
0050       m_cfg.inputMeasurementParticlesMap);
0051
0052   // setup the branches
0053   m_outputTree->Branch("event_id", &m_eventId);
0054   m_outputTree->Branch("measurement_id", &m_measurementId1, "measurement_id/l");
0055   m_outputTree->Branch("geometry_id", &m_geometryId1, "geometry_id/l");
0056   m_outputTree->Branch("measurement_id_2", &m_measurementId2,
0057                        "measurement_id_2/l");
0058   m_outputTree->Branch("geometry_id_2", &m_geometryId2, "geometry_id_2/l");
0059   m_outputTree->Branch("x", &m_x);
0060   m_outputTree->Branch("y", &m_y);
0061   m_outputTree->Branch("z", &m_z);
0062   m_outputTree->Branch("r", &m_r);
0063   m_outputTree->Branch("var_r", &m_var_r);
0064   m_outputTree->Branch("var_z", &m_var_z);
0065   if (m_inputMeasurementParticlesMap.isInitialized()) {
0066     m_outputTree->Branch("fake", &m_fake);
0067   }
0068 }
0069
0070 ActsExamples::RootSpacepointWriter::~RootSpacepointWriter() {
0071   if (m_outputFile != nullptr) {
0072     m_outputFile->Close();
0073   }
0074 }
0075
0076 ActsExamples::ProcessCode ActsExamples::RootSpacepointWriter::finalize() {
0077   m_outputFile->cd();
0078   m_outputTree->Write();
0079   m_outputFile->Close();
0080
0081   ACTS_VERBOSE("Wrote hits to tree '" << m_cfg.treeName << "' in '"
0082                                       << m_cfg.filePath << "'");
0083
0084   return ProcessCode::SUCCESS;
0085 }
0086
0087 ActsExamples::ProcessCode ActsExamples::RootSpacepointWriter::writeT(
0088     const AlgorithmContext& ctx,
0089     const ActsExamples::SimSpacePointContainer& spacepoints) {
0090   // ensure exclusive access to tree/file while writing
0091   std::lock_guard<std::mutex> lock(m_writeMutex);
0092
0093   MeasurementParticlesMap emptyMeasPartMap{};
0094   const auto& measPartMap = m_inputMeasurementParticlesMap.isInitialized()
0095                                 ? m_inputMeasurementParticlesMap(ctx)
0096                                 : emptyMeasPartMap;
0097
0098   // Get the event number
0099   m_eventId = ctx.eventNumber;
0100   for (const auto& sp : spacepoints) {
0101     const auto& sl1 = sp.sourceLinks().at(0).get<IndexSourceLink>();
0102     m_measurementId1 = sl1.index();
0103     m_geometryId1 = sl1.geometryId().value();
0104     if (sp.sourceLinks().size() == 2) {
0105       const auto& sl2 = sp.sourceLinks().at(1).get<IndexSourceLink>();
0106       m_measurementId2 = sl1.index();
0107       m_geometryId2 = sl2.geometryId().value();
0108     }
0109     // A spacepoint is fake if the measurements have no common particle
0110     if (m_inputMeasurementParticlesMap.isInitialized()) {
0111       if (sp.sourceLinks().size() == 1) {
0112         m_fake = false;
0113       } else {
0114         m_fake = true;
0115         auto [p1b, p1e] = measPartMap.equal_range(m_measurementId1);
0116         auto [p2b, p2e] = measPartMap.equal_range(m_measurementId2);
0117         for (auto it1 = p1b; it1 != p1e; ++it1) {
0118           for (auto it2 = p2b; it2 != p2e; ++it2) {
0119             if (*it1 == *it2) {
0120               m_fake = false;
0121             }
0122           }
0123         }
0124       }
0125     }
0126     // write sp position
0127     m_x = sp.x() / Acts::UnitConstants::mm;
0128     m_y = sp.y() / Acts::UnitConstants::mm;
0129     m_z = sp.z() / Acts::UnitConstants::mm;
0130     m_r = sp.r() / Acts::UnitConstants::mm;
0131     // write sp dimensions
0132     m_var_r = sp.varianceR() / Acts::UnitConstants::mm;
0133     m_var_z = sp.varianceZ() / Acts::UnitConstants::mm;
0134     // Fill the tree
0135     m_outputTree->Fill();
0136   }
0137   return ActsExamples::ProcessCode::SUCCESS;
0138 }