Root/src/RootMeasurementWriter.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/RootMeasurementWriter.hpp"
0010
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/Plugins/Root/RootMeasurementIo.hpp"
0013 #include "ActsExamples/EventData/AverageSimHits.hpp"
0014 #include "ActsExamples/EventData/Index.hpp"
0015 #include "ActsExamples/EventData/Measurement.hpp"
0016 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0017 #include "ActsExamples/Utilities/Range.hpp"
0018
0019 #include <ios>
0020 #include <limits>
0021 #include <memory>
0022 #include <stdexcept>
0023 #include <utility>
0024
0025 #include <TFile.h>
0026 #include <TTree.h>
0027
0028 namespace {
0029
0030 std::tuple<std::vector<double>, std::vector<double>, std::vector<unsigned int>>
0031 prepareBoundMeasurement(
0032     const ActsExamples::ConstVariableBoundMeasurementProxy& m) {
0033   std::vector<double> measurements = {};
0034   std::vector<double> variances = {};
0035   std::vector<unsigned int> subspaceIndex = {};
0036
0037   for (unsigned int i = 0; i < m.size(); ++i) {
0038     auto ib = m.subspaceIndexVector()[i];
0039
0040     measurements.push_back(m.parameters()[i]);
0041     variances.push_back(m.covariance()(i, i));
0042     subspaceIndex.push_back(static_cast<unsigned int>(ib));
0043   }
0044
0045   return {measurements, variances, subspaceIndex};
0046 }
0047
0048 std::vector<std::tuple<int, int, float>> prepareChannels(
0049     const ActsExamples::Cluster& c) {
0050   std::vector<std::tuple<int, int, float>> channels = {};
0051   for (auto ch : c.channels) {
0052     channels.emplace_back(static_cast<int>(ch.bin[0]),
0053                           static_cast<int>(ch.bin[1]),
0054                           static_cast<float>(ch.activation));
0055   }
0056   return channels;
0057 }
0058
0059 }  // namespace
0060
0061 namespace ActsExamples {
0062
0063 RootMeasurementWriter::RootMeasurementWriter(
0064     const RootMeasurementWriter::Config& config, Acts::Logging::Level level)
0065     : WriterT(config.inputMeasurements, "RootMeasurementWriter", level),
0066       m_cfg(config) {
0067   // Input container for measurements is already checked by base constructor
0068   if (m_cfg.inputSimHits.empty()) {
0069     throw std::invalid_argument("Missing simulated hits input collection");
0070   }
0071   if (m_cfg.inputMeasurementSimHitsMap.empty()) {
0072     throw std::invalid_argument(
0073         "Missing hit-to-simulated-hits map input collection");
0074   }
0075
0076   m_inputClusters.maybeInitialize(m_cfg.inputClusters);
0077   m_inputSimHits.initialize(m_cfg.inputSimHits);
0078   m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
0079
0080   if (m_cfg.surfaceByIdentifier.empty()) {
0081     throw std::invalid_argument("Missing Surface-GeoID association map");
0082   }
0083   // Setup ROOT File
0084   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0085   if (m_outputFile == nullptr) {
0086     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0087   }
0088
0089   m_outputFile->cd();
0090   m_outputTree = new TTree(m_cfg.treeName.c_str(), "Measurements");
0091   m_outputTree->Branch("particles", &m_particles);
0092
0093   Acts::RootMeasurementIo::Config treeCfg{m_cfg.boundIndices,
0094                                           m_cfg.clusterIndices};
0095   m_measurementIo = std::make_unique<Acts::RootMeasurementIo>(treeCfg);
0096   m_measurementIo->connectForWrite(*m_outputTree);
0097 }
0098
0099 RootMeasurementWriter::~RootMeasurementWriter() {
0100   if (m_outputFile != nullptr) {
0101     m_outputFile->Close();
0102   }
0103 }
0104
0105 ProcessCode RootMeasurementWriter::finalize() {
0106   /// Close the file if it's yours
0107   m_outputFile->cd();
0108   m_outputTree->Write();
0109   m_outputFile->Close();
0110
0111   return ProcessCode::SUCCESS;
0112 }
0113
0114 ProcessCode RootMeasurementWriter::writeT(
0115     const AlgorithmContext& ctx, const MeasurementContainer& measurements) {
0116   const auto& simHits = m_inputSimHits(ctx);
0117   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0118
0119   const ClusterContainer* clusters = nullptr;
0120   if (!m_cfg.inputClusters.empty()) {
0121     clusters = &m_inputClusters(ctx);
0122   }
0123
0124   // Exclusive access to the tree while writing
0125   std::lock_guard<std::mutex> lock(m_writeMutex);
0126
0127   for (Index hitIdx = 0u; hitIdx < measurements.size(); ++hitIdx) {
0128     const ConstVariableBoundMeasurementProxy meas =
0129         measurements.getMeasurement(hitIdx);
0130
0131     Acts::GeometryIdentifier geoId = meas.geometryId();
0132     // find the corresponding surface
0133     auto surfaceItr = m_cfg.surfaceByIdentifier.find(geoId);
0134     if (surfaceItr == m_cfg.surfaceByIdentifier.end()) {
0135       continue;
0136     }
0137     const Acts::Surface& surface = *(surfaceItr->second);
0138
0139     // Fill the identification
0140     m_measurementIo->fillIdentification(static_cast<int>(ctx.eventNumber),
0141                                         geoId);
0142
0143     // Find the contributing simulated hits
0144     auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0145     // Use average truth in the case of multiple contributing sim hits
0146     auto [local, pos4, dir] =
0147         averageSimHits(ctx.geoContext, surface, simHits, indices, logger());
0148     Acts::RotationMatrix3 rot =
0149         surface
0150             .referenceFrame(ctx.geoContext, pos4.segment<3>(Acts::ePos0), dir)
0151             .inverse();
0152     std::pair<double, double> angles =
0153         Acts::VectorHelpers::incidentAngles(dir, rot);
0154     for (auto [_, i] : indices) {
0155       m_particles.push_back(simHits.nth(i)->particleId().value());
0156     }
0157     m_measurementIo->fillTruthParameters(local, pos4, dir, angles);
0158
0159     // Fill the measurement parameters & clusters still
0160     auto [msm, vcs, ssi] = prepareBoundMeasurement(meas);
0161     m_measurementIo->fillBoundMeasurement(msm, vcs, ssi);
0162
0163     // Fill the cluster information if available
0164     if (clusters != nullptr) {
0165       const auto& cluster = (*clusters)[hitIdx];
0166       m_measurementIo->fillGlobalPosition(cluster.globalPosition);
0167       m_measurementIo->fillCluster(prepareChannels(cluster));
0168     }
0169
0170     m_outputTree->Fill();
0171     m_particles.clear();
0172     m_measurementIo->clear();
0173   }
0174
0175   return ProcessCode::SUCCESS;
0176 }
0177
0178 }  // namespace ActsExamples