Csv/src/CsvTrackParameterWriter.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2022 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/Csv/CsvTrackParameterWriter.hpp"
0010
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/EventData/GenericBoundTrackParameters.hpp"
0013 #include "ActsExamples/EventData/Trajectories.hpp"
0014 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0015 #include "ActsExamples/Utilities/Paths.hpp"
0016
0017 #include <optional>
0018 #include <stdexcept>
0019
0020 #include <dfe/dfe_io_dsv.hpp>
0021
0022 #include "CsvOutputData.hpp"
0023
0024 ActsExamples::CsvTrackParameterWriter::CsvTrackParameterWriter(
0025     const ActsExamples::CsvTrackParameterWriter::Config& config,
0026     Acts::Logging::Level level)
0027     : m_cfg(config),
0028       m_logger(Acts::getDefaultLogger("CsvTrackParameterWriter", level)) {
0029   if (m_cfg.inputTrackParameters.empty() == m_cfg.inputTracks.empty()) {
0030     throw std::invalid_argument(
0031         "You have to either provide track parameters or tracks");
0032   }
0033
0034   m_inputTrackParameters.maybeInitialize(m_cfg.inputTrackParameters);
0035   m_inputTracks.maybeInitialize(m_cfg.inputTracks);
0036 }
0037
0038 ActsExamples::CsvTrackParameterWriter::~CsvTrackParameterWriter() = default;
0039
0040 std::string ActsExamples::CsvTrackParameterWriter::name() const {
0041   return "CsvTrackParameterWriter";
0042 }
0043
0044 ActsExamples::ProcessCode ActsExamples::CsvTrackParameterWriter::finalize() {
0045   // Write the tree
0046   return ProcessCode::SUCCESS;
0047 }
0048
0049 ActsExamples::ProcessCode ActsExamples::CsvTrackParameterWriter::write(
0050     const AlgorithmContext& ctx) {
0051   TrackParametersContainer inputTrackParameters;
0052
0053   if (!m_cfg.inputTrackParameters.empty()) {
0054     const auto& tmp = m_inputTrackParameters(ctx);
0055     inputTrackParameters = tmp;
0056   } else {
0057     const auto& inputTracks = m_inputTracks(ctx);
0058
0059     for (const auto& track : inputTracks) {
0060       if (!track.hasReferenceSurface()) {
0061         continue;
0062       }
0063       auto trackParam = track.createParametersAtReference();
0064       inputTrackParameters.push_back(trackParam);
0065     }
0066   }
0067
0068   std::string path =
0069       perEventFilepath(m_cfg.outputDir, m_cfg.outputStem, ctx.eventNumber);
0070
0071   dfe::NamedTupleCsvWriter<TrackParameterData> writer(path,
0072                                                       m_cfg.outputPrecision);
0073
0074   TrackParameterData data{};
0075   for (const auto& tp : inputTrackParameters) {
0076     const auto& params = tp.parameters();
0077     const auto& cov = tp.covariance().value();
0078
0079     data.d0 = params[Acts::eBoundLoc0];
0080     data.z0 = params[Acts::eBoundLoc1];
0081     data.phi = params[Acts::eBoundPhi];
0082     data.theta = params[Acts::eBoundTheta];
0083     data.qop = params[Acts::eBoundQOverP];
0084
0085     data.var_d0 = cov(Acts::eBoundLoc0, Acts::eBoundLoc0);
0086     data.var_z0 = cov(Acts::eBoundLoc1, Acts::eBoundLoc1);
0087     data.var_phi = cov(Acts::eBoundPhi, Acts::eBoundPhi);
0088     data.var_theta = cov(Acts::eBoundTheta, Acts::eBoundTheta);
0089     data.var_qop = cov(Acts::eBoundQOverP, Acts::eBoundQOverP);
0090
0091     data.cov_d0z0 = cov(Acts::eBoundLoc0, Acts::eBoundLoc1);
0092     data.cov_d0phi = cov(Acts::eBoundLoc0, Acts::eBoundPhi);
0093     data.cov_d0theta = cov(Acts::eBoundLoc0, Acts::eBoundTheta);
0094     data.cov_d0qop = cov(Acts::eBoundLoc0, Acts::eBoundQOverP);
0095
0096     data.cov_z0d0 = cov(Acts::eBoundLoc1, Acts::eBoundLoc0);
0097     data.cov_z0phi = cov(Acts::eBoundLoc1, Acts::eBoundPhi);
0098     data.cov_z0theta = cov(Acts::eBoundLoc1, Acts::eBoundTheta);
0099     data.cov_z0qop = cov(Acts::eBoundLoc1, Acts::eBoundQOverP);
0100
0101     data.cov_phid0 = cov(Acts::eBoundPhi, Acts::eBoundLoc0);
0102     data.cov_phiz0 = cov(Acts::eBoundPhi, Acts::eBoundLoc1);
0103     data.cov_phitheta = cov(Acts::eBoundPhi, Acts::eBoundTheta);
0104     data.cov_phiqop = cov(Acts::eBoundPhi, Acts::eBoundQOverP);
0105
0106     data.cov_thetad0 = cov(Acts::eBoundTheta, Acts::eBoundLoc0);
0107     data.cov_thetaz0 = cov(Acts::eBoundTheta, Acts::eBoundLoc1);
0108     data.cov_thetaphi = cov(Acts::eBoundTheta, Acts::eBoundPhi);
0109     data.cov_thetaqop = cov(Acts::eBoundTheta, Acts::eBoundQOverP);
0110
0111     data.cov_qopd0 = cov(Acts::eBoundQOverP, Acts::eBoundLoc0);
0112     data.cov_qopz0 = cov(Acts::eBoundQOverP, Acts::eBoundLoc1);
0113     data.cov_qopphi = cov(Acts::eBoundQOverP, Acts::eBoundPhi);
0114     data.cov_qoptheta = cov(Acts::eBoundQOverP, Acts::eBoundTheta);
0115
0116     writer.append(data);
0117   }
0118
0119   return ActsExamples::ProcessCode::SUCCESS;
0120 }