Merge pull request #588 from krasznaa/MoreObjWriting-main-20240513

More (Wavefront) Obj Writing, main branch (2024.05.13.)

examples/run/cpu/seq_example.cpp

@@ -245,6 +245,11 @@ int seq_run(const traccc::opts::input_data& input_opts,
                 traccc::performance::timer timer{"Seeding", elapsedTimes};
                 seeds = sa(spacepoints_per_event);
             }
+            if (output_opts.directory != "") {
+                traccc::io::write(event, output_opts.directory,
+                                  output_opts.format, vecmem::get_data(seeds),
+                                  vecmem::get_data(spacepoints_per_event));
+            }
 
             /*----------------------------
               Track params estimation
@@ -265,6 +270,11 @@ int seq_run(const traccc::opts::input_data& input_opts,
                     track_candidates = finding_alg(
                         detector, field, measurements_per_event, params);
                 }
+                if (output_opts.directory != "") {
+                    traccc::io::write(
+                        event, output_opts.directory, output_opts.format,
+                        traccc::get_data(track_candidates), detector);
+                }
                 {
                     traccc::performance::timer timer{"Track fitting",
                                                      elapsedTimes};

io/CMakeLists.txt

@@ -73,8 +73,12 @@ traccc_add_library( traccc_io io TYPE SHARED
   "src/csv/make_particle_reader.cpp"
   "src/csv/read_particles.hpp"
   "src/csv/read_particles.cpp"
+  "src/obj/write_seeds.hpp"
+  "src/obj/write_seeds.cpp"
   "src/obj/write_spacepoints.hpp"
   "src/obj/write_spacepoints.cpp"
+  "src/obj/write_track_candidates.hpp"
+  "src/obj/write_track_candidates.cpp"
   )
 target_link_libraries( traccc_io
   PUBLIC vecmem::core traccc::core ActsCore

io/include/traccc/io/write.hpp

@@ -1,6 +1,6 @@
 /** TRACCC library, part of the ACTS project (R&D line)
  *
- * (c) 2021-2022 CERN for the benefit of the ACTS project
+ * (c) 2021-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
@@ -9,9 +9,14 @@
 
 // Project include(s).
 #include "traccc/edm/cell.hpp"
+#include "traccc/edm/seed.hpp"
 #include "traccc/edm/spacepoint.hpp"
+#include "traccc/edm/track_candidate.hpp"
 #include "traccc/io/data_format.hpp"
 
+// Detray include(s).
+#include "detray/core/detector.hpp"
+
 // System include(s).
 #include <cstddef>
 #include <string_view>
@@ -57,4 +62,29 @@ void write(std::size_t event, std::string_view directory,
            measurement_collection_types::const_view measurements,
            traccc::cell_module_collection_types::const_view modules);
 
+/// Function for seed writing
+///
+/// @param event is the event index
+/// @param directory is the directory for the output seed file
+/// @param format is the data format (obj only right now) of output file
+/// @param seeds is the seed collection to write
+/// @param spacepoints is the spacepoint collection the seeds are made of
+///
+void write(std::size_t event, std::string_view directory,
+           traccc::data_format format, seed_collection_types::const_view seeds,
+           spacepoint_collection_types::const_view spacepoints);
+
+/// Function for track candidate writing
+///
+/// @param event is the event index
+/// @param directory is the directory for the output seed file
+/// @param format is the data format (obj only right now) of output file
+/// @param tracks is the track candidate container to write
+/// @param detector is the Detray detector describing the geometry
+///
+void write(std::size_t event, std::string_view directory,
+           traccc::data_format format,
+           track_candidate_container_types::const_view tracks,
+           const detray::detector<>& detector);
+
 }  // namespace traccc::io

io/src/obj/write_seeds.cpp

@@ -0,0 +1,88 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+// Local include(s).
+#include "write_seeds.hpp"
+
+// System include(s).
+#include <fstream>
+
+namespace traccc::io::obj {
+
+void write_seeds(std::string_view filename,
+                 seed_collection_types::const_view seeds_view,
+                 spacepoint_collection_types::const_view spacepoints_view) {
+
+    // Open the output file.
+    std::ofstream file{filename.data()};
+    if (!file.is_open()) {
+        throw std::runtime_error("Failed to open file: " +
+                                 std::string(filename));
+    }
+
+    // Create device collections around the views.
+    const seed_collection_types::const_device seeds{seeds_view};
+    const spacepoint_collection_types::const_device spacepoints{
+        spacepoints_view};
+
+    // Map associating in-memory spacepoint indices to in-file ones.
+    std::map<std::size_t, std::size_t> spacepoint_indices;
+
+    // Helper lambda to write a spacepoint to the output file.
+    auto write_spacepoint = [&file, &spacepoints, &spacepoint_indices](
+                                std::size_t memory_index,
+                                std::size_t file_index) -> bool {
+        // Check whether this spacepoint has already been written.
+        if (spacepoint_indices.find(memory_index) != spacepoint_indices.end()) {
+            return false;
+        }
+        // Write the spacepoint.
+        const traccc::spacepoint& sp = spacepoints[memory_index];
+        file << "v " << sp.x() << " " << sp.y() << " " << sp.z() << "\n";
+        // Remember the mapping.
+        spacepoint_indices[memory_index] = file_index;
+        return true;
+    };
+
+    // First, write out all of the spacepoints as vertices. Making sure that we
+    // only write each spacepoint once. And remembering the indices of the
+    // spacepoints in the output file, to be used when making the seeds.
+    std::size_t file_index = 1;
+    file << "# Spacepoints from which the seeds are built\n";
+    for (const seed& s : seeds) {
+        if (write_spacepoint(s.spB_link, file_index)) {
+            file_index++;
+        }
+        if (write_spacepoint(s.spM_link, file_index)) {
+            file_index++;
+        }
+        if (write_spacepoint(s.spT_link, file_index)) {
+            file_index++;
+        }
+    }
+
+    // Helper lambda for getting an element of the spacepoint_indices map, in a
+    // "safe" way.
+    auto get_spacepoint_index =
+        [&spacepoint_indices](std::size_t memory_index) -> std::size_t {
+        auto it = spacepoint_indices.find(memory_index);
+        if (it == spacepoint_indices.end()) {
+            throw std::runtime_error("Spacepoint index not found");
+        }
+        return it->second;
+    };
+
+    // Now build the seeds as lines connecting the spacepoint vertices.
+    file << "# Seeds\n";
+    for (const seed& s : seeds) {
+        file << "l " << get_spacepoint_index(s.spB_link) << " "
+             << get_spacepoint_index(s.spM_link) << " "
+             << get_spacepoint_index(s.spT_link) << "\n";
+    }
+}
+
+}  // namespace traccc::io::obj

io/src/obj/write_seeds.hpp

@@ -0,0 +1,29 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/edm/seed.hpp"
+#include "traccc/edm/spacepoint.hpp"
+
+// System include(s).
+#include <string_view>
+
+namespace traccc::io::obj {
+
+/// Write a seed collection into a Wavefront OBJ file.
+///
+/// @param filename is the name of the output file
+/// @param seeds is the seed collection to write
+/// @param spacepoints is the spacepoint collection that the seeds reference
+///
+void write_seeds(std::string_view filename,
+                 seed_collection_types::const_view seeds,
+                 spacepoint_collection_types::const_view spacepoints);
+
+}  // namespace traccc::io::obj

io/src/obj/write_track_candidates.cpp

@@ -0,0 +1,81 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+// Local include(s).
+#include "write_track_candidates.hpp"
+
+// System include(s).
+#include <cassert>
+#include <fstream>
+
+namespace traccc::io::obj {
+
+void write_track_candidates(
+    std::string_view filename,
+    track_candidate_container_types::const_view tracks_view,
+    const detray::detector<>& detector) {
+
+    // Open the output file.
+    std::ofstream file{filename.data()};
+    if (!file.is_open()) {
+        throw std::runtime_error("Failed to open file: " +
+                                 std::string(filename));
+    }
+
+    // Create a device collection around the track container view.
+    const track_candidate_container_types::const_device tracks{tracks_view};
+
+    // Convenience type.
+    using size_type = track_candidate_container_types::const_device::size_type;
+
+    // First write out the measurements / spacepoints that the tracks are
+    // made from. Don't try to resolve the overlaps, just write out duplicate
+    // measurements if needed.
+    file << "# Measurements / spacepoints that the tracks are made out of\n";
+    for (size_type i = 0; i < tracks.size(); ++i) {
+
+        // The track candidate in question.
+        const track_candidate_container_types::const_device::const_element_view
+            track = tracks.at(i);
+
+        // Loop over the measurements that the track candidate is made out of.
+        for (const measurement& m : track.items) {
+
+            // Find the detector surface that this measurement sits on.
+            const detray::surface<detray::detector<> > surface{detector,
+                                                               m.surface_link};
+
+            // Calculate a position for this measurement in global 3D space.
+            const auto global = surface.bound_to_global({}, m.local, {});
+
+            // Write the 3D coordinates of the measurement / spacepoint.
+            assert(global.size() == 3);
+            file << "v " << global[0] << " " << global[1] << " " << global[2]
+                 << "\n";
+        }
+    }
+
+    // Now loop over the track candidates again, and creates lines for each of
+    // them using the measurements / spacepoints written out earlier.
+    file << "# Track candidates\n";
+    std::size_t vertex_counter = 1;
+    for (size_type i = 0; i < tracks.size(); ++i) {
+
+        // The track candidate in question.
+        const track_candidate_container_types::const_device::const_element_view
+            track = tracks.at(i);
+
+        // Construct the lines.
+        file << "l";
+        for (size_type j = 0; j < track.items.size(); ++j) {
+            file << " " << vertex_counter++;
+        }
+        file << "\n";
+    }
+}
+
+}  // namespace traccc::io::obj

io/src/obj/write_track_candidates.hpp

@@ -0,0 +1,31 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/edm/track_candidate.hpp"
+
+// Detray include(s).
+#include "detray/core/detector.hpp"
+
+// System include(s).
+#include <string_view>
+
+namespace traccc::io::obj {
+
+/// Write a track candidate container into a Wavefront OBJ file.
+///
+/// @param filename is the name of the output file
+/// @param tracks is the track candidate container to write
+/// @param detector is the Detray detector describing the geometry
+///
+void write_track_candidates(std::string_view filename,
+                            track_candidate_container_types::const_view tracks,
+                            const detray::detector<>& detector);
+
+}  // namespace traccc::io::obj

io/src/write.cpp

@@ -8,7 +8,9 @@
 // Local include(s).
 #include "traccc/io/write.hpp"
 
+#include "obj/write_seeds.hpp"
 #include "obj/write_spacepoints.hpp"
+#include "obj/write_track_candidates.hpp"
 #include "traccc/io/utils.hpp"
 #include "write_binary.hpp"
 
@@ -102,4 +104,41 @@ void write(std::size_t event, std::string_view directory,
     }
 }
 
+void write(std::size_t event, std::string_view directory,
+           traccc::data_format format, seed_collection_types::const_view seeds,
+           spacepoint_collection_types::const_view spacepoints) {
+
+    switch (format) {
+        case data_format::obj:
+            obj::write_seeds(
+                get_absolute_path((std::filesystem::path(directory) /
+                                   std::filesystem::path(
+                                       get_event_filename(event, "-seeds.obj")))
+                                      .native()),
+                seeds, spacepoints);
+            break;
+        default:
+            throw std::invalid_argument("Unsupported data format");
+    }
+}
+
+void write(std::size_t event, std::string_view directory,
+           traccc::data_format format,
+           track_candidate_container_types::const_view tracks,
+           const detray::detector<>& detector) {
+
+    switch (format) {
+        case data_format::obj:
+            obj::write_track_candidates(
+                get_absolute_path((std::filesystem::path(directory) /
+                                   std::filesystem::path(get_event_filename(
+                                       event, "-track-candidates.obj")))
+                                      .native()),
+                tracks, detector);
+            break;
+        default:
+            throw std::invalid_argument("Unsupported data format");
+    }
+}
+
 }  // namespace traccc::io