epic-analysis/tutorial/podio_frame_reader.C

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023 Christopher Dilks
0003
0004 // example how to read a ROOT file with `PODIO` `ROOTFrameReader`
0005
0006 R__LOAD_LIBRARY(fmt)
0007 #include "fmt/format.h"
0008
0009 R__LOAD_LIBRARY(podio)
0010 R__LOAD_LIBRARY(podioRootIO)
0011 #include "podio/ROOTFrameReader.h"
0012 #include "podio/Frame.h"
0013
0014 R__LOAD_LIBRARY(edm4hep)
0015 R__LOAD_LIBRARY(edm4eic)
0016 R__LOAD_LIBRARY(edm4hepDict)
0017 R__LOAD_LIBRARY(edm4eicDict)
0018 #include "edm4eic/ReconstructedParticleCollection.h"
0019 #include "edm4eic/MCRecoParticleAssociationCollection.h"
0020 #include "edm4hep/utils/kinematics.h"
0021
0022 // -----------------------------------------------------------------------------
0023
0024 void PrintParticle(const edm4hep::MCParticle& P);
0025 void PrintParticle(const edm4eic::ReconstructedParticle& P);
0026
0027 // -----------------------------------------------------------------------------
0028
0029 void podio_frame_reader(
0030     const std::string& root_file_name, // input ROOT file
0031     unsigned max_num_entries = 0       // maximum number of entries to read
0032     )
0033 {
0034
0035   // start ROOTFrameReader
0036   fmt::print("{:=^50}\nOpening file \"{}\"\n", "", root_file_name);
0037   auto reader = podio::ROOTFrameReader();
0038   reader.openFile(root_file_name);
0039
0040   // get number of events
0041   const std::string tree_name = "events"; // could be obtained from `reader.getAvailableCategories()`
0042   auto num_entries = reader.getEntries(tree_name);
0043   fmt::print("Number of entries in \"{}\": {}\n", tree_name, num_entries);
0044   if(max_num_entries>0) num_entries = std::min(num_entries, max_num_entries);
0045   fmt::print("Reading {} of them...\n", num_entries);
0046
0047   // event loop
0048   for(unsigned e=0; e<num_entries; e++) {
0049
0050     // next event
0051     auto frame = podio::Frame(reader.readNextEntry(tree_name));
0052
0053     // dump list of collections (from first event)
0054     if(e==0) {
0055       fmt::print("\n{:=^50}\n", " COLLECTIONS ");
0056       for(auto coll : frame.getAvailableCollections())
0057         fmt::print("  {}\n", coll);
0058       fmt::print("\n{:=^50}\n", " BEGIN EVENT LOOP ");
0059     }
0060
0061     // print event number
0062     fmt::print("\n{:=<50}\n", fmt::format("=== EVENT {} ",e));
0063
0064     // get collections
0065     std::string rec_parts_collname  = "ReconstructedChargedParticles";
0066     std::string rec_assocs_collname = "ReconstructedChargedParticlesAssociations";
0067     auto& rec_parts  = frame.get<edm4eic::ReconstructedParticleCollection>(rec_parts_collname);
0068     auto& rec_assocs = frame.get<edm4eic::MCRecoParticleAssociationCollection>(rec_assocs_collname);
0069
0070     // check for existence of these collections
0071     if( ! (rec_parts.isValid() && rec_assocs.isValid())) {
0072       fmt::print(stderr,"WARNING: missing collections for this event\n");
0073       continue;
0074     }
0075
0076     // get the number of entries in each collection
0077     fmt::print("Number of entries in collections:\n");
0078     fmt::print(" {:>50}: {}\n", rec_parts_collname,  rec_parts.size());
0079     fmt::print(" {:>50}: {}\n", rec_assocs_collname, rec_assocs.size());
0080
0081     // warn if number of reconstructed particles != number of associations
0082     // - you may not want this warning or check for this, because sometimes not
0083     //   every reconstructed particle will have an associated MC particle
0084     if(rec_parts.size() != rec_assocs.size())
0085       fmt::print(stderr,"WARNING: {} rec. particles != {} rec. particle associations\n", rec_parts.size(), rec_assocs.size());
0086
0087     // loop over reconstructed particles
0088     /*
0089     fmt::print("\n{:-^50}\n", fmt::format("Reconstructed Particles ",e));
0090     for(const auto& rec_part : rec_parts) {
0091       PrintParticle(rec_part);
0092       fmt::print("\n");
0093     }
0094     */
0095
0096     // loop over reconstructed particle associations
0097     fmt::print("\n{:-^50}\n", fmt::format("Reconstructed Particle Associations ",e));
0098     for(const auto& rec_assoc : rec_assocs) {
0099       PrintParticle(rec_assoc.getRec());
0100       PrintParticle(rec_assoc.getSim());
0101       fmt::print("\n");
0102     }
0103
0104   } // end event loop
0105
0106 }
0107
0108 // -----------------------------------------------------------------------------
0109
0110 void PrintParticle(const edm4hep::MCParticle& P) {
0111   fmt::print("=> True MC Particle\n");
0112   fmt::print("  {:>20}: {}\n",     "True PDG",     P.getPDG()             );
0113   fmt::print("  {:>20}: {}\n",     "Status",       P.getGeneratorStatus() );
0114   fmt::print("  {:>20}: {} GeV\n", "Energy",       P.getEnergy()          );
0115   fmt::print("  {:>20}: {} GeV\n", "p=|Momentum|", edm4hep::utils::p(P)   );
0116   fmt::print("  {:>20}: {} GeV\n", "pT_lab",       edm4hep::utils::pT(P)  );
0117   fmt::print("  {:>20}: ({}, {}, {}) GeV\n",
0118       "3-Momentum",
0119       P.getMomentum().x,
0120       P.getMomentum().y,
0121       P.getMomentum().z
0122       );
0123   fmt::print("  {:>20}: ({}, {}, {}) mm\n",
0124       "Vertex",
0125       P.getVertex().x,
0126       P.getVertex().y,
0127       P.getVertex().z
0128       );
0129   fmt::print("  {:>20}:\n", "Parents");
0130   if(P.parents_size()>0) {
0131     for(const auto& parent : P.getParents())
0132       fmt::print("    {:>20}: {}\n", "PDG", parent.getPDG());
0133   } else fmt::print("    {:>20}\n", "None");
0134   fmt::print("  {:>20}:\n", "Daughters");
0135   if(P.daughters_size()>0) {
0136     for(const auto& daughter : P.getDaughters())
0137       fmt::print("    {:>20}: {}\n", "PDG", daughter.getPDG());
0138   } else fmt::print("    {:>20}\n", "None");
0139 }
0140
0141 // -----------------------------------------------------------------------------
0142
0143 void PrintParticle(const edm4eic::ReconstructedParticle& P) {
0144   fmt::print("=> Reconstructed Particle:\n");
0145   fmt::print("  {:>20}: ", "ParticleIDUsed::PDG");
0146   if(P.getParticleIDUsed().isAvailable()) fmt::print("{}\n", P.getParticleIDUsed().getPDG());
0147   else fmt::print("???\n");
0148   fmt::print("  {:>20}: {} GeV\n", "Mass",         P.getMass()           );
0149   fmt::print("  {:>20}: {}\n",     "Charge",       P.getCharge()         );
0150   fmt::print("  {:>20}: {} GeV\n", "Energy",       P.getEnergy()         );
0151   fmt::print("  {:>20}: {} GeV\n", "p=|Momentum|", edm4hep::utils::p(P)  );
0152   fmt::print("  {:>20}: {} GeV\n", "pT_lab",       edm4hep::utils::pT(P) );
0153   fmt::print("  {:>20}: ({}, {}, {}) GeV\n",
0154       "3-Momentum",
0155       P.getMomentum().x,
0156       P.getMomentum().y,
0157       P.getMomentum().z
0158       );
0159   fmt::print("  {:>20}: {}\n", "# of clusters", P.clusters_size()    );
0160   fmt::print("  {:>20}: {}\n", "# of tracks",   P.tracks_size()      );
0161   fmt::print("  {:>20}: {}\n", "# of PIDs",     P.particleIDs_size() );
0162   fmt::print("  {:>20}: {}\n", "# of recParts", P.particles_size()   );
0163   // for(const auto& track : P.getTracks()) {
0164   //   // ...
0165   // }
0166   // for(const auto& cluster : P.getClusters()) {
0167   //   // ...
0168   // }
0169 }