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
@@ -81,7 +85,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();
 
     // Read user input

src/accel/LocalTransporter.hh

@@ -26,6 +26,7 @@ namespace detail
 {
 class HitManager;
 class OffloadWriter;
+class RootEventWriter;
 }  // 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::RootEventWriter>;
 
     struct HMFinalizer
     {

src/accel/SetupOptionsMessenger.cc

@@ -171,7 +171,7 @@ SetupOptionsMessenger::SetupOptionsMessenger(SetupOptions* options)
             "Filename for ROOT dump of physics data");
     add_cmd(&options->offload_output_file,
             "offloadOutputFile",
-            "Filename for HepMC3 copy of offloaded tracks as events");
+            "Filename for copy of offloaded tracks as events");
     add_cmd(&options->max_num_tracks,
             "maxNumTracks",
             "Number of track \"slots\" to be transported simultaneously");

src/celeritas/CMakeLists.txt

@@ -194,6 +194,8 @@ if(CELERITAS_USE_ROOT)
     ext/RootUniquePtr.root.cc
     ext/RootFileManager.cc
     user/RootStepWriter.cc
+    io/RootEventReader.cc
+    io/RootEventWriter.cc
     "${CMAKE_CURRENT_BINARY_DIR}/CeleritasRootInterface.cxx"
   )
 

src/celeritas/ext/RootExporter.hh

@@ -38,7 +38,7 @@ class RootExporter
     void operator()(ImportData const& data);
 
   private:
-    UPRootWritable<TFile> root_output_;
+    UPRootFileWritable root_output_;
 
   private:
     // ROOT TTree name

src/celeritas/ext/RootFileManager.cc

@@ -46,13 +46,13 @@ RootFileManager::RootFileManager(char const* filename)
  * To expand this class to write multiple root files (one per thread), add a
  * `tid` input parameter and call `tfile_[tid].get()`.
  */
-UPRootWritable<TTree>
+UPRootTreeWritable
 RootFileManager::make_tree(char const* name, char const* title)
 {
     CELER_EXPECT(tfile_->IsOpen());
 
     int const split_level = 99;
-    UPRootWritable<TTree> uptree;
+    UPRootTreeWritable uptree;
     uptree.reset(new TTree(name, title, split_level, tfile_.get()));
     return uptree;
 }

src/celeritas/ext/RootFileManager.hh

@@ -33,13 +33,13 @@ class RootFileManager
     explicit RootFileManager(char const* filename);
 
     // Create tree by passing a name and title
-    UPRootWritable<TTree> make_tree(char const* name, char const* title);
+    UPRootTreeWritable make_tree(char const* name, char const* title);
 
     // Manually write TFile
     void write();
 
   public:
-    UPRootWritable<TFile> tfile_;
+    UPRootFileWritable tfile_;
 };
 
 //---------------------------------------------------------------------------//
@@ -49,8 +49,7 @@ inline RootFileManager::RootFileManager(char const*)
     CELER_NOT_CONFIGURED("ROOT");
 }
 
-inline UPRootWritable<TTree>
-RootFileManager::make_tree(char const*, char const*)
+inline UPRootTreeWritable RootFileManager::make_tree(char const*, char const*)
 {
     CELER_NOT_CONFIGURED("ROOT");
 }

src/celeritas/ext/RootUniquePtr.hh

@@ -13,17 +13,28 @@
 #include "celeritas_config.h"
 #include "corecel/Assert.hh"
 
