tests/tracking_test/TrackingTest_processor.cc

0001
0002 #include "TrackingTest_processor.h"
0003
0004 #include <JANA/JApplication.h>
0005 #include <JANA/JApplicationFwd.h>
0006 #include <JANA/JEvent.h>
0007 #include <JANA/Services/JGlobalRootLock.h>
0008 #include <JANA/Services/JParameterManager.h>
0009 #include <Math/GenVector/Cartesian3D.h>
0010 #include <Math/GenVector/PxPyPzM4D.h>
0011 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0012 #include <edm4eic/ReconstructedParticleCollection.h>
0013 #include <edm4hep/MCParticleCollection.h>
0014 #include <edm4hep/Vector3d.h>
0015 #include <edm4hep/Vector3f.h>
0016 #include <fmt/core.h>
0017 #include <fmt/format.h>
0018 #include <podio/ObjectID.h>
0019 #include <spdlog/logger.h>
0020 #include <cstddef>
0021 #include <map>
0022 #include <string>
0023 #include <utility>
0024 #include <vector>
0025
0026 #include "services/io/podio/datamodel_glue.h" // IWYU pragma: keep (templated JEvent::GetCollection<T> needs PodioTypeMap)
0027 #include "services/log/Log_service.h"
0028 #include "services/rootfile/RootFile_service.h"
0029
0030 using namespace fmt;
0031
0032 //------------------
0033 // Init
0034 //------------------
0035 void TrackingTest_processor::Init() {
0036   std::string plugin_name = ("tracking_test");
0037
0038   // Get JANA application
0039   auto* app = GetApplication();
0040
0041   // Ask service locator a file to write histograms to
0042   auto root_file_service = app->GetService<RootFile_service>();
0043
0044   // Get TDirectory for histograms root file
0045   auto globalRootLock = app->GetService<JGlobalRootLock>();
0046   globalRootLock->acquire_write_lock();
0047   auto* file = root_file_service->GetHistFile();
0048   globalRootLock->release_lock();
0049
0050   // Create a directory for this plugin. And subdirectories for series of histograms
0051   m_dir_main = file->mkdir(plugin_name.c_str());
0052
0053   // Get log level from user parameter or default
0054   m_log = app->GetService<Log_service>()->logger(plugin_name);
0055   for (auto pair : app->GetJParameterManager()->GetAllParameters()) {
0056     m_log->info("{:<20} | {}", pair.first, pair.second->GetDescription());
0057   }
0058 }
0059
0060 //------------------
0061 // Process
0062 //------------------
0063 void TrackingTest_processor::Process(const std::shared_ptr<const JEvent>& event) {
0064   m_log->debug("---- TrackingTest_processor {} ----", event->GetEventNumber());
0065
0066   ProcessTrackingMatching(event);
0067 }
0068
0069 //------------------
0070 // Finish
0071 //------------------
0072 void TrackingTest_processor::Finish() { fmt::print("OccupancyAnalysis::Finish() called\n"); }
0073
0074 void TrackingTest_processor::ProcessTrackingResults(const std::shared_ptr<const JEvent>& event) {
0075   const auto* reco_particles =
0076       event->GetCollection<edm4eic::ReconstructedParticle>("outputParticles");
0077
0078   m_log->debug("Tracking reconstructed particles N={}: ", reco_particles->size());
0079   m_log->debug("   {:<5} {:>8} {:>8} {:>8} {:>8} {:>8}", "[i]", "[px]", "[py]", "[pz]", "[P]",
0080                "[P*3]");
0081
0082   for (std::size_t i = 0; i < reco_particles->size(); i++) {
0083     auto particle = (*reco_particles)[i];
0084
0085     double px = particle.getMomentum().x;
0086     double py = particle.getMomentum().y;
0087     double pz = particle.getMomentum().z;
0088     // ROOT::Math::PxPyPzM4D p4v(px, py, pz, particle.getMass());
0089     ROOT::Math::Cartesian3D p(px, py, pz);
0090     m_log->debug("   {:<5} {:>8.2f} {:>8.2f} {:>8.2f} {:>8.2f} {:>8.2f}", i, px, py, pz, p.R(),
0091                  p.R() * 3);
0092   }
0093
0094   auto mc_particles = event->Get<edm4hep::MCParticle>("MCParticles");
0095
0096   m_log->debug("MC particles N={}: ", mc_particles.size());
0097   m_log->debug("   {:<5} {:<6} {:<7} {:>8} {:>8} {:>8} {:>8}", "[i]", "status", "[PDG]", "[px]",
0098                "[py]", "[pz]", "[P]");
0099   for (std::size_t i = 0; i < mc_particles.size(); i++) {
0100
0101     const auto* particle = mc_particles[i];
0102
0103     if (particle->getGeneratorStatus() != 1) {
0104       continue;
0105     }
0106     //
0107     double px = particle->getMomentum().x;
0108     double py = particle->getMomentum().y;
0109     double pz = particle->getMomentum().z;
0110     ROOT::Math::PxPyPzM4D p4v(px, py, pz, particle->getMass());
0111     ROOT::Math::Cartesian3D p(px, py, pz);
0112     if (p.R() < 1) {
0113       continue;
0114     }
0115
0116     m_log->debug("   {:<5} {:<6} {:<7} {:>8.2f} {:>8.2f} {:>8.2f} {:>8.2f}", i,
0117                  particle->getGeneratorStatus(), particle->getPDG(), px, py, pz, p.R());
0118   }
0119 }
0120
0121 void TrackingTest_processor::ProcessTrackingMatching(const std::shared_ptr<const JEvent>& event) {
0122   m_log->debug("Associations [simId] [recID] [simE] [recE] [simPDG] [recPDG]");
0123
0124   const auto* associations = event->GetCollection<edm4eic::MCRecoParticleAssociation>(
0125       "ReconstructedChargedParticleAssociations");
0126
0127   for (auto assoc : *associations) {
0128     auto sim = assoc.getSim();
0129     auto rec = assoc.getRec();
0130
0131     m_log->debug("  {:<6} {:<6} {:>8d} {:>8d}", assoc.getSim().getObjectID().index,
0132                  assoc.getRec().getObjectID().index, sim.getPDG(), rec.getPDG());
0133   }
0134
0135   //    m_log->debug("Particles [objID] [PDG] [simE] [recE] [simPDG] [recPDG]");
0136   //    auto prt_with_assoc = event->GetSingle<edm4hep::ReconstructedParticle>("ChargedParticlesWithAssociations");
0137   //    for(auto part: prt_with_assoc->particles()) {
0138   //
0139   //        // auto sim = assoc->getSim();
0140   //        // auto rec = assoc->getRec();
0141   //
0142   //        m_log->debug("  {:<6} {:<6}  {:>8.2f} {:>8.2f}", part->getObjectID().index, part->getPDG(), part->getCharge(), part->getEnergy());
0143   //
0144   //    }
0145   //
0146   //
0147   //    m_log->debug("ReconstructedChargedParticles [objID] [PDG] [charge] [energy]");
0148   //    auto reco_charged_particles = event->Get<edm4eic::ReconstructedParticle>("ReconstructedChargedParticles");
0149   //    for(auto part: reco_charged_particles) {
0150   //        m_log->debug("  {:<6} {:<6}  {:>8.2f} {:>8.2f}", part->getObjectID().index, part->getPDG(), part->getCharge(), part->getEnergy());
0151   //    }
0152 }
0153
0154 void TrackingTest_processor::ProcessGloablMatching(const std::shared_ptr<const JEvent>& event) {
0155
0156   m_log->debug("ReconstructedParticles (FINAL) [objID] [PDG] [charge] [energy]");
0157   auto final_reco_particles =
0158       event->Get<edm4eic::ReconstructedParticle>("ReconstructedParticlesWithAssoc");
0159   for (const auto* part : final_reco_particles) {
0160     m_log->debug("  {:<6} {:<6}  {:>8.2f} {:>8.2f}", part->getObjectID().index, part->getPDG(),
0161                  part->getCharge(), part->getEnergy());
0162   }
0163 }