celeritas/user/RootStepWriter.hh

0001 //------------------------------- -*- C++ -*- -------------------------------//
0002 // Copyright Celeritas contributors: see top-level COPYRIGHT file for details
0003 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0004 //---------------------------------------------------------------------------//
0005 //! \file celeritas/user/RootStepWriter.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include <array>
0010
0011 #include "corecel/Config.hh"
0012
0013 #include "corecel/Assert.hh"
0014 #include "celeritas/ext/RootUniquePtr.hh"
0015
0016 #include "RootStepWriterInput.hh"
0017 #include "StepInterface.hh"
0018
0019 class TTree;
0020
0021 namespace celeritas
0022 {
0023 //---------------------------------------------------------------------------//
0024 class ParticleParams;
0025 class RootFileManager;
0026
0027 //---------------------------------------------------------------------------//
0028 /*!
0029  * Write "MC truth" data to ROOT at every step.
0030  *
0031  * `TTree::Fill()` is called for each step and thread id, making each ROOT
0032  * entry a step. Since the ROOT data is stored in branches with primitive types
0033  * instead of a full struct, no dictionaries are needed for reading the output
0034  * file.
0035  *
0036  * The step data that is written to the ROOT file can be filtered by providing
0037  * a user-defined `WriteFilter` function.
0038  */
0039 class RootStepWriter final : public StepInterface
0040 {
0041   public:
0042     // Unspecified step attribute data value
0043     static constexpr size_type unspecified{static_cast<size_type>(-1)};
0044
0045     //! Truth step point data; Naming convention must match StepPointStateData
0046     struct TStepPoint
0047     {
0048         size_type volume_id = unspecified;
0049         real_type energy = 0;  //!< [MeV]
0050         real_type time = 0;  //!< [time]
0051         std::array<real_type, 3> pos{0, 0, 0};  //!< [len]
0052         std::array<real_type, 3> dir{0, 0, 0};
0053     };
0054
0055     //! Truth step data; Naming convention must match StepStateData
0056     struct TStepData
0057     {
0058         size_type event_id = unspecified;
0059         size_type track_id = unspecified;
0060         size_type parent_id = unspecified;
0061         size_type action_id = unspecified;
0062         size_type track_step_count = unspecified;
0063         int particle = 0;  //!< PDG number
0064         real_type energy_deposition = 0;  //!< [MeV]
0065         real_type step_length = 0;  //!< [len]
0066         EnumArray<StepPoint, TStepPoint> points;
0067     };
0068
0069   public:
0070     //!@{
0071     //! \name Type aliases
0072     using SPRootFileManager = std::shared_ptr<RootFileManager>;
0073     using SPParticleParams = std::shared_ptr<ParticleParams const>;
0074     using WriteFilter = std::function<bool(TStepData const&)>;
0075     //!@}
0076
0077     // Construct with step data writer filter
0078     RootStepWriter(SPRootFileManager root_manager,
0079                    SPParticleParams particle_params,
0080                    StepSelection selection,
0081                    WriteFilter filter);
0082
0083     // Construct and store all step data
0084     RootStepWriter(SPRootFileManager root_manager,
0085                    SPParticleParams particle_params,
0086                    StepSelection selection);
0087
0088     // Set number of entries stored in memory before being flushed to disk
0089     void set_auto_flush(long num_entries);
0090
0091     // Process step data on the host and fill step tree
0092     void process_steps(HostStepState) final;
0093
0094     // Device execution is not currently implemented
0095     void process_steps(DeviceStepState) final
0096     {
0097         CELER_NOT_IMPLEMENTED("RootStepWriter with device data");
0098     }
0099
0100     // Selection of data to be stored
0101     StepSelection selection() const final { return selection_; }
0102
0103     // No detector filtering selection is implemented
0104     Filters filters() const final { return {}; }
0105
0106   private:
0107     // Create steps tree based on selection_ booleans
0108     void make_tree();
0109
0110   private:
0111     SPRootFileManager root_manager_;
0112     SPParticleParams particles_;
0113     StepSelection selection_;
0114     UPRootTreeWritable tstep_tree_;
0115     TStepData tstep_;  // Members are used as refs of the TTree branches
0116     std::function<bool(TStepData const&)> filter_;
0117 };
0118
0119 //---------------------------------------------------------------------------//
0120 // FREE FUNCTIONS
0121 //---------------------------------------------------------------------------//
0122 // Create a write filter for some simple IDs
0123 RootStepWriter::WriteFilter make_write_filter(SimpleRootFilterInput const&);
0124
0125 //---------------------------------------------------------------------------//
0126 #if !CELERITAS_USE_ROOT
0127 inline RootStepWriter::RootStepWriter(SPRootFileManager,
0128                                       SPParticleParams,
0129                                       StepSelection,
0130                                       WriteFilter)
0131 {
0132     CELER_NOT_CONFIGURED("ROOT");
0133 }
0134
0135 inline void RootStepWriter::process_steps(HostStepState)
0136 {
0137     CELER_NOT_CONFIGURED("ROOT");
0138 }
0139
0140 inline RootStepWriter::WriteFilter
0141 make_write_filter(SimpleRootFilterInput const&)
0142 {
0143     return nullptr;
0144 }
0145
0146 #endif
0147
0148 //---------------------------------------------------------------------------//
0149 }  // namespace celeritas