dd4hep-plugin/src/FirebirdTrajectoryWriterEventAction.cpp

0001 //==========================================================================
0002 //  AIDA Detector description implementation
0003 //--------------------------------------------------------------------------
0004 //  FirebirdTrajectoryWriterEventAction — DD4hep/DDG4 plugin
0005 //
0006 //  Writes Geant4 trajectories as Firebird-format JSON for event display.
0007 //  Supports rich-trajectory time extraction, configurable filtering,
0008 //  and per-point verbose diagnostics.
0009 //==========================================================================
0010
0011 // DD4hep / DDG4
0012 #include "DDG4/Geant4EventAction.h"
0013 #include "DDG4/Geant4Kernel.h"
0014 #include "DD4hep/Printout.h"
0015
0016 // Geant4
0017 #include "G4Event.hh"
0018 #include "G4TrajectoryContainer.hh"
0019 #include "G4VTrajectory.hh"
0020 #include "G4VTrajectoryPoint.hh"
0021 #include "G4RichTrajectory.hh"
0022 #include "G4RichTrajectoryPoint.hh"
0023 #include "G4SystemOfUnits.hh"
0024 #include "G4AttValue.hh"
0025 #include "G4AttDef.hh"
0026
0027 // C++ standard
0028 #include <algorithm>
0029 #include <cmath>
0030 #include <fstream>
0031 #include <limits>
0032 #include <sstream>
0033 #include <string>
0034 #include <vector>
0035
0036 #include <fmt/core.h>
0037 #include <fmt/format.h>
0038
0039 namespace dd4hep::sim {
0040
0041   // ──────────────────────────────────────────────────────────────────────
0042   //  FirebirdTrajectoryWriterEventAction
0043   // ──────────────────────────────────────────────────────────────────────
0044
0045   /// Writes filtered Geant4 trajectories to a JSON file in the Firebird
0046   /// event-display format, with robust time extraction from G4RichTrajectoryPoint.
0047   ///
0048   /// \version 1.2
0049   /// \ingroup DD4HEP_SIMULATION
0050   class FirebirdTrajectoryWriterEventAction : public Geant4EventAction {
0051
0052     // ── configuration ──────────────────────────────────────────────────
0053
0054     /// Output path for the JSON file.
0055     std::string m_outputFile{"trajectories.firebird.json"};
0056
0057     /// Logical component name written into the JSON.
0058     std::string m_componentName{"Geant4Trajectories"};
0059
0060     // Particle / track filters
0061     bool              m_saveOptical{false};   ///< Always save optical photons
0062     bool              m_onlyPrimary{false};   ///< Keep only ParentID == 0
0063     std::vector<int>  m_saveParticles{};      ///< PDG whitelist (empty ⇒ all)
0064     double            m_minMomentum{150};     ///< Lower momentum cut [MeV/c]
0065     double            m_maxMomentum{1e12};    ///< Upper momentum cut [MeV/c]
0066     double            m_minTrackLength{0};    ///< Minimum path length  [mm]
0067
0068     // Vertex position cut
0069     bool   m_vertexCut{false};
0070     double m_vertexZMin{-5000};   ///< [mm]
0071     double m_vertexZMax{ 5000};   ///< [mm]
0072
0073     // Step (point) position cut
0074     bool   m_stepCut{false};
0075     double m_stepZMin{-5000};     ///< [mm]
0076     double m_stepZMax{ 5000};     ///< [mm]
0077     double m_stepRMax{ 5000};     ///< [mm]
0078
0079     // Time extraction behaviour
0080     bool m_requireRichTrajectory{true};
0081
0082     // Diagnostics
0083     bool m_verboseTimeExtraction{false};  ///< Log time-extraction internals
0084     bool m_verboseSteps{false};           ///< Log every point with full details
0085
0086     // ── statistics ─────────────────────────────────────────────────────
0087
0088     long m_totalTrajectories{0};
0089     long m_filteredTrajectories{0};
0090     long m_savedTrajectories{0};
0091     long m_trajectoryWithoutTime{0};
0092     long m_stepsFiltered{0};
0093
0094     // ── per-run buffer ─────────────────────────────────────────────────
0095
0096     std::vector<std::string> m_entries;   ///< One JSON string per event
0097
0098     // ── helpers ────────────────────────────────────────────────────────
0099
0100     /// Return true when `v` is a finite number safe for JSON serialisation.
0101     static bool isFinite(double v) { return !std::isinf(v) && !std::isnan(v); }
0102
0103     /// Replace non-finite values with `fallback`.
0104     static double sanitise(double v, double fallback = 0.0) {
0105       return isFinite(v) ? v : fallback;
0106     }
0107
0108     /// Transverse distance from the beam axis.
0109     static double rxy(const G4ThreeVector& p) {
0110       return std::sqrt(p.x() * p.x() + p.y() * p.y());
0111     }
0112
0113     // ── time extraction ────────────────────────────────────────────────
0114
0115     /// Convert a Geant4 time string (value + optional unit) to internal
0116     /// units.  Handles ns, ps, us/µs, ms, s.  Falls back to ns.
0117     static G4double parseTimeString(const std::string& raw) {
0118       std::istringstream iss(raw);
0119       G4double value{};
0120       std::string unit;
0121       iss >> value >> unit;
0122
0123       if      (unit == "ps")                   return value * CLHEP::picosecond;
0124       else if (unit == "ns")                   return value * CLHEP::ns;
0125       else if (unit == "us" || unit == "µs")   return value * CLHEP::microsecond;
0126       else if (unit == "ms")                   return value * CLHEP::ms;
0127       else if (unit == "s")                    return value * CLHEP::s;
0128       else                                     return value * CLHEP::ns;  // default
0129     }
0130
0131     /// Try to read PreT (index 0) or PostT (index > 0) from a
0132     /// G4RichTrajectoryPoint.  Returns the time in internal units or
0133     /// –1 on failure.
0134     ///
0135     /// When `trajectory` is non-null and `m_verboseSteps` is set the
0136     /// method prints a one-line diagnostic for each point.
0137     G4double extractTimeFromPoint(G4VTrajectoryPoint* point,
0138                                   int pointIndex,
0139                                   G4VTrajectory* trajectory = nullptr)
0140     {
0141       auto* richPoint = dynamic_cast<G4RichTrajectoryPoint*>(point);
0142       if (!richPoint) {
0143         if (m_requireRichTrajectory) {
0144           if (m_verboseTimeExtraction) {
0145             warning("[firebird-writer] Point %d: not a G4RichTrajectoryPoint", pointIndex);
0146           }
0147           ++m_trajectoryWithoutTime;
0148           return -1.0;
0149         }
0150         return pointIndex * 0.1 * CLHEP::ns;   // synthetic fallback
0151       }
0152
0153       auto* attValues = richPoint->CreateAttValues();
0154       if (!attValues) {
0155         if (m_verboseTimeExtraction) {
0156           warning("[firebird-writer] Point %d: CreateAttValues() returned null", pointIndex);
0157         }
0158         return -1.0;
0159       }
0160
0161       const std::string targetAttr = (pointIndex == 0) ? "PreT" : "PostT";
0162       G4double extractedTime = -1.0;
0163
0164       for (const auto& av : *attValues) {
0165         if (av.GetName() == targetAttr) {
0166           extractedTime = parseTimeString(av.GetValue());
0167
0168           if (m_verboseTimeExtraction) {
0169             info("[firebird-writer] Point %d: %s raw=\"%s\" → %.6f ns",
0170                  pointIndex, targetAttr.c_str(),
0171                  av.GetValue().c_str(), extractedTime / CLHEP::ns);
0172           }
0173           break;
0174         }
0175       }
0176       delete attValues;
0177
0178       // ── verbose per-step dump ──────────────────────────────────────
0179       if (m_verboseSteps) {
0180         const auto pos = point->GetPosition();
0181         if (trajectory) {
0182           info("[firebird-steps] trk PDG=%d  pt=%d  pos=(%.3f, %.3f, %.3f) mm  "
0183                "t=%.6f ns  attr=%s",
0184                trajectory->GetPDGEncoding(), pointIndex,
0185                pos.x() / CLHEP::mm, pos.y() / CLHEP::mm, pos.z() / CLHEP::mm,
0186                (extractedTime >= 0 ? extractedTime / CLHEP::ns : -1.0),
0187                targetAttr.c_str());
0188         } else {
0189           info("[firebird-steps] pt=%d  pos=(%.3f, %.3f, %.3f) mm  t=%.6f ns  attr=%s",
0190                pointIndex,
0191                pos.x() / CLHEP::mm, pos.y() / CLHEP::mm, pos.z() / CLHEP::mm,
0192                (extractedTime >= 0 ? extractedTime / CLHEP::ns : -1.0),
0193                targetAttr.c_str());
0194         }
0195       }
0196
0197       // Handle extraction failure
0198       if (extractedTime < 0) {
0199         if (m_requireRichTrajectory) {
0200           if (m_verboseTimeExtraction) {
0201             warning("[firebird-writer] Point %d: attribute %s not found", pointIndex, targetAttr.c_str());
0202           }
0203           ++m_trajectoryWithoutTime;
0204           return -1.0;
0205         }
0206         return pointIndex * 0.1 * CLHEP::ns;
0207       }
0208       return extractedTime;
0209     }
0210
0211     // ── filtering ──────────────────────────────────────────────────────
0212
0213     bool passesFilters(G4VTrajectory* trj) const {
0214       const int    pdg      = trj->GetPDGEncoding();
0215       const int    parentID = trj->GetParentID();
0216       const double p_MeV    = trj->GetInitialMomentum().mag() / CLHEP::MeV;
0217
0218       // Optical photons bypass everything when requested
0219       if (m_saveOptical && trj->GetParticleName() == "opticalphoton") return true;
0220
0221       if (m_onlyPrimary && parentID != 0) return false;
0222
0223       if (p_MeV < m_minMomentum || p_MeV > m_maxMomentum) return false;
0224
0225       // PDG whitelist
0226       if (!m_saveParticles.empty()) {
0227         if (std::find(m_saveParticles.begin(), m_saveParticles.end(), pdg)
0228             == m_saveParticles.end())
0229           return false;
0230       }
0231
0232       // Minimum track length
0233       if (m_minTrackLength > 0) {
0234         const int npts = trj->GetPointEntries();
0235         if (npts <= 1) return false;
0236
0237         double length = 0;
0238         auto prev = trj->GetPoint(0)->GetPosition();
0239         for (int i = 1; i < npts; ++i) {
0240           auto cur = trj->GetPoint(i)->GetPosition();
0241           length += (cur - prev).mag();
0242           prev = cur;
0243         }
0244         if (length / CLHEP::mm < m_minTrackLength) return false;
0245       }
0246
0247       // Vertex Z window
0248       if (m_vertexCut && trj->GetPointEntries() > 0) {
0249         const double vz = trj->GetPoint(0)->GetPosition().z() / CLHEP::mm;
0250         if (vz < m_vertexZMin || vz > m_vertexZMax) return false;
0251       }
0252
0253       return true;
0254     }
0255
0256     /// Additional check that requires mutable state (time extraction
0257     /// counters), so it's separate from the const filter above.
0258     bool passesRichTrajectoryCheck(G4VTrajectory* trj) {
0259       if (!m_requireRichTrajectory) return true;
0260
0261       if (!dynamic_cast<G4RichTrajectory*>(trj)) {
0262         if (m_verboseTimeExtraction)
0263           warning("[firebird-writer] Trajectory is not G4RichTrajectory — skipped");
0264         return false;
0265       }
0266       if (trj->GetPointEntries() > 0) {
0267         if (extractTimeFromPoint(trj->GetPoint(0), 0, trj) < 0) {
0268           if (m_verboseTimeExtraction)
0269             warning("[firebird-writer] First point has no time — trajectory skipped");
0270           return false;
0271         }
0272       }
0273       return true;
0274     }
0275
0276     // ── JSON builders ──────────────────────────────────────────────────
0277
0278     /// Produce the per-trajectory parameter array (JSON).
0279     std::string buildParamsJson(G4VTrajectory* trj) {
0280       const auto mom = trj->GetInitialMomentum();
0281       const double p = std::max(mom.mag(), 1e-10);  // avoid /0
0282
0283       const int    pdg    = trj->GetPDGEncoding();
0284       const double charge = trj->GetCharge();
0285       const double theta  = mom.theta();
0286       const double phi    = mom.phi();
0287       const double qOverP = charge / (p / CLHEP::GeV);
0288
0289       double vx = 0, vy = 0, vz = 0, time = 0;
0290       if (trj->GetPointEntries() > 0) {
0291         auto* pt0 = trj->GetPoint(0);
0292         const auto pos = pt0->GetPosition();
0293         vx = pos.x() / CLHEP::mm;
0294         vy = pos.y() / CLHEP::mm;
0295         vz = pos.z() / CLHEP::mm;
0296
0297         double t = extractTimeFromPoint(pt0, 0, trj);
0298         time = (t >= 0 ? t : 0.0) / CLHEP::ns;
0299       }
0300
0301       // pdg, type, charge, px, py, pz, vx, vy, vz, theta, phi, q/p, locA, locB, t
0302       return fmt::format("[{},\"{}\",{},{},{},{},{},{},{},{},{},{},{},{},{}]",
0303                          pdg, trj->GetParticleName(), sanitise(charge),
0304                          sanitise(mom.x() / CLHEP::MeV),
0305                          sanitise(mom.y() / CLHEP::MeV),
0306                          sanitise(mom.z() / CLHEP::MeV),
0307                          sanitise(vx), sanitise(vy), sanitise(vz),
0308                          sanitise(theta), sanitise(phi), sanitise(qOverP),
0309                          0.0, 0.0, sanitise(time));
0310     }
0311
0312     /// Produce the points array (JSON) for one trajectory.
0313     std::string buildPointsJson(G4VTrajectory* trj) {
0314       const int npts = trj->GetPointEntries();
0315       if (npts == 0) return "[]";
0316
0317       std::string out = "[";
0318       bool first = true;
0319
0320       for (int i = 0; i < npts; ++i) {
0321         auto* pt  = trj->GetPoint(i);
0322         auto  pos = pt->GetPosition();
0323
0324         if (m_stepCut) {
0325           const double z = pos.z() / CLHEP::mm;
0326           const double r = rxy(pos) / CLHEP::mm;
0327           if (z < m_stepZMin || z > m_stepZMax || r > m_stepRMax) {
0328             ++m_stepsFiltered;
0329             continue;
0330           }
0331         }
0332
0333         double t = extractTimeFromPoint(pt, i, trj);
0334         if (t < 0) t = i * 0.1 * CLHEP::ns;
0335         t /= CLHEP::ns;
0336
0337         if (!first) out += ',';
0338         first = false;
0339
0340         out += fmt::format("[{},{},{},{},{}]",
0341                            sanitise(pos.x() / CLHEP::mm),
0342                            sanitise(pos.y() / CLHEP::mm),
0343                            sanitise(pos.z() / CLHEP::mm),
0344                            sanitise(t), 0);
0345       }
0346       out += ']';
0347       return out;
0348     }
0349
0350     // ── file output ────────────────────────────────────────────────────
0351
0352     void writeJsonFile() const {
0353       if (m_entries.empty()) {
0354         warning("[firebird-writer] No events collected — output file not created.");
0355         return;
0356       }
0357
0358       std::ofstream out(m_outputFile);
0359       if (!out.is_open()) {
0360         error("[firebird-writer] Cannot open output file: %s", m_outputFile.c_str());
0361         return;
0362       }
0363
0364       out << fmt::format(
0365         R"({{"type":"firebird-dex-json","version":"0.04",)"
0366         R"("origin":{{"file":"{}","entries_count":{}}},)"
0367         R"("events":[)", m_outputFile, m_entries.size());
0368
0369       for (size_t i = 0; i < m_entries.size(); ++i) {
0370         if (i) out << ',';
0371         out << m_entries[i];
0372       }
0373       out << "]}";
0374       out.close();
0375
0376       info("[firebird-writer] Wrote %zu event(s) to %s",
0377            m_entries.size(), m_outputFile.c_str());
0378     }
0379
0380     void printStatistics() const {
0381       auto pct = [&](long n) {
0382         return m_totalTrajectories > 0
0383                  ? n * 100.0 / m_totalTrajectories : 0.0;
0384       };
0385       info("[firebird-writer] ── statistics ──────────────────────");
0386       info("[firebird-writer]  Total trajectories  : %ld", m_totalTrajectories);
0387       info("[firebird-writer]  Filtered (skipped)   : %ld (%.1f%%)",
0388            m_filteredTrajectories, pct(m_filteredTrajectories));
0389       info("[firebird-writer]  Saved                : %ld (%.1f%%)",
0390            m_savedTrajectories, pct(m_savedTrajectories));
0391       if (m_stepCut)
0392         info("[firebird-writer]  Step points filtered : %ld", m_stepsFiltered);
0393       if (m_requireRichTrajectory)
0394         info("[firebird-writer]  Missing time info    : %ld", m_trajectoryWithoutTime);
0395     }
0396
0397     void logConfiguration() const {
0398       info("[firebird-writer] ── configuration ───────────────────");
0399       info("[firebird-writer]  OutputFile            : %s", m_outputFile.c_str());
0400       info("[firebird-writer]  ComponentName         : %s", m_componentName.c_str());
0401       info("[firebird-writer]  SaveOptical           : %s", m_saveOptical ? "true" : "false");
0402       info("[firebird-writer]  OnlyPrimary           : %s", m_onlyPrimary ? "true" : "false");
0403       info("[firebird-writer]  VertexCut             : %s (Z: %.2f – %.2f mm)",
0404            m_vertexCut ? "true" : "false", m_vertexZMin, m_vertexZMax);
0405       info("[firebird-writer]  StepCut               : %s (Z: %.2f – %.2f mm, R < %.2f mm)",
0406            m_stepCut ? "true" : "false", m_stepZMin, m_stepZMax, m_stepRMax);
0407       info("[firebird-writer]  Momentum              : %.3f – %.3g MeV/c", m_minMomentum, m_maxMomentum);
0408       info("[firebird-writer]  MinTrackLength        : %.2f mm", m_minTrackLength);
0409       info("[firebird-writer]  RequireRichTrajectory : %s", m_requireRichTrajectory ? "true" : "false");
0410       info("[firebird-writer]  VerboseTimeExtraction : %s", m_verboseTimeExtraction ? "true" : "false");
0411       info("[firebird-writer]  VerboseSteps          : %s", m_verboseSteps ? "true" : "false");
0412
0413       if (m_saveParticles.empty()) {
0414         info("[firebird-writer]  SaveParticles         : [all]");
0415       } else {
0416         std::ostringstream ss;
0417         for (size_t i = 0; i < m_saveParticles.size(); ++i) {
0418           if (i) ss << ", ";
0419           ss << m_saveParticles[i];
0420         }
0421         info("[firebird-writer]  SaveParticles         : %s", ss.str().c_str());
0422       }
0423     }
0424
0425   public:
0426
0427     // ── lifecycle ──────────────────────────────────────────────────────
0428
0429     FirebirdTrajectoryWriterEventAction(Geant4Context* context,
0430                                         const std::string& name = "FirebirdTrajectoryWriterEventAction")
0431       : Geant4EventAction(context, name)
0432     {
0433       declareProperty("OutputFile",              m_outputFile);
0434       declareProperty("ComponentName",           m_componentName);
0435       declareProperty("SaveOptical",             m_saveOptical);
0436       declareProperty("OnlyPrimary",             m_onlyPrimary);
0437       declareProperty("VertexCut",               m_vertexCut);
0438       declareProperty("VertexZMin",              m_vertexZMin);
0439       declareProperty("VertexZMax",              m_vertexZMax);
0440       declareProperty("StepCut",                 m_stepCut);
0441       declareProperty("StepZMin",                m_stepZMin);
0442       declareProperty("StepZMax",                m_stepZMax);
0443       declareProperty("StepRMax",                m_stepRMax);
0444       declareProperty("MomentumMin",             m_minMomentum);
0445       declareProperty("MomentumMax",             m_maxMomentum);
0446       declareProperty("TrackLengthMin",          m_minTrackLength);
0447       declareProperty("SaveParticles",           m_saveParticles);
0448       declareProperty("RequireRichTrajectory",   m_requireRichTrajectory);
0449       declareProperty("VerboseTimeExtraction",   m_verboseTimeExtraction);
0450       declareProperty("VerboseSteps",            m_verboseSteps);
0451
0452       logConfiguration();
0453     }
0454
0455     ~FirebirdTrajectoryWriterEventAction() override {
0456       writeJsonFile();
0457       printStatistics();
0458     }
0459
0460     // ── event callbacks ────────────────────────────────────────────────
0461
0462     void begin(const G4Event* /*event*/) override { /* nothing */ }
0463
0464     void end(const G4Event* event) override {
0465       auto* container = event->GetTrajectoryContainer();
0466       if (!container || container->entries() == 0) {
0467         warning("[firebird-writer] Event %d: no trajectories", event->GetEventID());
0468         return;
0469       }
0470
0471       const int nTrajectories = container->entries();
0472       m_totalTrajectories += nTrajectories;
0473
0474       int nFiltered = 0, nSaved = 0;
0475
0476       // ── build event JSON ───────────────────────────────────────────
0477       std::string ev = fmt::format(
0478         R"({{"id":{},"groups":[{{"name":"{}","type":"PointTrajectory",)"
0479         R"("origin":{{"type":["G4VTrajectory","G4VTrajectoryPoint"]}},)"
0480         R"("paramColumns":["pdg","type","charge","px","py","pz",)"
0481         R"("vx","vy","vz","theta","phi","q_over_p","loc_a","loc_b","time"],)"
0482         R"("pointColumns":["x","y","z","t","aux"],)"
0483         R"("trajectories":[)",
0484         event->GetEventID(), m_componentName);
0485
0486       bool first = true;
0487
0488       for (int i = 0; i < nTrajectories; ++i) {
0489         auto* trj = (*container)[i];
0490
0491         if (!passesFilters(trj) || !passesRichTrajectoryCheck(trj)) {
0492           ++nFiltered;
0493           continue;
0494         }
0495
0496         auto points = buildPointsJson(trj);
0497         if (points == "[]") { ++nFiltered; continue; }
0498
0499         if (!first) ev += ',';
0500         first = false;
0501
0502         ev += fmt::format(R"({{"points":{},"params":{}}})", points, buildParamsJson(trj));
0503         ++nSaved;
0504       }
0505
0506       ev += "]}]}";
0507
0508       if (nSaved > 0)
0509         m_entries.push_back(std::move(ev));
0510
0511       m_filteredTrajectories += nFiltered;
0512       m_savedTrajectories    += nSaved;
0513
0514       info("[firebird-writer] Event %d: %d trajectories, %d filtered, %d saved",
0515            event->GetEventID(), nTrajectories, nFiltered, nSaved);
0516     }
0517   };
0518
0519 } // namespace dd4hep::sim
0520
0521 // Plugin registration
0522 #include "DDG4/Factories.h"
0523 DECLARE_GEANT4ACTION(FirebirdTrajectoryWriterEventAction)