Back to home page

EIC code displayed by LXR

 
 

    


Warning, file /include/celeritas/user/RootStepWriter.hh was not indexed or was modified since last indexation (in which case cross-reference links may be missing, inaccurate or erroneous).

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 primary_id = unspecified;
0062         size_type post_step_action_id = unspecified;
0063         size_type track_step_count = unspecified;
0064         real_type step_length = 0;  //!< [len]
0065         real_type energy_deposition = 0;  //!< [MeV]
0066         int particle = 0;  //!< PDG number
0067         EnumArray<StepPoint, TStepPoint> points;
0068     };
0069 
0070   public:
0071     //!@{
0072     //! \name Type aliases
0073     using SPRootFileManager = std::shared_ptr<RootFileManager>;
0074     using SPParticleParams = std::shared_ptr<ParticleParams const>;
0075     using WriteFilter = std::function<bool(TStepData const&)>;
0076     //!@}
0077 
0078     // Construct with step data writer filter
0079     RootStepWriter(SPRootFileManager root_manager,
0080                    SPParticleParams particle_params,
0081                    StepSelection selection,
0082                    WriteFilter filter);
0083 
0084     // Construct and store all step data
0085     RootStepWriter(SPRootFileManager root_manager,
0086                    SPParticleParams particle_params,
0087                    StepSelection selection);
0088 
0089     // Set number of entries stored in memory before being flushed to disk
0090     void set_auto_flush(long num_entries);
0091 
0092     // Process step data on the host and fill step tree
0093     void process_steps(HostStepState) final;
0094 
0095     // Device execution is not currently implemented
0096     void process_steps(DeviceStepState) final
0097     {
0098         CELER_NOT_IMPLEMENTED("RootStepWriter with device data");
0099     }
0100 
0101     // Selection of data to be stored
0102     StepSelection selection() const final { return selection_; }
0103 
0104     // No detector filtering selection is implemented
0105     Filters filters() const final { return {}; }
0106 
0107   private:
0108     // Create steps tree based on selection_ booleans
0109     void make_tree();
0110 
0111   private:
0112     SPRootFileManager root_manager_;
0113     SPParticleParams particles_;
0114     StepSelection selection_;
0115     UPRootTreeWritable tstep_tree_;
0116     TStepData tstep_;  // Members are used as refs of the TTree branches
0117     std::function<bool(TStepData const&)> filter_;
0118 };
0119 
0120 //---------------------------------------------------------------------------//
0121 // FREE FUNCTIONS
0122 //---------------------------------------------------------------------------//
0123 // Create a write filter for some simple IDs
0124 RootStepWriter::WriteFilter make_write_filter(SimpleRootFilterInput const&);
0125 
0126 //---------------------------------------------------------------------------//
0127 #if !CELERITAS_USE_ROOT
0128 inline RootStepWriter::RootStepWriter(
0129     SPRootFileManager, SPParticleParams, StepSelection, WriteFilter)
0130 {
0131     CELER_NOT_CONFIGURED("ROOT");
0132 }
0133 
0134 inline void RootStepWriter::process_steps(HostStepState)
0135 {
0136     CELER_NOT_CONFIGURED("ROOT");
0137 }
0138 
0139 inline RootStepWriter::WriteFilter make_write_filter(
0140     SimpleRootFilterInput const&)
0141 {
0142     return nullptr;
0143 }
0144 
0145 #endif
0146 
0147 //---------------------------------------------------------------------------//
0148 }  // namespace celeritas