Root/src/RootPlanarClusterWriter.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/RootPlanarClusterWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Definitions/TrackParametrization.hpp"
0013 #include "Acts/Definitions/Units.hpp"
0014 #include "Acts/Digitization/DigitizationCell.hpp"
0015 #include "Acts/Digitization/DigitizationModule.hpp"
0016 #include "Acts/Digitization/DigitizationSourceLink.hpp"
0017 #include "Acts/Digitization/PlanarModuleCluster.hpp"
0018 #include "Acts/Digitization/Segmentation.hpp"
0019 #include "Acts/EventData/SourceLink.hpp"
0020 #include "Acts/Geometry/GeometryIdentifier.hpp"
0021 #include "Acts/Geometry/TrackingGeometry.hpp"
0022 #include "Acts/Plugins/Identification/IdentifiedDetectorElement.hpp"
0023 #include "Acts/Surfaces/Surface.hpp"
0024 #include "Acts/Utilities/Result.hpp"
0025 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0026 #include "ActsExamples/Utilities/GroupBy.hpp"
0027 #include "ActsExamples/Utilities/Range.hpp"
0028 #include "ActsFatras/EventData/Barcode.hpp"
0029 #include "ActsFatras/EventData/Hit.hpp"
0030
0031 #include <cstdint>
0032 #include <ios>
0033 #include <ostream>
0034 #include <stdexcept>
0035 #include <utility>
0036
0037 #include <TFile.h>
0038 #include <TTree.h>
0039
0040 ActsExamples::RootPlanarClusterWriter::RootPlanarClusterWriter(
0041     const ActsExamples::RootPlanarClusterWriter::Config& config,
0042     Acts::Logging::Level level)
0043     : WriterT(config.inputClusters, "RootPlanarClusterWriter", level),
0044       m_cfg(config) {
0045   // inputClusters is already checked by base constructor
0046   if (m_cfg.inputSimHits.empty()) {
0047     throw std::invalid_argument("Missing simulated hits input collection");
0048   }
0049
0050   m_inputSimHits.initialize(m_cfg.inputSimHits);
0051
0052   if (m_cfg.treeName.empty()) {
0053     throw std::invalid_argument("Missing tree name");
0054   }
0055   if (!m_cfg.trackingGeometry) {
0056     throw std::invalid_argument("Missing tracking geometry");
0057   }
0058   // Setup ROOT I/O
0059   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0060   if (m_outputFile == nullptr) {
0061     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0062   }
0063   m_outputFile->cd();
0064   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0065   if (m_outputTree == nullptr) {
0066     throw std::bad_alloc();
0067   }
0068
0069   // Set the branches
0070   m_outputTree->Branch("event_nr", &m_eventNr);
0071   m_outputTree->Branch("volume_id", &m_volumeID);
0072   m_outputTree->Branch("layer_id", &m_layerID);
0073   m_outputTree->Branch("surface_id", &m_surfaceID);
0074   m_outputTree->Branch("g_x", &m_x);
0075   m_outputTree->Branch("g_y", &m_y);
0076   m_outputTree->Branch("g_z", &m_z);
0077   m_outputTree->Branch("g_t", &m_t);
0078   m_outputTree->Branch("l_x", &m_lx);
0079   m_outputTree->Branch("l_y", &m_ly);
0080   m_outputTree->Branch("cov_l_x", &m_cov_lx);
0081   m_outputTree->Branch("cov_l_y", &m_cov_ly);
0082   m_outputTree->Branch("cell_ID_x", &m_cell_IDx);
0083   m_outputTree->Branch("cell_ID_y", &m_cell_IDy);
0084   m_outputTree->Branch("cell_l_x", &m_cell_lx);
0085   m_outputTree->Branch("cell_l_y", &m_cell_ly);
0086   m_outputTree->Branch("cell_data", &m_cell_data);
0087   m_outputTree->Branch("truth_g_x", &m_t_gx);
0088   m_outputTree->Branch("truth_g_y", &m_t_gy);
0089   m_outputTree->Branch("truth_g_z", &m_t_gz);
0090   m_outputTree->Branch("truth_g_t", &m_t_gt);
0091   m_outputTree->Branch("truth_l_x", &m_t_lx);
0092   m_outputTree->Branch("truth_l_y", &m_t_ly);
0093   m_outputTree->Branch("truth_barcode", &m_t_barcode);
0094 }
0095
0096 ActsExamples::RootPlanarClusterWriter::~RootPlanarClusterWriter() {
0097   if (m_outputFile != nullptr) {
0098     m_outputFile->Close();
0099   }
0100 }
0101
0102 ActsExamples::ProcessCode ActsExamples::RootPlanarClusterWriter::finalize() {
0103   // Write the tree
0104   m_outputFile->cd();
0105   m_outputTree->Write();
0106   m_outputFile->Close();
0107
0108   ACTS_INFO("Wrote clusters to tree '" << m_cfg.treeName << "' in '"
0109                                        << m_cfg.filePath << "'");
0110
0111   return ProcessCode::SUCCESS;
0112 }
0113
0114 ActsExamples::ProcessCode ActsExamples::RootPlanarClusterWriter::writeT(
0115     const AlgorithmContext& ctx,
0116     const ActsExamples::GeometryIdMultimap<Acts::PlanarModuleCluster>&
0117         clusters) {
0118   // retrieve simulated hits
0119   const auto& simHits = m_inputSimHits(ctx);
0120
0121   // Exclusive access to the tree while writing
0122   std::lock_guard<std::mutex> lock(m_writeMutex);
0123   // Get the event number
0124   m_eventNr = ctx.eventNumber;
0125
0126   // Loop over the planar clusters in this event
0127   for (auto [moduleGeoId, moduleClusters] : groupByModule(clusters)) {
0128     const Acts::Surface* surfacePtr =
0129         m_cfg.trackingGeometry->findSurface(moduleGeoId);
0130     if (surfacePtr == nullptr) {
0131       ACTS_ERROR("Could not find surface for " << moduleGeoId);
0132       return ProcessCode::ABORT;
0133     }
0134     const Acts::Surface& surface = *surfacePtr;
0135
0136     // geometry identification is the same for all clusters on this module
0137     m_volumeID = moduleGeoId.volume();
0138     m_layerID = moduleGeoId.layer();
0139     m_surfaceID = moduleGeoId.sensitive();
0140
0141     for (const auto& entry : moduleClusters) {
0142       const Acts::PlanarModuleCluster& cluster = entry.second;
0143       // local cluster information: position, @todo coveraiance
0144       auto parameters = cluster.parameters();
0145       Acts::Vector2 local(parameters[Acts::BoundIndices::eBoundLoc0],
0146                           parameters[Acts::BoundIndices::eBoundLoc1]);
0147
0148       /// prepare for calculating the
0149       Acts::Vector3 mom(1, 1, 1);
0150       // transform local into global position information
0151       Acts::Vector3 pos = surface.localToGlobal(ctx.geoContext, local, mom);
0152       m_x = pos.x() / Acts::UnitConstants::mm;
0153       m_y = pos.y() / Acts::UnitConstants::mm;
0154       m_z = pos.z() / Acts::UnitConstants::mm;
0155       m_t = parameters[2] / Acts::UnitConstants::mm;
0156       m_lx = local.x();
0157       m_ly = local.y();
0158       m_cov_lx = 0.;  // @todo fill in
0159       m_cov_ly = 0.;  // @todo fill in
0160       // get the cells and run through them
0161       const auto& cells = cluster.digitizationCells();
0162       auto detectorElement =
0163           dynamic_cast<const Acts::IdentifiedDetectorElement*>(
0164               surface.associatedDetectorElement());
0165       for (auto& cell : cells) {
0166         // cell identification
0167         m_cell_IDx.push_back(cell.channel0);
0168         m_cell_IDy.push_back(cell.channel1);
0169         m_cell_data.push_back(cell.data);
0170         // for more we need the digitization module
0171         if ((detectorElement != nullptr) &&
0172             detectorElement->digitizationModule()) {
0173           auto digitationModule = detectorElement->digitizationModule();
0174           const Acts::Segmentation& segmentation =
0175               digitationModule->segmentation();
0176           // get the cell positions
0177           auto cellLocalPosition = segmentation.cellPosition(cell);
0178           m_cell_lx.push_back(cellLocalPosition.x());
0179           m_cell_ly.push_back(cellLocalPosition.y());
0180         }
0181       }
0182       // write hit-particle truth association
0183       // each hit can have multiple particles, e.g. in a dense environment
0184       const auto& sl = cluster.sourceLink().get<Acts::DigitizationSourceLink>();
0185       for (auto idx : sl.indices()) {
0186         auto it = simHits.nth(idx);
0187         if (it == simHits.end()) {
0188           ACTS_FATAL("Simulation hit with index " << idx << " does not exist");
0189           return ProcessCode::ABORT;
0190         }
0191         const auto& simHit = *it;
0192
0193         // local position to be calculated
0194         Acts::Vector2 lPosition{0., 0.};
0195         auto lpResult = surface.globalToLocal(ctx.geoContext, simHit.position(),
0196                                               simHit.direction());
0197         if (!lpResult.ok()) {
0198           ACTS_FATAL("Global to local transformation did not succeed.");
0199           return ProcessCode::ABORT;
0200         }
0201         lPosition = lpResult.value();
0202         // fill the variables
0203         m_t_gx.push_back(simHit.position().x());
0204         m_t_gy.push_back(simHit.position().y());
0205         m_t_gz.push_back(simHit.position().z());
0206         m_t_gt.push_back(simHit.time());
0207         m_t_lx.push_back(lPosition.x());
0208         m_t_ly.push_back(lPosition.y());
0209         m_t_barcode.push_back(simHit.particleId().value());
0210       }
0211       // fill the tree
0212       m_outputTree->Fill();
0213       // now reset
0214       m_cell_IDx.clear();
0215       m_cell_IDy.clear();
0216       m_cell_lx.clear();
0217       m_cell_ly.clear();
0218       m_cell_data.clear();
0219       m_t_gx.clear();
0220       m_t_gy.clear();
0221       m_t_gz.clear();
0222       m_t_gt.clear();
0223       m_t_lx.clear();
0224       m_t_ly.clear();
0225       m_t_barcode.clear();
0226     }
0227   }
0228   return ActsExamples::ProcessCode::SUCCESS;
0229 }