Add ROOT-based event exporter

app/celer-g4/celer-g4.cc

@@ -18,6 +18,10 @@
 #include <G4UImanager.hh>
 #include <G4Version.hh>
 
+#include "accel/HepMC3PrimaryGenerator.hh"
+
+#include "GlobalSetup.hh"
+
 #if G4VERSION_NUMBER >= 1100
 #    include <G4RunManagerFactory.hh>
 #else
@@ -80,7 +84,7 @@ void run(int argc, char** argv)
     CELER_LOG(info) << "Run manager type: "
                     << TypeDemangler<G4RunManager>{}(*run_manager);
 
-    // Make global setup commands available to UI
+    // Construct singleton, also making it available to UI
     GlobalSetup::Instance();
 
     G4UImanager* ui = G4UImanager::GetUIpointer();

src/accel/CMakeLists.txt

@@ -35,6 +35,12 @@ endif()
 if(CELERITAS_USE_HepMC3)
   list(APPEND PRIVATE_DEPS HepMC3::HepMC3)
   list(APPEND SOURCES HepMC3PrimaryGenerator.cc)
+
+  if (CELERITAS_USE_ROOT)
+    list(APPEND PRIVATE_DEPS ROOT::Tree)
+    list(APPEND SOURCES detail/RootOffloadWriter.cc)
+    list(APPEND SOURCES detail/RootOffloadReader.cc)
+  endif()
 endif()
 
 #-----------------------------------------------------------------------------#

src/accel/LocalTransporter.cc

@@ -27,6 +27,7 @@
 #include "SharedParams.hh"
 #include "detail/HitManager.hh"
 #include "detail/OffloadWriter.hh"