+class TFile;
+class TTree;
+
 namespace celeritas
 {
 //---------------------------------------------------------------------------//
 /*!
- * Call `TObject->Write()` before deletion. Used by TFile and TTree writer
- * classes.
+ * Call `TFile->Write()` before deletion.
  */
-template<class T>
-struct RootWritableDeleter
+struct RootFileWritableDeleter
 {
-    void operator()(T* ptr);
+    void operator()(TFile* ptr);
+};
+
+//---------------------------------------------------------------------------//
+/*!
+ * Call `TTree->AutoSave()` before deletion in order to update `gDirectory`
+ * accordingly before writing the TTree.
+ */
+struct RootTreeAutoSaveDeleter
+{
+    void operator()(TTree* ptr);
 };
 
 //---------------------------------------------------------------------------//
@@ -39,15 +50,20 @@ struct ExternDeleter
 
 //---------------------------------------------------------------------------//
 // Type aliases
-template<class T>
-using UPRootWritable = std::unique_ptr<T, RootWritableDeleter<T>>;
+using UPRootFileWritable = std::unique_ptr<TFile, RootFileWritableDeleter>;
+using UPRootTreeWritable = std::unique_ptr<TTree, RootTreeAutoSaveDeleter>;
+
 template<class T>
 using UPExtern = std::unique_ptr<T, ExternDeleter<T>>;
 
 //---------------------------------------------------------------------------//
 #if !CELERITAS_USE_ROOT
-template<class T>
-void RootWritableDeleter<T>::operator()(T*)
+inline void RootFileWritableDeleter::operator()(TFile*)
+{
+    CELER_NOT_CONFIGURED("ROOT");
+}
+
+inline void RootTreeAutoSaveDeleter::operator()(TTree*)
 {
     CELER_NOT_CONFIGURED("ROOT");
 }

src/celeritas/ext/RootUniquePtr.root.cc

@@ -16,11 +16,9 @@ namespace celeritas
 {
 //---------------------------------------------------------------------------//
 /*!
- * Call `TObject->Write()` before deletion. Used by TFile and TTree writer
- * classes.
+ * Call `TFile->Write()` before deletion.
  */
-template<class T>
-void RootWritableDeleter<T>::operator()(T* ptr)
+void RootFileWritableDeleter::operator()(TFile* ptr)
 {
     CELER_EXPECT(ptr);
     CELER_LOG(debug) << "Writing " << ptr->ClassName() << " '"
@@ -29,6 +27,20 @@ void RootWritableDeleter<T>::operator()(T* ptr)
     delete ptr;
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Call `TTree->AutoSave()` before deletion. This ensures that `gDirectory` is
+ * updated accordingly when multiple TFiles are open at the same time.
+ */
+void RootTreeAutoSaveDeleter::operator()(TTree* ptr)
+{
+    CELER_EXPECT(ptr);
+    CELER_LOG(debug) << "Autosaving " << ptr->ClassName() << " '"
+                     << ptr->GetName() << "'";
+    ptr->AutoSave();
+    delete ptr;
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Custom deleter to avoid propagating any dependency-specific implementation
@@ -46,9 +58,6 @@ void ExternDeleter<T>::operator()(T* ptr)
 //---------------------------------------------------------------------------//
 // EXPLICIT INSTANTIATIONS
 //---------------------------------------------------------------------------//
-template struct RootWritableDeleter<TFile>;
-template struct RootWritableDeleter<TTree>;
-
 template struct ExternDeleter<TFile>;
 template struct ExternDeleter<TTree>;
 

src/celeritas/io/RootEventReader.cc

@@ -0,0 +1,118 @@
+//----------------------------------*-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 celeritas/io/RootEventReader.cc
+//---------------------------------------------------------------------------//
+#include "RootEventReader.hh"
+
+#include <TFile.h>
+#include <TLeaf.h>
+#include <TTree.h>
+
+#include "corecel/io/Logger.hh"
+#include "celeritas/ext/ScopedRootErrorHandler.hh"
+#include "celeritas/phys/ParticleParams.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Construct with ROOT input filename.
+ */
+RootEventReader::RootEventReader(std::string const& filename,
+                                 SPConstParticles params)
+    : params_(std::move(params))
+{
+    CELER_EXPECT(!filename.empty());
+    ScopedRootErrorHandler scoped_root_error;
+
+    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);
+
+    scoped_root_error.throw_if_errors();
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Read single event from the primaries tree.
+ */
+auto RootEventReader::operator()() -> result_type
+{
+    CELER_EXPECT(entry_count_ <= num_entries_);
+    ScopedRootErrorHandler scoped_root_error;
+
+    EventId expected_evt_id{0};
+    TrackId track_id{0};
+    result_type primaries;
+
+    for (; entry_count_ < num_entries_; entry_count_++)
+    {
+        ttree_->GetEntry(entry_count_);
+
+        auto entry_evt_id = EventId{this->from_leaf<size_type>("event_id")};
+        if (primaries.empty())
+        {
+            // First entry; set current event id
+            expected_evt_id = entry_evt_id;
+        }
+        if (entry_evt_id != expected_evt_id)
+        {
+            // End of primaries in this event
+            break;
+        }
+
+        Primary primary;
+        primary.track_id = track_id;
+        primary.event_id = expected_evt_id;
+        primary.particle_id
+            = params_->find(PDGNumber{this->from_leaf<int>("particle")});
+        primary.energy
+            = units::MevEnergy{this->from_leaf<real_type>("energy")};
+        primary.time = this->from_leaf<real_type>("time");
+        primary.position = this->from_array_leaf("pos");
+        primary.direction = this->from_array_leaf("dir");
+        primaries.push_back(std::move(primary));
+
+        track_id++;
+    }
+
+    scoped_root_error.throw_if_errors();
+    CELER_LOG_LOCAL(debug) << "Read event " << expected_evt_id.get()
+                           << " with " << primaries.size() << " primaries";
+    return primaries;
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Helper function to fetch leaves.
+ */
+template<class T>
+auto RootEventReader::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 RootEventReader::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 celeritas