feat: examples CSV track param writer (#1717)

adding a track param writer

similar to #1678

Examples/Io/Csv/CMakeLists.txt

@@ -13,6 +13,7 @@ add_library(
   src/CsvSpacePointReader.cpp
   src/CsvTrackingGeometryWriter.cpp
   src/CsvTrackParameterReader.cpp
+  src/CsvTrackParameterWriter.cpp
   src/CsvMultiTrajectoryWriter.cpp 
   src/CsvSpacePointWriter.cpp
   src/CsvBFieldWriter.cpp)

Examples/Io/Csv/include/ActsExamples/Io/Csv/CsvTrackParameterReader.hpp

@@ -8,8 +8,8 @@
 
 #pragma once
 
+#include "Acts/Utilities/Logger.hpp"
 #include "ActsExamples/Framework/IReader.hpp"
-#include <Acts/Utilities/Logger.hpp>
 
 #include <array>
 #include <memory>

Examples/Io/Csv/include/ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp

@@ -0,0 +1,68 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2022 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include "Acts/Utilities/Helpers.hpp"
+#include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/Framework/IWriter.hpp"
+
+#include <string>
+
+namespace ActsExamples {
+
+/// Write track parameters in comma-separated-value format.
+///
+/// This writes one file per event in the configured input directory
+/// and filename. Files are assumed to be named using the following schema
+///
+///     event000000001-<stem>.csv
+///     event000000002-<stem>.csv
+///
+/// and each line in the file corresponds to one track parameter.
+class CsvTrackParameterWriter final : public IWriter {
+ public:
+  struct Config {
+    /// Optional. Input track parameters collection
+    std::string inputTrackParameters;
+    /// Optional. Input trajectories container.
+    std::string inputTrajectories;
+    /// Where to place output files
+    std::string outputDir;
+    /// Input filename stem.
+    std::string outputStem;
+    /// Number of decimal digits for floating point precision in output.
+    size_t outputPrecision = std::numeric_limits<float>::max_digits10;
+  };
+
+  /// Constructor with
+  /// @param config configuration struct
+  /// @param level logging level
+  CsvTrackParameterWriter(const Config& config, Acts::Logging::Level level);
+
+  /// Virtual destructor
+  ~CsvTrackParameterWriter() override;
+
+  /// Provide the name of the writer
+  std::string name() const override;
+
+  /// Read the object and call the type-specific member function.
+  ProcessCode write(const AlgorithmContext& ctx) override;
+
+  /// No-op default implementation.
+  ProcessCode endRun() override;
+
+  /// Get readonly access to the config parameters
+  const Config& config() const { return m_cfg; }
+
+ private:
+  Config m_cfg;
+  std::unique_ptr<const Acts::Logger> m_logger;
+};
+
+}  // namespace ActsExamples

Examples/Io/Csv/src/CsvTrackParameterWriter.cpp

@@ -0,0 +1,124 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2022 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include "ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp"
+
+#include "Acts/Definitions/TrackParametrization.hpp"
+#include "Acts/Definitions/Units.hpp"
+#include "ActsExamples/EventData/Track.hpp"
+#include "ActsExamples/EventData/Trajectories.hpp"
+#include "ActsExamples/Framework/WhiteBoard.hpp"
+#include "ActsExamples/Utilities/Paths.hpp"
+
+#include <ios>
+#include <optional>
+#include <stdexcept>
+
+#include <dfe/dfe_io_dsv.hpp>
+
+#include "CsvOutputData.hpp"
+
+ActsExamples::CsvTrackParameterWriter::CsvTrackParameterWriter(
+    const ActsExamples::CsvTrackParameterWriter::Config& config,
+    Acts::Logging::Level level)
+    : m_cfg(config),
+      m_logger(Acts::getDefaultLogger("CsvTrackParameterWriter", level)) {
+  if (m_cfg.inputTrackParameters.empty() == m_cfg.inputTrajectories.empty()) {
+    throw std::invalid_argument(
+        "You have to either provide track parameters or trajectories");
+  }
+}
+
+ActsExamples::CsvTrackParameterWriter::~CsvTrackParameterWriter() = default;
+
+std::string ActsExamples::CsvTrackParameterWriter::name() const {
+  return "CsvTrackParameterWriter";
+}
+
+ActsExamples::ProcessCode ActsExamples::CsvTrackParameterWriter::endRun() {
+  // Write the tree
+  return ProcessCode::SUCCESS;
+}
+
+ActsExamples::ProcessCode ActsExamples::CsvTrackParameterWriter::write(
+    const AlgorithmContext& ctx) {
+  std::vector<Acts::BoundTrackParameters> inputTrackParameters;
+
+  if (!m_cfg.inputTrackParameters.empty()) {
+    const auto& tmp =
+        ctx.eventStore.get<std::vector<Acts::BoundTrackParameters>>(
+            m_cfg.inputTrackParameters);
+    inputTrackParameters = tmp;
+  } else {
+    const auto& inputTrajectories =
+        ctx.eventStore.get<TrajectoriesContainer>(m_cfg.inputTrajectories);
+
+    for (const auto& trajectories : inputTrajectories) {
+      for (auto tip : trajectories.tips()) {
+        if (!trajectories.hasTrackParameters(tip)) {
+          continue;
+        }
+        const auto& trackParam = trajectories.trackParameters(tip);
+        inputTrackParameters.push_back(trackParam);
+      }
+    }
+  }
+
+  std::string path =
+      perEventFilepath(m_cfg.outputDir, m_cfg.outputStem, ctx.eventNumber);
+
+  dfe::NamedTupleCsvWriter<TrackParameterData> writer(path,
+                                                      m_cfg.outputPrecision);
+
+  TrackParameterData data{};
+  for (const auto& tp : inputTrackParameters) {
+    const auto& params = tp.parameters();
+    const auto& cov = tp.covariance().value();
+
+    data.d0 = params[Acts::eBoundLoc0];
+    data.z0 = params[Acts::eBoundLoc1];
+    data.phi = params[Acts::eBoundPhi];
+    data.theta = params[Acts::eBoundTheta];
+    data.qop = params[Acts::eBoundQOverP];
+
+    data.var_d0 = cov(Acts::eBoundLoc0, Acts::eBoundLoc0);
+    data.var_z0 = cov(Acts::eBoundLoc1, Acts::eBoundLoc1);
+    data.var_phi = cov(Acts::eBoundPhi, Acts::eBoundPhi);
+    data.var_theta = cov(Acts::eBoundTheta, Acts::eBoundTheta);
+    data.var_qop = cov(Acts::eBoundQOverP, Acts::eBoundQOverP);
+
+    data.cov_d0z0 = cov(Acts::eBoundLoc0, Acts::eBoundLoc1);
+    data.cov_d0phi = cov(Acts::eBoundLoc0, Acts::eBoundPhi);
+    data.cov_d0theta = cov(Acts::eBoundLoc0, Acts::eBoundTheta);
+    data.cov_d0qop = cov(Acts::eBoundLoc0, Acts::eBoundQOverP);
+
+    data.cov_z0d0 = cov(Acts::eBoundLoc1, Acts::eBoundLoc0);
+    data.cov_z0phi = cov(Acts::eBoundLoc1, Acts::eBoundPhi);
+    data.cov_z0theta = cov(Acts::eBoundLoc1, Acts::eBoundTheta);
+    data.cov_z0qop = cov(Acts::eBoundLoc1, Acts::eBoundQOverP);
+
+    data.cov_phid0 = cov(Acts::eBoundPhi, Acts::eBoundLoc0);
+    data.cov_phiz0 = cov(Acts::eBoundPhi, Acts::eBoundLoc1);
+    data.cov_phitheta = cov(Acts::eBoundPhi, Acts::eBoundTheta);
+    data.cov_phiqop = cov(Acts::eBoundPhi, Acts::eBoundQOverP);
+
+    data.cov_thetad0 = cov(Acts::eBoundTheta, Acts::eBoundLoc0);
+    data.cov_thetaz0 = cov(Acts::eBoundTheta, Acts::eBoundLoc1);
+    data.cov_thetaphi = cov(Acts::eBoundTheta, Acts::eBoundPhi);
+    data.cov_thetaqop = cov(Acts::eBoundTheta, Acts::eBoundQOverP);
+
+    data.cov_qopd0 = cov(Acts::eBoundQOverP, Acts::eBoundLoc0);
+    data.cov_qopz0 = cov(Acts::eBoundQOverP, Acts::eBoundLoc1);
+    data.cov_qopphi = cov(Acts::eBoundQOverP, Acts::eBoundPhi);
+    data.cov_qoptheta = cov(Acts::eBoundQOverP, Acts::eBoundTheta);
+
+    writer.append(data);
+  }
+
+  return ActsExamples::ProcessCode::SUCCESS;
+}

Examples/Python/src/Output.cpp

@@ -14,6 +14,7 @@
 #include "ActsExamples/Io/Csv/CsvPlanarClusterWriter.hpp"
 #include "ActsExamples/Io/Csv/CsvSimHitWriter.hpp"
 #include "ActsExamples/Io/Csv/CsvSpacepointWriter.hpp"
+#include "ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp"
 #include "ActsExamples/Io/Csv/CsvTrackingGeometryWriter.hpp"
 #include "ActsExamples/Io/NuclearInteractions/RootNuclearInteractionParametersWriter.hpp"
 #include "ActsExamples/Io/Performance/CKFPerformanceWriter.hpp"
@@ -343,6 +344,11 @@ void addOutput(Context& ctx) {
       invariantMassBins, multiplicityMax, writeOptionalHistograms,
       nSimulatedEvents);
 
+  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvTrackParameterWriter, mex,
+                             "CsvTrackParameterWriter", inputTrackParameters,
+                             inputTrajectories, outputDir, outputStem,
+                             outputPrecision);
+
   {
     using Writer = ActsExamples::CsvBFieldWriter;