+#include "detail/RootOffloadWriter.hh"
 
 namespace celeritas
 {
@@ -39,6 +40,7 @@ LocalTransporter::LocalTransporter(SetupOptions const& options,
     : auto_flush_(options.max_num_tracks)
     , max_steps_(options.max_steps)
     , dump_primaries_{params.offload_writer()}
+    , dump_primaries_root_{params.root_offload_writer()}
     , hit_manager_{params.hit_manager()}
 {
     CELER_VALIDATE(params,
@@ -151,9 +153,14 @@ void LocalTransporter::Flush()
 
     if (dump_primaries_)
     {
-        // Write offload particles if user requested
+        // Write offload particles to HepMC3 if user requested
         (*dump_primaries_)(buffer_);
     }
+    if (dump_primaries_root_)
+    {
+        // Write offload particles to ROOT if user requested
+        (*dump_primaries_root_)(buffer_);
+    }
 
     // Abort cleanly for interrupt and user-defined signals
     ScopedSignalHandler interrupted{SIGINT, SIGUSR2};

src/accel/LocalTransporter.hh

@@ -26,6 +26,7 @@ namespace detail
 {
 class HitManager;
 class OffloadWriter;
+class RootOffloadWriter;
 }  // namespace detail
 
 struct SetupOptions;
@@ -76,6 +77,7 @@ class LocalTransporter
   private:
     using SPHitManger = std::shared_ptr<detail::HitManager>;
     using SPOffloadWriter = std::shared_ptr<detail::OffloadWriter>;
+    using SPRootOffloadWriter = std::shared_ptr<detail::RootOffloadWriter>;
 
     struct HMFinalizer
     {
@@ -95,6 +97,7 @@ class LocalTransporter
 
     // Shared across threads to write flushed particles
     SPOffloadWriter dump_primaries_;
+    SPRootOffloadWriter dump_primaries_root_;
     // Shared pointer across threads, "finalize" called when clearing
     InitializedValue<SPHitManger, HMFinalizer> hit_manager_;
 };

src/accel/SetupOptions.hh

@@ -112,6 +112,8 @@ struct SetupOptions
     std::string physics_output_file;
     //! Filename to dump a HepMC3 copy of offloaded tracks as events
     std::string offload_output_file;
+    //! Filename to dump a ROOT copy of offloaded tracks as events
+    std::string offload_root_output_file;
     //!@}
 
     //!@{

src/accel/SetupOptionsMessenger.cc

@@ -172,6 +172,9 @@ SetupOptionsMessenger::SetupOptionsMessenger(SetupOptions* options)
     add_cmd(&options->offload_output_file,
             "offloadOutputFile",
             "Filename for HepMC3 copy of offloaded tracks as events");
+    add_cmd(&options->offload_root_output_file,
+            "offloadRootOutputFile",
+            "Filename for ROOT copy of offloaded tracks as events");
     add_cmd(&options->max_num_tracks,
             "maxNumTracks",
             "Number of track \"slots\" to be transported simultaneously");

src/accel/SharedParams.cc

@@ -51,6 +51,7 @@
 #include "SetupOptions.hh"
 #include "detail/HitManager.hh"
 #include "detail/OffloadWriter.hh"
+#include "detail/RootOffloadWriter.hh"
 
 namespace celeritas
 {
@@ -173,6 +174,12 @@ SharedParams::SharedParams(SetupOptions const& options)
             options.offload_output_file, params_->particle());
     }
 
+    if (!options.offload_root_output_file.empty())
+    {
+        root_offload_writer_ = std::make_shared<detail::RootOffloadWriter>(
+            options.offload_root_output_file, params_->particle());
+    }
+
     CELER_ENSURE(*this);
 }
 

src/accel/SharedParams.hh

@@ -21,6 +21,7 @@ namespace detail
 {
 class HitManager;
 class OffloadWriter;
+class RootOffloadWriter;
 }  // namespace detail
 class CoreParams;
 struct Primary;
@@ -87,13 +88,20 @@ class SharedParams
 
     using SPHitManager = std::shared_ptr<detail::HitManager>;
     using SPOffloadWriter = std::shared_ptr<detail::OffloadWriter>;
+    using SPRootOffloadWriter = std::shared_ptr<detail::RootOffloadWriter>;
 
     //! Hit manager, to be used only by LocalTransporter
     SPHitManager const& hit_manager() const { return hit_manager_; }
 
     //! Optional offload writer, only for use by LocalTransporter
     SPOffloadWriter const& offload_writer() const { return offload_writer_; }
 
+    //! Optional ROOT offload writer, only for use by LocalTransporter
+    SPRootOffloadWriter const& root_offload_writer() const
+    {
+        return root_offload_writer_;
+    }
+
     //!@}
 
   private:
@@ -105,6 +113,7 @@ class SharedParams
     VecG4ParticleDef particles_;
     std::string output_filename_;
     SPOffloadWriter offload_writer_;
+    SPRootOffloadWriter root_offload_writer_;
 
     //// HELPER FUNCTIONS ////
 

src/accel/detail/RootOffloadReader.cc

@@ -0,0 +1,105 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2023 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file accel/detail/RootOffloadReader.cc
+//---------------------------------------------------------------------------//
+#include "RootOffloadReader.hh"
+
+#include <TFile.h>
+#include <TLeaf.h>
+#include <TTree.h>
+
+#include "corecel/io/Logger.hh"
+#include "celeritas/phys/ParticleParams.hh"
+
+namespace celeritas
+{
+namespace detail
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Construct with ROOT input filename.
+ */
+RootOffloadReader::RootOffloadReader(std::string const& filename,
+                                     SPConstParticles params)
+    : params_(std::move(params))
+{
+    CELER_EXPECT(!filename.empty());
+    tfile_.reset(TFile::Open(filename.c_str(), "read"));
+    CELER_ASSERT(tfile_->IsOpen());
+    ttree_.reset(tfile_->Get<TTree>(tree_name()));
+    CELER_ASSERT(ttree_);
+
+    num_entries_ = ttree_->GetEntries();
+    CELER_ASSERT(num_entries_ > 0);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Read single event from the primaries tree.
+ */
+auto RootOffloadReader::operator()() -> result_type
+{
+    std::lock_guard scoped_lock{read_mutex_};
+
+    CELER_EXPECT(entry_count_ <= num_entries_);
+    ttree_->GetEntry(entry_count_);
+    auto const this_evt_id = ttree_->GetLeaf("event_id")->GetValue();
+
+    result_type primaries;
+    for (; entry_count_ < num_entries_; entry_count_++)
+    {
+        ttree_->GetEntry(entry_count_);
+        if (ttree_->GetLeaf("event_id")->GetValue() != this_evt_id)
+        {
+            break;
+        }
+
+        Primary primary;
+        primary.event_id = EventId{this->from_leaf<std::size_t>("event_id")};
+        primary.track_id = TrackId{this->from_leaf<std::size_t>("track_id")};
+        primary.particle_id
+            = params_->find(PDGNumber{this->from_leaf<int>("particle")});
+        primary.energy = units::MevEnergy{this->from_leaf<double>("energy")};
+        primary.time = this->from_leaf<double>("time");
+        primary.position = this->from_array_leaf("pos");
+        primary.direction = this->from_array_leaf("dir");
+        primaries.push_back(std::move(primary));
+    }
+
+    CELER_LOG_LOCAL(info) << "Read event " << this_evt_id << " with "
+                          << primaries.size() << " primaries";
+    return primaries;
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Helper function to fetch leaves.
+ */
+template<class T>
+auto RootOffloadReader::from_leaf(char const* leaf_name) -> T
+{
+    CELER_EXPECT(ttree_);
+    auto const leaf = ttree_->GetLeaf(leaf_name);
+    CELER_ASSERT(leaf);
+    return static_cast<T>(leaf->GetValue());
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Helper function to fetch leaves containing `std::array<double, 3>`.
+ */
+Real3 RootOffloadReader::from_array_leaf(char const* leaf_name)
+{
+    CELER_EXPECT(ttree_);
+    auto const leaf = ttree_->GetLeaf(leaf_name);
+    CELER_ASSERT(leaf);
+    CELER_ASSERT(leaf->GetLen() == 3);
+    return {leaf->GetValue(0), leaf->GetValue(1), leaf->GetValue(2)};
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas

src/accel/detail/RootOffloadReader.hh

@@ -0,0 +1,111 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2023 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file accel/detail/RootOffloadReader.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include <string>
+
+#include "corecel/Macros.hh"
+#include "celeritas/ext/RootUniquePtr.hh"
+#include "celeritas/phys/Primary.hh"
+
+// ROOT forward declaration
+class TFile;
+class TTree;
+
+namespace celeritas
+{
+class ParticleParams;
+
+namespace detail
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Read ROOT file generated by \c RootOffloadWriter .
+ *
+ * Use \c operator() to read new primaries:
+ * \code
+    RootOffloadReader read("primaries.root", particle_params);
+    auto event = read();
+    while (!event.empty())
+    {
+        event.read();
+    }
+ * \endcode
+ */
+class RootOffloadReader
+{
+  public:
+    //!@{
+    //! \name Type aliases
+    using SPConstParticles = std::shared_ptr<ParticleParams const>;
+    using result_type = std::vector<Primary>;
+    //!@}
+
+    // Construct with ROOT filename
+    explicit RootOffloadReader(std::string const& filename,
+                               SPConstParticles params);
+
+    // Read a single event from the ROOT file
+    result_type operator()();
+
+  private:
+    //// DATA ////
+
+    SPConstParticles params_;
+    std::size_t entry_count_{0};  // Current TTree entry
+    std::size_t num_entries_;  // Total number of entries in the TTree
+    UPExtern<TFile> tfile_;
+    UPExtern<TTree> ttree_;
+    std::mutex read_mutex_;
+
+    //// HELPER FUNCTIONS ////
+
+    // Hardcoded ROOT TTree name defined by RootOffloadWriter
+    char const* tree_name() { return "primaries"; }
+
+    // Fetch basic data types from leaves
+    template<class T>
+    auto from_leaf(char const* leaf_name) -> T;
+
+    // Fetch arrays from leaves
+    Real3 from_array_leaf(char const* leaf_name);
+};
+
+//---------------------------------------------------------------------------//
+#if !CELERITAS_USE_ROOT
+inline RootOffloadReader::RootOffloadReader(std::string const&)
+{
+    (void)sizeof(num_entries_);
+    (void)sizeof(entry_count_);
+    (void)sizeof(tfile_);
+    (void)sizeof(ttree_);
+    (void)sizeof(read_mutex_);
+    (void)sizeof(event_);
+    CELER_NOT_CONFIGURED("ROOT");
+}
+
+inline RootOffloadReader::result_type RootOffloadReader::operator()()
+{
+    CELER_ASSERT_UNREACHABLE();
+}
+
+inline template<class T>
+auto RootOffloadReader::from_leaf(char const*) -> T
+{
+    CELER_ASSERT_UNREACHABLE();
+}
+
+Real3 RootOffloadReader::from_leaf(char const*)
+{
+    CELER_ASSERT_UNREACHABLE();
+}
+#endif
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas