dd4hep-plugin/src/TextDumpingSteppingAction.cpp

0001 //==========================================================================
0002 //  AIDA Detector description implementation
0003 //--------------------------------------------------------------------------
0004 // Copyright (C) Organisation europeenne pour la Recherche nucleaire (CERN)
0005 // All rights reserved.
0006 //
0007 // For the licensing terms see $DD4hepINSTALL/LICENSE.
0008 // For the list of contributors see $DD4hepINSTALL/doc/CREDITS.
0009 //
0010 //==========================================================================
0011
0012 // Framework include files
0013 #include "DDG4/Geant4SteppingAction.h"
0014
0015 #include <G4Step.hh>
0016 #include <G4RunManager.hh>
0017 #include <G4TrackStatus.hh>
0018 #include <G4Run.hh>
0019 #include <G4Event.hh>
0020 #include <G4UnitsTable.hh>
0021 #include <stack>
0022 #include <CLHEP/Units/SystemOfUnits.h>
0023
0024 #include <fmt/core.h>
0025
0026 /// Namespace for the AIDA detector description toolkit
0027 namespace dd4hep {
0028
0029     /// Namespace for the Geant4 based simulation part of the AIDA detector description toolkit
0030     namespace sim {
0031
0032         /// Class to count steps and suspend tracks every 5 steps
0033         /** Class to count steps and suspens tracks every 5 steps
0034          * Count steps and suspend
0035          *
0036          *  \version 1.0
0037          *  \ingroup DD4HEP_SIMULATION
0038          */
0039         class TextDumpingSteppingAction : public Geant4SteppingAction {
0040             std::size_t m_total_steps {0UL };
0041             std::size_t m_total_tracks {0UL };
0042             std::size_t m_total_events {0UL };
0043             std::string  m_file_name {"events_stepping.txt"};
0044             std::string m_class_name;
0045             std::ofstream m_output_file;
0046             std::once_flag m_init_flag;
0047             double m_vertex_z_min = -4500;
0048             double m_vertex_z_max = 4500;
0049             bool m_vertex_cut = true;
0050
0051             bool m_save_optical = false;                    // When TRUE - If OpticalPhoton - it will be saved. All other cuts don't affect optical photons when this is TRUE. I.e. Optical photon? - SAVE
0052             bool m_only_primary = true;                     // Keep geant tracks and their steps only if geant track has ParentID=0
0053             int m_prev_event = -1;
0054             int m_prev_track_id = -1;
0055             // std::stack<std::size_t> m_gen_part_ids;
0056             bool m_skipping_track = false;
0057             double m_min_momentum = 300 * CLHEP::MeV;      // Minimal momentum for track to be saved
0058             double m_max_momentum = 10000 * CLHEP::TeV;    // Maximum momentum for track to be saved
0059
0060         public:
0061             /// Standard constructor
0062             TextDumpingSteppingAction(Geant4Context* context, const std::string& nam)
0063                     : Geant4SteppingAction(context, nam)
0064             {
0065                 declareProperty("OutputFile",    m_file_name);
0066                 declareProperty("MomentumMin",   m_min_momentum);
0067                 declareProperty("SaveOptical",   m_save_optical);
0068                 declareProperty("OnlyPrimary",   m_only_primary);
0069                 declareProperty("VertexCut",     m_vertex_cut);
0070                 declareProperty("VertexZMin",    m_vertex_z_min);
0071                 declareProperty("VertexZMax",    m_vertex_z_max);
0072                 m_class_name = __func__;        // Must be in constructor to be the easiest way to get demangled name
0073             }
0074             /// Default destructor
0075             ~TextDumpingSteppingAction() override
0076             {
0077                 info("+++ Total Steps: %ld", m_total_steps);
0078                 info("+++ Total Tracks written: %ld", m_total_tracks);
0079                 info("+++ Total Events written: %ld", m_total_events);
0080                 // info("+++ Per Event steps: %ld, tracks: %ld", (int) m_total_steps/m_total_events, (int) m_total_tracks/m_total_events);
0081
0082                 // Do we need it here? It should be done automatically
0083                 if(m_output_file.is_open() && m_output_file.good()) {
0084                     m_output_file.close();
0085                 }
0086             }
0087
0088             /// Checks m_output_file is open and is ready for write or throws an exception
0089             void ensure_output_writable() {
0090                 if(!m_output_file.is_open() || !m_output_file.good()) {
0091                     auto err_msg = fmt::format( "Failed to open the file or file stream is in a bad state. File name: '{}'", m_file_name);
0092                     error(err_msg.c_str());
0093                     throw std::runtime_error(err_msg);
0094                 }
0095             }
0096
0097             void write_point(G4StepPoint *point){
0098                 auto row = fmt::format("P {} {} {} {}",
0099                             point->GetPosition().x(),
0100                             point->GetPosition().y(),
0101                             point->GetPosition().z(),
0102                             point->GetGlobalTime()
0103                             );
0104                 m_output_file<<row<<std::endl;
0105             }
0106
0107             /// stepping callback
0108             void operator()(const G4Step* step, G4SteppingManager*) override
0109             {
0110
0111                 // Get run and event number
0112                 auto run_num = context()->run().run().GetRunID();
0113                 auto event_num = context()->event().event().GetEventID();
0114
0115
0116                 // First time in this function. Open a file
0117                 std::call_once(m_init_flag, [this](){
0118
0119                     fmt::print("Plugin {} info: \n", m_class_name);
0120                     fmt::print("   OutputFile     {}\n", m_file_name);
0121                     fmt::print("   MomentumMin    {}\n", m_min_momentum);
0122                     fmt::print("   MomentumMax    {}\n", m_max_momentum);
0123                     fmt::print("   SaveOptical    {}\n", m_save_optical);
0124                     fmt::print("   OnlyPrimary    {}\n", m_only_primary);
0125                     fmt::print("   VertexCut      {}\n", m_vertex_cut);
0126                     fmt::print("   VertexZMin     {}\n", m_vertex_z_min);
0127                     fmt::print("   VertexZMax     {}\n", m_vertex_z_max);
0128
0129                     m_output_file = std::ofstream(m_file_name);
0130
0131                     ensure_output_writable();
0132                     m_output_file<<"#Format description" << std::endl;
0133                     m_output_file<<"#  E - event: run_num event_num" << std::endl;
0134                     m_output_file<<"#  T - track: id, status, pdg, pdg_name, eta, phi, qOverP, px, py, pz, vx, vy, vz" << std::endl;
0135                     m_output_file<<"#  P - point: x, y, z, t" << std::endl;
0136                 });
0137
0138                 // Check file is open and writable
0139                 ensure_output_writable();
0140
0141                 if(event_num != m_prev_event) {
0142                     // We've got a new event
0143                     m_prev_event = event_num;
0144                     std::string evt_str = fmt::format("E {} {} ", run_num, event_num);
0145                     m_prev_track_id = -1;   // Reset track id, so any track is new
0146                     m_output_file << evt_str << std::endl;
0147                     m_total_events++;
0148                     fmt::print("+\n+\n+\n+\n+\n=====================================================\n");
0149                 }
0150
0151                 // Check if this is new track
0152                 auto track = step->GetTrack();
0153                 if(track->GetTrackID() != m_prev_track_id)
0154                 {
0155                     m_prev_track_id = track->GetTrackID();
0156
0157                     // New track
0158                     auto id = track->GetTrackID();
0159                     auto pdg = track->GetParticleDefinition()->GetPDGEncoding();
0160                     auto pdg_name = track->GetParticleDefinition()->GetParticleName();
0161                     auto charge = track->GetParticleDefinition()->GetPDGCharge();
0162                     auto eta = track->GetMomentum().eta();
0163                     auto phi = track->GetMomentum().phi();
0164                     auto qOverP = track->GetParticleDefinition()->GetPDGCharge()/track->GetMomentum().mag();
0165                     auto px = track->GetMomentum().x();
0166                     auto py = track->GetMomentum().y();
0167                     auto pz = track->GetMomentum().z();
0168                     auto p = track->GetMomentum().mag();
0169                     auto vx = track->GetVertexPosition().x();
0170                     auto vy = track->GetVertexPosition().y();
0171                     auto vz = track->GetVertexPosition().z();
0172
0173                     bool filter = m_vertex_cut && (vz < m_vertex_z_min || vz > m_vertex_z_max);
0174
0175                     if(m_prev_track_id < 1000)
0176                     {
0177                         fmt::print("track: {:<5}, {:<10} vtx: {:12.5f} {:12.5f} {:12.5f}  {:5} {:12.5f} {:12.5f} {:5}\n",
0178                             id, pdg_name, vx,vy,vz, filter, m_vertex_z_min, m_vertex_z_max, m_vertex_cut);
0179                     }
0180
0181
0182                     //if(track->GetParentID() == 0 || pdg_name == "opticalphoton") {
0183
0184                     // First we filter by track types
0185                     bool track_is_ok = true;
0186
0187                     if(m_only_primary)
0188                     {
0189                         // >oO fmt::print("{} ID: {} ParentID: {} CrModelID: {}\n", pdg_name, id, track->GetParentID(), track->GetCreatorModelID());
0190                         track_is_ok = track->GetParentID() == 0;   // Only tracks without parents are OK
0191                     }
0192
0193                     if(track_is_ok && (p < m_min_momentum || p > m_max_momentum))
0194                     {
0195                         track_is_ok = false;
0196                     }
0197
0198
0199                     if(m_vertex_cut && (vz < m_vertex_z_min || vz > m_vertex_z_max))
0200                     {
0201                         track_is_ok = false;
0202                     }
0203
0204                     if(m_save_optical && pdg_name == "opticalphoton")
0205                     {
0206                         track_is_ok = true;     // Force to save track whatever
0207                     }
0208
0209                     m_skipping_track = !track_is_ok;
0210
0211
0212                     if(m_skipping_track) return;        //  <= nothing to do more here if the track is to be skipped
0213
0214                     std::string trk_str = fmt::format("T {} {} {} {} {} {} {} {} {} {} {} {} {}",
0215                                                       id, pdg, pdg_name, charge, eta, phi, qOverP, px, py, pz, vx, vy, vz);
0216                     m_output_file << trk_str << std::endl;
0217                     m_total_tracks++;
0218
0219                     // We always write GetPostStepPoint so the first time we have to write GetPreStepPoint
0220                     write_point(step->GetPreStepPoint());
0221
0222                 }
0223
0224                 // Maybe we skipping this track and we just return
0225                 if(m_skipping_track) return;
0226
0227                 // Write post point info
0228                 write_point(step->GetPostStepPoint());
0229
0230                 // Statistics
0231                 m_total_steps++;
0232             }
0233         };
0234     }    // End namespace sim
0235 }      // End namespace dd4hep
0236
0237 #include "DDG4/Factories.h"
0238 DECLARE_GEANT4ACTION_NS(dd4hep::sim,TextDumpingSteppingAction)