Root/src/RootMeasurementWriter.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2017 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 http://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Io/Root/RootMeasurementWriter.hpp"
0010
0011 #include "Acts/Geometry/TrackingGeometry.hpp"
0012 #include "ActsExamples/EventData/AverageSimHits.hpp"
0013 #include "ActsExamples/EventData/Index.hpp"
0014 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0015 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0016 #include "ActsExamples/Utilities/Range.hpp"
0017
0018 #include <cstddef>
0019 #include <ios>
0020 #include <stdexcept>
0021 #include <utility>
0022 #include <variant>
0023
0024 #include <TFile.h>
0025
0026 namespace Acts {
0027 class Surface;
0028 }  // namespace Acts
0029
0030 ActsExamples::RootMeasurementWriter::RootMeasurementWriter(
0031     const ActsExamples::RootMeasurementWriter::Config& config,
0032     Acts::Logging::Level level)
0033     : WriterT(config.inputMeasurements, "RootMeasurementWriter", level),
0034       m_cfg(config) {
0035   // Input container for measurements is already checked by base constructor
0036   if (m_cfg.inputSimHits.empty()) {
0037     throw std::invalid_argument("Missing simulated hits input collection");
0038   }
0039   if (m_cfg.inputMeasurementSimHitsMap.empty()) {
0040     throw std::invalid_argument(
0041         "Missing hit-to-simulated-hits map input collection");
0042   }
0043
0044   m_inputSimHits.initialize(m_cfg.inputSimHits);
0045   m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
0046   m_inputClusters.maybeInitialize(m_cfg.inputClusters);
0047
0048   if (!m_cfg.trackingGeometry) {
0049     throw std::invalid_argument("Missing tracking geometry");
0050   }
0051   // Setup ROOT File
0052   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0053   if (m_outputFile == nullptr) {
0054     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0055   }
0056
0057   m_outputFile->cd();
0058
0059   // Analyze the smearers
0060   std::vector<
0061       std::pair<Acts::GeometryIdentifier, std::unique_ptr<DigitizationTree>>>
0062       dTrees;
0063   if (!m_cfg.boundIndices.empty()) {
0064     ACTS_DEBUG("Bound indices are declared, preparing trees.");
0065     for (std::size_t ikv = 0; ikv < m_cfg.boundIndices.size(); ++ikv) {
0066       auto geoID = m_cfg.boundIndices.idAt(ikv);
0067       auto bIndices = m_cfg.boundIndices.valueAt(ikv);
0068       auto dTree = std::make_unique<DigitizationTree>(geoID);
0069       for (const auto& bIndex : bIndices) {
0070         ACTS_VERBOSE("- setup branch for index: " << bIndex);
0071         dTree->setupBoundRecBranch(bIndex);
0072       }
0073       if (!m_cfg.inputClusters.empty()) {
0074         dTree->setupClusterBranch(bIndices);
0075       }
0076       dTrees.push_back({geoID, std::move(dTree)});
0077     }
0078   } else {
0079     ACTS_DEBUG("Bound indices are not declared, no reco setup.")
0080   }
0081
0082   m_outputTrees = Acts::GeometryHierarchyMap<std::unique_ptr<DigitizationTree>>(
0083       std::move(dTrees));
0084 }
0085
0086 ActsExamples::RootMeasurementWriter::~RootMeasurementWriter() {
0087   if (m_outputFile != nullptr) {
0088     m_outputFile->Close();
0089   }
0090 }
0091
0092 ActsExamples::ProcessCode ActsExamples::RootMeasurementWriter::finalize() {
0093   /// Close the file if it's yours
0094   m_outputFile->cd();
0095   for (auto dTree = m_outputTrees.begin(); dTree != m_outputTrees.end();
0096        ++dTree) {
0097     (*dTree)->tree->Write();
0098   }
0099   m_outputFile->Close();
0100
0101   return ProcessCode::SUCCESS;
0102 }
0103
0104 ActsExamples::ProcessCode ActsExamples::RootMeasurementWriter::writeT(
0105     const AlgorithmContext& ctx, const MeasurementContainer& measurements) {
0106   const auto& simHits = m_inputSimHits(ctx);
0107   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0108
0109   ClusterContainer clusters;
0110   if (!m_cfg.inputClusters.empty()) {
0111     clusters = m_inputClusters(ctx);
0112   }
0113
0114   // Exclusive access to the tree while writing
0115   std::lock_guard<std::mutex> lock(m_writeMutex);
0116
0117   for (Index hitIdx = 0u; hitIdx < measurements.size(); ++hitIdx) {
0118     const auto& meas = measurements[hitIdx];
0119
0120     std::visit(
0121         [&](const auto& m) {
0122           Acts::GeometryIdentifier geoId =
0123               m.sourceLink().template get<IndexSourceLink>().geometryId();
0124           // find the corresponding surface
0125           const Acts::Surface* surfacePtr =
0126               m_cfg.trackingGeometry->findSurface(geoId);
0127           if (!surfacePtr) {
0128             return;
0129           }
0130           const Acts::Surface& surface = *surfacePtr;
0131           // find the corresponding output tree
0132           auto dTreeItr = m_outputTrees.find(geoId);
0133           if (dTreeItr == m_outputTrees.end()) {
0134             return;
0135           }
0136           auto& dTree = *dTreeItr;
0137
0138           // Fill the identification
0139           dTree->fillIdentification(ctx.eventNumber, geoId);
0140
0141           // Find the contributing simulated hits
0142           auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0143           // Use average truth in the case of multiple contributing sim hits
0144           auto [local, pos4, dir] = averageSimHits(ctx.geoContext, surface,
0145                                                    simHits, indices, logger());
0146           Acts::RotationMatrix3 rot =
0147               surface
0148                   .referenceFrame(ctx.geoContext, pos4.segment<3>(Acts::ePos0),
0149                                   dir)
0150                   .inverse();
0151           std::pair<double, double> angles =
0152               Acts::VectorHelpers::incidentAngles(dir, rot);
0153           dTree->fillTruthParameters(local, pos4, dir, angles);
0154           dTree->fillBoundMeasurement(m);
0155           if (!clusters.empty()) {
0156             const auto& c = clusters[hitIdx];
0157             dTree->fillCluster(c);
0158           }
0159           dTree->tree->Fill();
0160           if (dTree->chValue != nullptr) {
0161             dTree->chValue->clear();
0162           }
0163           if (dTree->chId[0] != nullptr) {
0164             dTree->chId[0]->clear();
0165           }
0166           if (dTree->chId[1] != nullptr) {
0167             dTree->chId[1]->clear();
0168           }
0169         },
0170         meas);
0171   }
0172
0173   return ActsExamples::ProcessCode::SUCCESS;
0174 }