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 }  // namespace
0048
0049 RootMeasurementWriter::RootMeasurementWriter(
0050     const RootMeasurementWriter::Config& config, Acts::Logging::Level level)
0051     : WriterT(config.inputMeasurements, "RootMeasurementWriter", level),
0052       m_cfg(config) {
0053   // Input container for measurements is already checked by base constructor
0054   if (m_cfg.inputSimHits.empty()) {
0055     throw std::invalid_argument("Missing simulated hits input collection");
0056   }
0057   if (m_cfg.inputMeasurementSimHitsMap.empty()) {
0058     throw std::invalid_argument(
0059         "Missing hit-to-simulated-hits map input collection");
0060   }
0061
0062   m_inputSimHits.initialize(m_cfg.inputSimHits);
0063   m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
0064
0065   if (m_cfg.surfaceByIdentifier.empty()) {
0066     throw std::invalid_argument("Missing Surface-GeoID association map");
0067   }
0068   // Setup ROOT File
0069   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0070   if (m_outputFile == nullptr) {
0071     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0072   }
0073
0074   m_outputFile->cd();
0075   m_outputTree = new TTree(m_cfg.treeName.c_str(), "Measurements");
0076   m_outputTree->Branch("particles_vertex_primary", &m_particleVertexPrimary);
0077   m_outputTree->Branch("particles_vertex_secondary",
0078                        &m_particleVertexSecondary);
0079   m_outputTree->Branch("particles_particle", &m_particleParticle);
0080   m_outputTree->Branch("particles_generation", &m_particleGeneration);
0081   m_outputTree->Branch("particles_sub_particle", &m_particleSubParticle);
0082
0083   ActsPlugins::RootMeasurementIo::Config treeCfg{m_cfg.boundIndices};
0084   m_measurementIo = std::make_unique<ActsPlugins::RootMeasurementIo>(treeCfg);
0085   m_measurementIo->connectForWrite(*m_outputTree);
0086 }
0087
0088 RootMeasurementWriter::~RootMeasurementWriter() {
0089   if (m_outputFile != nullptr) {
0090     m_outputFile->Close();
0091   }
0092 }
0093
0094 ProcessCode RootMeasurementWriter::finalize() {
0095   /// Close the file if it's yours
0096   m_outputFile->cd();
0097   m_outputTree->Write();
0098   m_outputFile->Close();
0099
0100   return ProcessCode::SUCCESS;
0101 }
0102
0103 ProcessCode RootMeasurementWriter::writeT(
0104     const AlgorithmContext& ctx, const MeasurementContainer& measurements) {
0105   const auto& simHits = m_inputSimHits(ctx);
0106   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0107
0108   // Exclusive access to the tree while writing
0109   std::lock_guard<std::mutex> lock(m_writeMutex);
0110
0111   for (Index hitIdx = 0u; hitIdx < measurements.size(); ++hitIdx) {
0112     const ConstVariableBoundMeasurementProxy meas =
0113         measurements.getMeasurement(hitIdx);
0114
0115     Acts::GeometryIdentifier geoId = meas.geometryId();
0116     // find the corresponding surface
0117     auto surfaceItr = m_cfg.surfaceByIdentifier.find(geoId);
0118     if (surfaceItr == m_cfg.surfaceByIdentifier.end()) {
0119       continue;
0120     }
0121     const Acts::Surface& surface = *(surfaceItr->second);
0122
0123     // Fill the identification
0124     m_measurementIo->fillIdentification(static_cast<int>(ctx.eventNumber),
0125                                         hitIdx, geoId);
0126
0127     // Find the contributing simulated hits
0128     auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0129     // Use average truth in the case of multiple contributing sim hits
0130     auto [local, pos4, dir] =
0131         averageSimHits(ctx.simGeoContext, surface, simHits, indices, logger());
0132     Acts::RotationMatrix3 rot =
0133         surface
0134             .referenceFrame(ctx.simGeoContext, pos4.segment<3>(Acts::ePos0),
0135                             dir)
0136             .inverse();
0137     std::pair<double, double> angles =
0138         Acts::VectorHelpers::incidentAngles(dir, rot);
0139     for (auto [_, i] : indices) {
0140       const auto barcode = simHits.nth(i)->particleId();
0141       m_particleVertexPrimary.push_back(barcode.vertexPrimary());
0142       m_particleVertexSecondary.push_back(barcode.vertexSecondary());
0143       m_particleParticle.push_back(barcode.particle());
0144       m_particleGeneration.push_back(barcode.generation());
0145       m_particleSubParticle.push_back(barcode.subParticle());
0146     }
0147     m_measurementIo->fillTruthParameters(local, pos4, dir, angles);
0148
0149     // Fill the measurement parameters
0150     auto [msm, vcs, ssi] = prepareBoundMeasurement(meas);
0151     m_measurementIo->fillBoundMeasurement(msm, vcs, ssi);
0152
0153     m_outputTree->Fill();
0154     m_particleVertexPrimary.clear();
0155     m_particleVertexSecondary.clear();
0156     m_particleParticle.clear();
0157     m_particleGeneration.clear();
0158     m_particleSubParticle.clear();
0159     m_measurementIo->clear();
0160   }
0161
0162   return ProcessCode::SUCCESS;
0163 }
0164
0165 }  // namespace ActsExamples