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 "ActsExamples/EventData/IndexSourceLink.hpp"
0013 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0014
0015 #include <ios>
0016 #include <stdexcept>
0017
0018 #include <TFile.h>
0019 #include <TTree.h>
0020
0021 namespace ActsExamples {
0022
0023 RootSpacePointWriter::RootSpacePointWriter(
0024     const RootSpacePointWriter::Config& config, Acts::Logging::Level level)
0025     : WriterT(config.inputSpacePoints, "RootSpacePointWriter", level),
0026       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   m_inputMeasurementParticlesMap.maybeInitialize(
0047       m_cfg.inputMeasurementParticlesMap);
0048
0049   // setup the branches
0050   m_outputTree->Branch("event_id", &m_eventId);
0051   m_outputTree->Branch("measurement_id", &m_measurementId1, "measurement_id/l");
0052   m_outputTree->Branch("geometry_id", &m_geometryId1, "geometry_id/l");
0053   m_outputTree->Branch("measurement_id_2", &m_measurementId2,
0054                        "measurement_id_2/l");
0055   m_outputTree->Branch("geometry_id_2", &m_geometryId2, "geometry_id_2/l");
0056   m_outputTree->Branch("x", &m_x);
0057   m_outputTree->Branch("y", &m_y);
0058   m_outputTree->Branch("z", &m_z);
0059   m_outputTree->Branch("r", &m_r);
0060   m_outputTree->Branch("t", &m_t);
0061   m_outputTree->Branch("var_r", &m_var_r);
0062   m_outputTree->Branch("var_z", &m_var_z);
0063   if (m_inputMeasurementParticlesMap.isInitialized()) {
0064     m_outputTree->Branch("fake", &m_fake);
0065   }
0066 }
0067
0068 RootSpacePointWriter::~RootSpacePointWriter() {
0069   if (m_outputFile != nullptr) {
0070     m_outputFile->Close();
0071   }
0072 }
0073
0074 ProcessCode RootSpacePointWriter::finalize() {
0075   m_outputFile->cd();
0076   m_outputTree->Write();
0077   m_outputFile->Close();
0078
0079   ACTS_VERBOSE("Wrote hits to tree '" << m_cfg.treeName << "' in '"
0080                                       << m_cfg.filePath << "'");
0081
0082   return ProcessCode::SUCCESS;
0083 }
0084
0085 ProcessCode RootSpacePointWriter::writeT(
0086     const AlgorithmContext& ctx, const SpacePointContainer& spacePoints) {
0087   // ensure exclusive access to tree/file while writing
0088   std::lock_guard<std::mutex> lock(m_writeMutex);
0089
0090   MeasurementParticlesMap emptyMeasPartMap{};
0091   const auto& measPartMap = m_inputMeasurementParticlesMap.isInitialized()
0092                                 ? m_inputMeasurementParticlesMap(ctx)
0093                                 : emptyMeasPartMap;
0094
0095   // Get the event number
0096   m_eventId = ctx.eventNumber;
0097   for (const auto& sp : spacePoints) {
0098     const auto& sl1 = sp.sourceLinks()[0].get<IndexSourceLink>();
0099     m_measurementId1 = sl1.index();
0100     m_geometryId1 = sl1.geometryId().value();
0101     if (sp.sourceLinks().size() == 2) {
0102       const auto& sl2 = sp.sourceLinks()[1].get<IndexSourceLink>();
0103       m_measurementId2 = sl2.index();
0104       m_geometryId2 = sl2.geometryId().value();
0105     }
0106     // A space point is fake if the measurements have no common particle
0107     if (m_inputMeasurementParticlesMap.isInitialized()) {
0108       if (sp.sourceLinks().size() == 1) {
0109         m_fake = false;
0110       } else {
0111         m_fake = true;
0112         auto [p1b, p1e] = measPartMap.equal_range(m_measurementId1);
0113         auto [p2b, p2e] = measPartMap.equal_range(m_measurementId2);
0114         for (auto it1 = p1b; it1 != p1e; ++it1) {
0115           for (auto it2 = p2b; it2 != p2e; ++it2) {
0116             if (*it1 == *it2) {
0117               m_fake = false;
0118             }
0119           }
0120         }
0121       }
0122     }
0123     // write sp position
0124     m_x = sp.x() / Acts::UnitConstants::mm;
0125     m_y = sp.y() / Acts::UnitConstants::mm;
0126     m_z = sp.z() / Acts::UnitConstants::mm;
0127     m_r = sp.r() / Acts::UnitConstants::mm;
0128     m_t = sp.time() / Acts::UnitConstants::ns;
0129     // write sp dimensions
0130     m_var_r =
0131         sp.varianceR() / (Acts::UnitConstants::mm * Acts::UnitConstants::mm);
0132     m_var_z =
0133         sp.varianceZ() / (Acts::UnitConstants::mm * Acts::UnitConstants::mm);
0134     // Fill the tree
0135     m_outputTree->Fill();
0136   }
0137   return ProcessCode::SUCCESS;
0138 }
0139
0140 }  // namespace ActsExamples