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 "ActsExamples/EventData/AverageSimHits.hpp"
0013 #include "ActsExamples/EventData/Index.hpp"
0014 #include "ActsExamples/EventData/Measurement.hpp"
0015 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0016 #include "ActsExamples/Utilities/Range.hpp"
0017 #include "ActsPlugins/Root/RootMeasurementIo.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_vertex_primary", &m_particleVertexPrimary);
0092   m_outputTree->Branch("particles_vertex_secondary",
0093                        &m_particleVertexSecondary);
0094   m_outputTree->Branch("particles_particle", &m_particleParticle);
0095   m_outputTree->Branch("particles_generation", &m_particleGeneration);
0096   m_outputTree->Branch("particles_sub_particle", &m_particleSubParticle);
0097
0098   ActsPlugins::RootMeasurementIo::Config treeCfg{m_cfg.boundIndices,
0099                                                  m_cfg.clusterIndices};
0100   m_measurementIo = std::make_unique<ActsPlugins::RootMeasurementIo>(treeCfg);
0101   m_measurementIo->connectForWrite(*m_outputTree);
0102 }
0103
0104 RootMeasurementWriter::~RootMeasurementWriter() {
0105   if (m_outputFile != nullptr) {
0106     m_outputFile->Close();
0107   }
0108 }
0109
0110 ProcessCode RootMeasurementWriter::finalize() {
0111   /// Close the file if it's yours
0112   m_outputFile->cd();
0113   m_outputTree->Write();
0114   m_outputFile->Close();
0115
0116   return ProcessCode::SUCCESS;
0117 }
0118
0119 ProcessCode RootMeasurementWriter::writeT(
0120     const AlgorithmContext& ctx, const MeasurementContainer& measurements) {
0121   const auto& simHits = m_inputSimHits(ctx);
0122   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0123
0124   const ClusterContainer* clusters = nullptr;
0125   if (!m_cfg.inputClusters.empty()) {
0126     clusters = &m_inputClusters(ctx);
0127   }
0128
0129   // Exclusive access to the tree while writing
0130   std::lock_guard<std::mutex> lock(m_writeMutex);
0131
0132   for (Index hitIdx = 0u; hitIdx < measurements.size(); ++hitIdx) {
0133     const ConstVariableBoundMeasurementProxy meas =
0134         measurements.getMeasurement(hitIdx);
0135
0136     Acts::GeometryIdentifier geoId = meas.geometryId();
0137     // find the corresponding surface
0138     auto surfaceItr = m_cfg.surfaceByIdentifier.find(geoId);
0139     if (surfaceItr == m_cfg.surfaceByIdentifier.end()) {
0140       continue;
0141     }
0142     const Acts::Surface& surface = *(surfaceItr->second);
0143
0144     // Fill the identification
0145     m_measurementIo->fillIdentification(static_cast<int>(ctx.eventNumber),
0146                                         geoId);
0147
0148     // Find the contributing simulated hits
0149     auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0150     // Use average truth in the case of multiple contributing sim hits
0151     auto [local, pos4, dir] =
0152         averageSimHits(ctx.geoContext, surface, simHits, indices, logger());
0153     Acts::RotationMatrix3 rot =
0154         surface
0155             .referenceFrame(ctx.geoContext, pos4.segment<3>(Acts::ePos0), dir)
0156             .inverse();
0157     std::pair<double, double> angles =
0158         Acts::VectorHelpers::incidentAngles(dir, rot);
0159     for (auto [_, i] : indices) {
0160       const auto barcode = simHits.nth(i)->particleId();
0161       m_particleVertexPrimary.push_back(barcode.vertexPrimary());
0162       m_particleVertexSecondary.push_back(barcode.vertexSecondary());
0163       m_particleParticle.push_back(barcode.particle());
0164       m_particleGeneration.push_back(barcode.generation());
0165       m_particleSubParticle.push_back(barcode.subParticle());
0166     }
0167     m_measurementIo->fillTruthParameters(local, pos4, dir, angles);
0168
0169     // Fill the measurement parameters & clusters still
0170     auto [msm, vcs, ssi] = prepareBoundMeasurement(meas);
0171     m_measurementIo->fillBoundMeasurement(msm, vcs, ssi);
0172
0173     // Fill the cluster information if available
0174     if (clusters != nullptr) {
0175       const auto& cluster = (*clusters)[hitIdx];
0176       m_measurementIo->fillGlobalPosition(cluster.globalPosition);
0177       m_measurementIo->fillCluster(prepareChannels(cluster));
0178     }
0179
0180     m_outputTree->Fill();
0181     m_particleVertexPrimary.clear();
0182     m_particleVertexSecondary.clear();
0183     m_particleParticle.clear();
0184     m_particleGeneration.clear();
0185     m_particleSubParticle.clear();
0186     m_measurementIo->clear();
0187   }
0188
0189   return ProcessCode::SUCCESS;
0190 }
0191
0192 }  // namespace ActsExamples