io/podio/JEventSourcePODIO.cc

0001 // Copyright 2022, David Lawrence
0002 // Subject to the terms in the LICENSE file found in the top-level directory.
0003 //
0004 //
0005 // This is a JANA event source that uses PODIO to read from a ROOT
0006 // file created using the EDM4hep Data Model.
0007
0008 #include "JEventSourcePODIO.h"
0009
0010 #include <JANA/JApplication.h>
0011 #include <JANA/JEvent.h>
0012 #include <JANA/JException.h>
0013 #include <JANA/Utils/JTypeInfo.h>
0014 #include <TFile.h>
0015 #include <TObject.h>
0016 #include <fmt/core.h>
0017 #include <fmt/ostream.h>
0018 #include <podio/CollectionBase.h>
0019 #include <podio/Frame.h>
0020 #include <podio/podioVersion.h>
0021 #include <algorithm>
0022 #include <exception>
0023 #include <iostream>
0024 #include <map>
0025 #include <memory>
0026 #include <utility>
0027 #include <vector>
0028
0029 // These files are generated automatically by make_datamodel_glue.py
0030 #include "services/io/podio/datamodel_glue.h"
0031 #include "services/io/podio/datamodel_includes.h" // IWYU pragma: keep
0032 #include "services/log/Log_service.h"
0033
0034 // Formatter for podio::version::Version
0035 template <> struct fmt::formatter<podio::version::Version> : ostream_formatter {};
0036
0037 //------------------------------------------------------------------------------
0038 // InsertingVisitor
0039 //
0040 /// This datamodel visitor will insert a PODIO collection into a JEvent.
0041 /// This allows us to access the PODIO data through JEvent::Get and JEvent::GetCollection.
0042 /// This makes it transparent to downstream factories whether the data was loaded from file, or calculated.
0043 /// InsertingVisitor is called in GetEvent()
0044 ///
0045 /// \param event             JANA JEvent to copy the data objects into
0046 /// \param collection_name   name of the collection which will be used as the factory tag for these objects
0047 //------------------------------------------------------------------------------
0048 struct InsertingVisitor {
0049   JEvent& m_event;
0050   const std::string& m_collection_name;
0051
0052   InsertingVisitor(JEvent& event, const std::string& collection_name)
0053       : m_event(event), m_collection_name(collection_name){};
0054
0055   template <typename T> void operator()(const T& collection) {
0056
0057     using ContentsT = decltype(collection[0]);
0058     m_event.InsertCollectionAlreadyInFrame<ContentsT>(&collection, m_collection_name);
0059   }
0060 };
0061
0062 //------------------------------------------------------------------------------
0063 // Constructor
0064 //
0065 ///
0066 /// \param resource_name  Name of root file to open (n.b. file is not opened until Open() is called)
0067 /// \param app            JApplication
0068 //------------------------------------------------------------------------------
0069 JEventSourcePODIO::JEventSourcePODIO(std::string resource_name, JApplication* app)
0070     : JEventSource(resource_name, app) {
0071   SetTypeName(NAME_OF_THIS); // Provide JANA with class name
0072 #if JANA_NEW_CALLBACK_STYLE
0073   SetCallbackStyle(CallbackStyle::ExpertMode); // Use new, exception-free Emit() callback
0074 #endif
0075
0076   // Get Logger
0077   m_log = GetApplication()->GetService<Log_service>()->logger("JEventSourcePODIO");
0078
0079   // Tell JANA that we want it to call the FinishEvent() method.
0080   // EnableFinishEvent();
0081
0082   // Allow user to specify to recycle events forever
0083   GetApplication()->SetDefaultParameter("podio:run_forever", m_run_forever,
0084                                         "set to true to recycle through events continuously");
0085
0086   bool print_type_table = false;
0087   GetApplication()->SetDefaultParameter("podio:print_type_table", print_type_table,
0088                                         "Print list of collection names and their types");
0089
0090   // Hopefully we won't need to reimplement background event merging. Using podio frames, it looks like we would
0091   // have to do a deep copy of all data in order to insert it into the same frame, which would probably be
0092   // quite inefficient.
0093   /*
0094     std::string background_filename;
0095     GetApplication()->SetDefaultParameter(
0096             "podio:background_filename",
0097             background_filename,
0098             "Name of file containing background events to merge in (default is not to merge any background)"
0099             );
0100
0101     int num_background_events=1;
0102     GetApplication()->SetDefaultParameter(
0103             "podio:num_background_events",
0104             num_background_events,
0105             "Number of background events to add to every primary event."
0106     );
0107     */
0108 }
0109
0110 //------------------------------------------------------------------------------
0111 // Destructor
0112 //------------------------------------------------------------------------------
0113 JEventSourcePODIO::~JEventSourcePODIO() {
0114   m_log->info("Closing Event Source for {}", GetResourceName());
0115 }
0116
0117 //------------------------------------------------------------------------------
0118 // Open
0119 //
0120 /// Open the root file and read in metadata.
0121 //------------------------------------------------------------------------------
0122 void JEventSourcePODIO::Open() {
0123
0124   bool print_type_table = GetApplication()->GetParameterValue<bool>("podio:print_type_table");
0125   // std::string background_filename = GetApplication()->GetParameterValue<std::string>("podio:background_filename");;
0126   // int num_background_events = GetApplication()->GetParameterValue<int>("podio:num_background_events");;
0127
0128   // Open primary events file
0129   try {
0130
0131     m_reader.openFile(GetResourceName());
0132
0133     auto version          = m_reader.currentFileVersion();
0134     bool version_mismatch = version.major > podio::version::build_version.major;
0135     version_mismatch |= (version.major == podio::version::build_version.major) &&
0136                         (version.minor > podio::version::build_version.minor);
0137     if (version_mismatch) {
0138       std::stringstream ss;
0139       ss << "Mismatch in PODIO versions! " << version << " > " << podio::version::build_version;
0140       // FIXME: The podio ROOTReader is somehow failing to read in the correct version numbers from the file
0141       //            throw JException(ss.str());
0142     }
0143
0144     m_log->info("PODIO version: file={} (executable={})", version, podio::version::build_version);
0145
0146     Nevents_in_file = m_reader.getEntries("events");
0147     m_log->info("Opened PODIO Frame file \"{}\" with {} events", GetResourceName(),
0148                 Nevents_in_file);
0149
0150     if (print_type_table) {
0151       PrintCollectionTypeTable();
0152     }
0153
0154   } catch (std::exception& e) {
0155     m_log->error(e.what());
0156     throw JException(fmt::format("Problem opening file \"{}\"", GetResourceName()));
0157   }
0158 }
0159
0160 //------------------------------------------------------------------------------
0161 // Close
0162 //
0163 /// Cleanly close the resource when JANA is terminated via Ctrl-C or jana:nevents
0164 ///
0165 /// \param event
0166 //------------------------------------------------------------------------------
0167 void JEventSourcePODIO::Close() {
0168   // m_reader.close();
0169   // TODO: ROOTReader does not appear to have a close() method.
0170 }
0171
0172 //------------------------------------------------------------------------------
0173 // GetEvent
0174 //
0175 /// Read next event from file and copy its objects into the given JEvent.
0176 ///
0177 /// \param event
0178 //------------------------------------------------------------------------------
0179 #if JANA_NEW_CALLBACK_STYLE
0180 JEventSourcePODIO::Result JEventSourcePODIO::Emit(JEvent& event) {
0181 #else
0182 void JEventSourcePODIO::GetEvent(std::shared_ptr<JEvent> _event) {
0183   auto& event = *_event;
0184 #endif
0185
0186   /// Calls to GetEvent are synchronized with each other, which means they can
0187   /// read and write state on the JEventSource without causing race conditions.
0188
0189   // Check if we have exhausted events from file
0190   if (Nevents_read >= Nevents_in_file) {
0191     if (m_run_forever) {
0192       Nevents_read = 0;
0193     } else {
0194 #if JANA_NEW_CALLBACK_STYLE
0195       return Result::FailureFinished;
0196 #else
0197       throw RETURN_STATUS::kNO_MORE_EVENTS;
0198 #endif
0199     }
0200   }
0201
0202   auto frame_data = m_reader.readEntry("events", Nevents_read);
0203   auto frame      = std::make_unique<podio::Frame>(std::move(frame_data));
0204
0205   if (m_use_event_headers) {
0206     const auto& event_headers = frame->get<edm4hep::EventHeaderCollection>("EventHeader");
0207     if (event_headers.size() != 1) {
0208       m_log->warn("Missing or bad event headers: Entry {} contains {} items, but 1 expected. Will "
0209                   "not use event and run numbers from header",
0210                   Nevents_read, event_headers.size());
0211       m_use_event_headers = false;
0212     } else {
0213       event.SetEventNumber(event_headers[0].getEventNumber());
0214       event.SetRunNumber(event_headers[0].getRunNumber());
0215     }
0216   }
0217
0218   // Insert contents odf frame into JFactories
0219   VisitPodioCollection<InsertingVisitor> visit;
0220   for (const std::string& coll_name : frame->getAvailableCollections()) {
0221     const podio::CollectionBase* collection = frame->get(coll_name);
0222     InsertingVisitor visitor(event, coll_name);
0223     visit(visitor, *collection);
0224   }
0225
0226   event.Insert(frame.release()); // Transfer ownership from unique_ptr to JFactoryT<podio::Frame>
0227   Nevents_read += 1;
0228 #if JANA_NEW_CALLBACK_STYLE
0229   return Result::Success;
0230 #endif
0231 }
0232
0233 //------------------------------------------------------------------------------
0234 // GetDescription
0235 //------------------------------------------------------------------------------
0236 std::string JEventSourcePODIO::GetDescription() {
0237
0238   /// GetDescription() helps JANA explain to the user what is going on
0239   return "PODIO root file (Frames, podio >= v0.16.3)";
0240 }
0241
0242 //------------------------------------------------------------------------------
0243 // CheckOpenable
0244 //
0245 /// Return a value from 0-1 indicating probability that this source will be
0246 /// able to read this root file. Currently, it simply checks that the file
0247 /// name contains the string ".root" and if does, returns a small number (0.02).
0248 /// This will need to be made more sophisticated if the alternative root file
0249 /// formats need to be supported by other event sources.
0250 ///
0251 /// \param resource_name name of root file to evaluate.
0252 /// \return              value from 0-1 indicating confidence that this source can open the given file
0253 //------------------------------------------------------------------------------
0254 template <>
0255 double JEventSourceGeneratorT<JEventSourcePODIO>::CheckOpenable(std::string resource_name) {
0256
0257   // PODIO Frame reader gets slightly higher precedence than PODIO Legacy reader, but only if the file
0258   // contains a 'podio_metadata' TTree. If the file doesn't exist, this will return 0. The "file not found"
0259   // error will hopefully be generated by the PODIO legacy reader instead.
0260   if (resource_name.find(".root") == std::string::npos) {
0261     return 0.0;
0262   }
0263
0264   // PODIO FrameReader segfaults on legacy input files, so we use ROOT to validate beforehand. Of course,
0265   // we can't validate if ROOT can't read the file.
0266   std::unique_ptr<TFile> file = std::unique_ptr<TFile>{TFile::Open(resource_name.c_str())};
0267   if (!file || file->IsZombie()) {
0268     return 0.0;
0269   }
0270
0271   // We test the format the same way that PODIO's python API does. See python/podio/reading.py
0272   TObject* tree = file->Get("podio_metadata");
0273   if (tree == nullptr) {
0274     return 0.0;
0275   }
0276   return 0.03;
0277 }
0278
0279 //------------------------------------------------------------------------------
0280 // PrintCollectionTypeTable
0281 //
0282 /// Print the list of collection names from the currently open file along
0283 /// with their types. This will be called automatically when the file is
0284 /// open if the PODIO:PRINT_TYPE_TABLE variable is set to a non-zero value
0285 //------------------------------------------------------------------------------
0286 void JEventSourcePODIO::PrintCollectionTypeTable() {
0287
0288   // Read the zeroth entry. This assumes that m_reader has already been initialized with a valid filename
0289   auto frame_data = m_reader.readEntry("events", 0);
0290   auto frame      = std::make_unique<podio::Frame>(std::move(frame_data));
0291
0292   std::map<std::string, std::string> collectionNames;
0293   std::size_t max_name_len = 0;
0294   std::size_t max_type_len = 0;
0295
0296   // Record all (collection name, value type name) pairs
0297   // Record the maximum length of both strings so that we can print nicely aligned columns.
0298   for (const std::string& name : frame->getAvailableCollections()) {
0299     const podio::CollectionBase* coll = frame->get(name);
0300     const auto type                   = coll->getTypeName();
0301     max_name_len                      = std::max(max_name_len, name.length());
0302     max_type_len                      = std::max(max_type_len, type.length());
0303     collectionNames[name]             = std::string(type);
0304   }
0305
0306   // Print table
0307   std::cout << std::endl;
0308   std::cout << "Available Collections" << std::endl;
0309   std::cout << std::endl;
0310   std::cout << "Collection Name"
0311             << std::string(max_name_len + 2 - std::string("Collection Name").length(), ' ')
0312             << "Data Type" << std::endl;
0313   std::cout << std::string(max_name_len, '-') << "  " << std::string(max_name_len, '-')
0314             << std::endl;
0315   for (auto& [name, type] : collectionNames) {
0316     std::cout << name + std::string(max_name_len + 2 - name.length(), ' ');
0317     std::cout << type << std::endl;
0318   }
0319   std::cout << std::endl;
0320 }