tests/tracking_test/TrackingTest_processor.cc

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