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