Add Track tree

acts-project · Apr 13, 2024 · 4418cd1 · 4418cd1
1 parent b51d395
commit 4418cd1
Show file tree

Hide file tree

Showing 9 changed files with 282 additions and 2 deletions.
diff --git a/examples/run/cpu/seeding_example.cpp b/examples/run/cpu/seeding_example.cpp
@@ -25,6 +25,7 @@
 // performance
 #include "traccc/efficiency/finding_performance_writer.hpp"
 #include "traccc/efficiency/seeding_performance_writer.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/resolution/fitting_performance_writer.hpp"
 
 // options
@@ -79,6 +80,8 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
     vecmem::host_memory_resource host_mr;
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::seeding_performance_writer sd_performance_writer(
         traccc::seeding_performance_writer::config{});
     traccc::finding_performance_writer find_performance_writer(
@@ -229,6 +232,9 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
             traccc::event_map2 evt_map(event, input_opts.directory,
                                        input_opts.directory,
                                        input_opts.directory);
+
+            evt_tree_writer.write(traccc::get_data(track_states), evt_map);
+
             sd_performance_writer.write(vecmem::get_data(seeds),
                                         vecmem::get_data(spacepoints_per_event),
                                         evt_map);
@@ -253,6 +259,7 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         sd_performance_writer.finalize();
         find_performance_writer.finalize();
         fit_performance_writer.finalize();

diff --git a/examples/run/cpu/truth_finding_example.cpp b/examples/run/cpu/truth_finding_example.cpp
@@ -9,6 +9,7 @@
 #include "traccc/definitions/common.hpp"
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/efficiency/finding_performance_writer.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/finding/finding_algorithm.hpp"
 #include "traccc/fitting/fitting_algorithm.hpp"
 #include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
@@ -67,6 +68,8 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
     vecmem::host_memory_resource host_mr;
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::finding_performance_writer find_performance_writer(
         traccc::finding_performance_writer::config{});
     traccc::fitting_performance_writer fit_performance_writer(
@@ -173,6 +176,9 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
         const unsigned int n_fitted_tracks = track_states.size();
 
         if (performance_opts.run) {
+
+            evt_tree_writer.write(traccc::get_data(track_states), evt_map2);
+
             find_performance_writer.write(traccc::get_data(track_candidates),
                                           evt_map2);
 
@@ -188,6 +194,7 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         find_performance_writer.finalize();
         fit_performance_writer.finalize();
     }

diff --git a/examples/run/cpu/truth_fitting_example.cpp b/examples/run/cpu/truth_fitting_example.cpp
@@ -8,6 +8,7 @@
 // Project include(s).
 #include "traccc/definitions/common.hpp"
 #include "traccc/definitions/primitives.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/fitting/fitting_algorithm.hpp"
 #include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
 #include "traccc/io/read_geometry.hpp"
@@ -74,6 +75,8 @@ int main(int argc, char* argv[]) {
     vecmem::host_memory_resource host_mr;
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::fitting_performance_writer fit_performance_writer(
         traccc::fitting_performance_writer::config{});
 
@@ -145,6 +148,8 @@ int main(int argc, char* argv[]) {
 
         if (performance_opts.run) {
 
+            evt_tree_writer.write(traccc::get_data(track_states), evt_map2);
+
             for (unsigned int i = 0; i < n_fitted_tracks; i++) {
                 const auto& trk_states_per_track = track_states.at(i).items;
 
@@ -157,6 +162,7 @@ int main(int argc, char* argv[]) {
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         fit_performance_writer.finalize();
     }
 

diff --git a/examples/run/cuda/seeding_example_cuda.cpp b/examples/run/cuda/seeding_example_cuda.cpp
@@ -17,6 +17,7 @@
 #include "traccc/efficiency/nseed_performance_writer.hpp"
 #include "traccc/efficiency/seeding_performance_writer.hpp"
 #include "traccc/efficiency/track_filter.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/finding/finding_algorithm.hpp"
 #include "traccc/fitting/fitting_algorithm.hpp"
 #include "traccc/io/read_geometry.hpp"
@@ -95,6 +96,8 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
     traccc::memory_resource mr{device_mr, &cuda_host_mr};
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::seeding_performance_writer sd_performance_writer(
         traccc::seeding_performance_writer::config{});
     traccc::finding_performance_writer find_performance_writer(
@@ -436,6 +439,9 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
             traccc::event_map2 evt_map(event, input_opts.directory,
                                        input_opts.directory,
                                        input_opts.directory);
+
+            evt_tree_writer.write(traccc::get_data(track_states), evt_map);
+
             sd_performance_writer.write(
                 vecmem::get_data(seeds_cuda),
                 vecmem::get_data(sp_reader_output.spacepoints), evt_map);
@@ -456,6 +462,7 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         sd_performance_writer.finalize();
         nsd_performance_writer.finalize();
         find_performance_writer.finalize();

diff --git a/examples/run/cuda/truth_finding_example_cuda.cpp b/examples/run/cuda/truth_finding_example_cuda.cpp
@@ -14,6 +14,7 @@
 #include "traccc/device/container_d2h_copy_alg.hpp"
 #include "traccc/device/container_h2d_copy_alg.hpp"
 #include "traccc/efficiency/finding_performance_writer.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/finding/finding_algorithm.hpp"
 #include "traccc/fitting/fitting_algorithm.hpp"
 #include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
@@ -89,6 +90,8 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
     traccc::memory_resource mr{device_mr, &cuda_host_mr};
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::finding_performance_writer find_performance_writer(
         traccc::finding_performance_writer::config{});
     traccc::fitting_performance_writer fit_performance_writer(
@@ -319,6 +322,7 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
         n_fitted_tracks_cuda += track_states_cuda.size();
 
         if (performance_opts.run) {
+            evt_tree_writer.write(traccc::get_data(track_states), evt_map2);
             find_performance_writer.write(
                 traccc::get_data(track_candidates_cuda), evt_map2);
 
@@ -335,6 +339,7 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         find_performance_writer.finalize();
         fit_performance_writer.finalize();
     }

diff --git a/examples/run/cuda/truth_fitting_example_cuda.cpp b/examples/run/cuda/truth_fitting_example_cuda.cpp
@@ -12,6 +12,7 @@
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/device/container_d2h_copy_alg.hpp"
 #include "traccc/device/container_h2d_copy_alg.hpp"
+#include "traccc/event/track_tree_writer.hpp"
 #include "traccc/fitting/fitting_algorithm.hpp"
 #include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
 #include "traccc/io/read_geometry.hpp"
@@ -96,6 +97,8 @@ int main(int argc, char* argv[]) {
     traccc::memory_resource mr{device_mr, &cuda_host_mr};
 
     // Performance writer
+    traccc::track_tree_writer evt_tree_writer(
+        traccc::track_tree_writer::config{});
     traccc::fitting_performance_writer fit_performance_writer(
         traccc::fitting_performance_writer::config{});
 
@@ -237,6 +240,10 @@ int main(int argc, char* argv[]) {
         n_fitted_tracks_cuda += track_states_cuda.size();
 
         if (performance_opts.run) {
+
+            evt_tree_writer.write(traccc::get_data(track_states_cuda),
+                                  evt_map2);
+
             for (unsigned int i = 0; i < track_states_cuda.size(); i++) {
                 const auto& trk_states_per_track =
                     track_states_cuda.at(i).items;
@@ -250,6 +257,7 @@ int main(int argc, char* argv[]) {
     }
 
     if (performance_opts.run) {
+        evt_tree_writer.finalize();
         fit_performance_writer.finalize();
     }
 

diff --git a/performance/CMakeLists.txt b/performance/CMakeLists.txt
@@ -9,6 +9,9 @@ include( traccc-compiler-options-cpp )
 
 # Set up the build of the traccc::performance library.
 traccc_add_library( traccc_performance performance TYPE SHARED
+   # Event plot code
+   "src/event/track_tree_writer.cpp"
+   "include/traccc/event/track_tree_writer.hpp"
    # Efficiency calculation code.
    "src/efficiency/duplication_plot_tool.hpp"
    "src/efficiency/eff_plot_tool.hpp"
@@ -68,9 +71,9 @@ target_link_libraries( traccc_performance
 
 # Use ROOT in traccc::performance, if requested.
 if( TRACCC_USE_ROOT )
-   find_package( ROOT COMPONENTS Core RIO Hist REQUIRED )
+   find_package( ROOT COMPONENTS Core RIO Hist Tree REQUIRED )
    target_link_libraries( traccc_performance
-      PRIVATE ROOT::Core ROOT::RIO ROOT::Hist )
+      PRIVATE ROOT::Core ROOT::RIO ROOT::Hist ROOT::Tree )
    target_compile_definitions( traccc_performance
       PRIVATE TRACCC_HAVE_ROOT )
 endif()

diff --git a/performance/include/traccc/event/track_tree_writer.hpp b/performance/include/traccc/event/track_tree_writer.hpp
@@ -0,0 +1,60 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2023 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Local include(s).
+#include "traccc/utils/helpers.hpp"
+
+// Project include(s).
+#include "traccc/edm/track_state.hpp"
+#include "traccc/io/event_map2.hpp"
+
+namespace traccc {
+namespace details {
+
+/// Data members that should not pollute the API of
+/// @c traccc::track_tree_writer
+struct track_tree_writer_data;
+
+}  // namespace details
+
+class track_tree_writer {
+
+    public:
+    /// Configuration for the tool
+    struct config {
+
+        /// Output filename.
+        std::string file_path = "track_tree.root";
+        /// Output file mode
+        std::string file_mode = "RECREATE";
+    };
+
+    /// Construct from configuration and log level.
+    /// @param cfg The configuration
+    ///
+    track_tree_writer(const config& cfg);
+
+    /// Destructor
+    ~track_tree_writer();
+
+    void write(const track_state_container_types::const_view& track_states_view,
+               const event_map2& evt_map);
+
+    void finalize();
+
+    private:
+    /// Configuration for the tool
+    config m_cfg;
+
+    /// Opaque data members for the class
+    std::unique_ptr<details::track_tree_writer_data> m_data;
+
+};  // class track_tree_writer
+
+}  // namespace traccc