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("t", &m_t);
0064   m_outputTree->Branch("var_r", &m_var_r);
0065   m_outputTree->Branch("var_z", &m_var_z);
0066   if (m_inputMeasurementParticlesMap.isInitialized()) {
0067     m_outputTree->Branch("fake", &m_fake);
0068   }
0069 }
0070
0071 ActsExamples::RootSpacepointWriter::~RootSpacepointWriter() {
0072   if (m_outputFile != nullptr) {
0073     m_outputFile->Close();
0074   }
0075 }
0076
0077 ActsExamples::ProcessCode ActsExamples::RootSpacepointWriter::finalize() {
0078   m_outputFile->cd();
0079   m_outputTree->Write();
0080   m_outputFile->Close();
0081
0082   ACTS_VERBOSE("Wrote hits to tree '" << m_cfg.treeName << "' in '"
0083                                       << m_cfg.filePath << "'");
0084
0085   return ProcessCode::SUCCESS;
0086 }
0087
0088 ActsExamples::ProcessCode ActsExamples::RootSpacepointWriter::writeT(
0089     const AlgorithmContext& ctx,
0090     const ActsExamples::SimSpacePointContainer& spacepoints) {
0091   // ensure exclusive access to tree/file while writing
0092   std::lock_guard<std::mutex> lock(m_writeMutex);
0093
0094   MeasurementParticlesMap emptyMeasPartMap{};
0095   const auto& measPartMap = m_inputMeasurementParticlesMap.isInitialized()
0096                                 ? m_inputMeasurementParticlesMap(ctx)
0097                                 : emptyMeasPartMap;
0098
0099   // Get the event number
0100   m_eventId = ctx.eventNumber;
0101   for (const auto& sp : spacepoints) {
0102     const auto& sl1 = sp.sourceLinks().at(0).get<IndexSourceLink>();
0103     m_measurementId1 = sl1.index();
0104     m_geometryId1 = sl1.geometryId().value();
0105     if (sp.sourceLinks().size() == 2) {
0106       const auto& sl2 = sp.sourceLinks().at(1).get<IndexSourceLink>();
0107       m_measurementId2 = sl2.index();
0108       m_geometryId2 = sl2.geometryId().value();
0109     }
0110     // A spacepoint is fake if the measurements have no common particle
0111     if (m_inputMeasurementParticlesMap.isInitialized()) {
0112       if (sp.sourceLinks().size() == 1) {
0113         m_fake = false;
0114       } else {
0115         m_fake = true;
0116         auto [p1b, p1e] = measPartMap.equal_range(m_measurementId1);
0117         auto [p2b, p2e] = measPartMap.equal_range(m_measurementId2);
0118         for (auto it1 = p1b; it1 != p1e; ++it1) {
0119           for (auto it2 = p2b; it2 != p2e; ++it2) {
0120             if (*it1 == *it2) {
0121               m_fake = false;
0122             }
0123           }
0124         }
0125       }
0126     }
0127     // write sp position
0128     m_x = sp.x() / Acts::UnitConstants::mm;
0129     m_y = sp.y() / Acts::UnitConstants::mm;
0130     m_z = sp.z() / Acts::UnitConstants::mm;
0131     m_r = sp.r() / Acts::UnitConstants::mm;
0132     m_t = sp.t() ? *sp.t() / Acts::UnitConstants::ns
0133                  : std::numeric_limits<double>::quiet_NaN();
0134     // write sp dimensions
0135     m_var_r = sp.varianceR() / Acts::UnitConstants::mm;
0136     m_var_z = sp.varianceZ() / Acts::UnitConstants::mm;
0137     // Fill the tree
0138     m_outputTree->Fill();
0139   }
0140   return ActsExamples::ProcessCode::SUCCESS;
0141 }