plugins/rntuple/NanoAODRNTuples.cc

0001 #include "NanoAODRNTuples.h"
0002
0003 #include "DataFormats/NanoAOD/interface/MergeableCounterTable.h"
0004 #include "FWCore/Framework/interface/RunForOutput.h"
0005
0006 #include <ROOT/RNTuple.hxx>
0007 #include <ROOT/RNTupleModel.hxx>
0008 #include <ROOT/RPageStorageFile.hxx>
0009 using ROOT::Experimental::RNTupleModel;
0010 #if ROOT_VERSION_CODE < ROOT_VERSION(6, 31, 0)
0011 using ROOT::Experimental::RNTupleWriter;
0012 using ROOT::Experimental::Detail::RPageSinkFile;
0013 #define MakeRNTupleWriter std::make_unique<RNTupleWriter>
0014 #include <ROOT/RNTupleOptions.hxx>
0015 #else
0016 using ROOT::Experimental::Internal::RPageSinkFile;
0017 #define MakeRNTupleWriter ROOT::Experimental::Internal::CreateRNTupleWriter
0018 #include <ROOT/RNTupleWriteOptions.hxx>
0019 #endif
0020 using ROOT::Experimental::RNTupleWriteOptions;
0021
0022 #include "RNTupleFieldPtr.h"
0023 #include "SummaryTableOutputFields.h"
0024
0025 void LumiNTuple::createFields(const edm::LuminosityBlockID& id, TFile& file) {
0026   auto model = RNTupleModel::Create();
0027   m_run = RNTupleFieldPtr<UInt_t>("run", "", *model);
0028   m_luminosityBlock = RNTupleFieldPtr<UInt_t>("luminosityBlock", "", *model);
0029   // TODO use Append when we bump our RNTuple version:
0030   // m_ntuple = RNTupleWriter::Append(std::move(model), "LuminosityBlocks", file);
0031   RNTupleWriteOptions options;
0032   options.SetCompression(file.GetCompressionSettings());
0033   m_ntuple = MakeRNTupleWriter(std::move(model), std::make_unique<RPageSinkFile>("LuminosityBlocks", file, options));
0034 }
0035
0036 void LumiNTuple::fill(const edm::LuminosityBlockID& id, TFile& file) {
0037   if (!m_ntuple) {
0038     createFields(id, file);
0039   }
0040   m_run.fill(id.run());
0041   m_luminosityBlock.fill(id.value());
0042   m_ntuple->Fill();
0043 }
0044
0045 void LumiNTuple::finalizeWrite() { m_ntuple.reset(); }
0046
0047 void RunNTuple::registerToken(const edm::EDGetToken& token) { m_tokens.push_back(token); }
0048
0049 void RunNTuple::createFields(const edm::RunForOutput& iRun, TFile& file) {
0050   auto model = RNTupleModel::Create();
0051   m_run = RNTupleFieldPtr<UInt_t>("run", "", *model);
0052
0053   edm::Handle<nanoaod::MergeableCounterTable> handle;
0054   for (const auto& token : m_tokens) {
0055     iRun.getByToken(token, handle);
0056     const nanoaod::MergeableCounterTable& tab = *handle;
0057     m_tables.push_back(SummaryTableOutputFields(tab, *model));
0058   }
0059
0060   // TODO use Append when we bump our RNTuple version
0061   RNTupleWriteOptions options;
0062   options.SetCompression(file.GetCompressionSettings());
0063   m_ntuple = MakeRNTupleWriter(std::move(model), std::make_unique<RPageSinkFile>("Runs", file, options));
0064 }
0065
0066 void RunNTuple::fill(const edm::RunForOutput& iRun, TFile& file) {
0067   if (!m_ntuple) {
0068     createFields(iRun, file);
0069   }
0070   m_run.fill(iRun.id().run());
0071   edm::Handle<nanoaod::MergeableCounterTable> handle;
0072   for (std::size_t i = 0; i < m_tokens.size(); i++) {
0073     iRun.getByToken(m_tokens.at(i), handle);
0074     const nanoaod::MergeableCounterTable& tab = *handle;
0075     m_tables.at(i).fill(tab);
0076   }
0077   m_ntuple->Fill();
0078 }
0079
0080 void RunNTuple::finalizeWrite() { m_ntuple.reset(); }
0081
0082 void PSetNTuple::createFields(TFile& file) {
0083   // use a collection to emulate std::pair
0084   auto pairModel = RNTupleModel::Create();
0085   m_psetId = RNTupleFieldPtr<std::string>("first", "", *pairModel);
0086   m_psetBlob = RNTupleFieldPtr<std::string>("second", "", *pairModel);
0087   auto model = RNTupleModel::Create();
0088   m_collection = model->MakeCollection(edm::poolNames::idToParameterSetBlobsBranchName(), std::move(pairModel));
0089   // TODO use Append when we bump our RNTuple version
0090   RNTupleWriteOptions options;
0091   options.SetCompression(file.GetCompressionSettings());
0092   m_ntuple = MakeRNTupleWriter(std::move(model),
0093                                std::make_unique<RPageSinkFile>(edm::poolNames::parameterSetsTreeName(), file, options));
0094 }
0095
0096 void PSetNTuple::fill(edm::pset::Registry* pset, TFile& file) {
0097   if (!pset) {
0098     throw cms::Exception("LogicError", "null edm::pset::Registry::Instance pointer");
0099   }
0100   if (!m_ntuple) {
0101     createFields(file);
0102   }
0103   for (const auto& ps : *pset) {
0104     std::ostringstream oss;
0105     oss << ps.first;
0106     m_psetId.fill(oss.str());
0107     m_psetBlob.fill(ps.second.toString());
0108     m_collection->Fill();
0109     m_ntuple->Fill();
0110   }
0111 }
0112
0113 void PSetNTuple::finalizeWrite() { m_ntuple.reset(); }
0114
0115 // TODO blocked on RNTuple typedef member field support
0116 void MetadataNTuple::createFields(TFile& file) {
0117   auto procHistModel = RNTupleModel::Create();
0118   // ProcessHistory.transients_.phid_ replacement
0119   m_phId = RNTupleFieldPtr<std::string>("transients_phid_", "", *procHistModel);
0120   auto model = RNTupleModel::Create();
0121   m_procHist = model->MakeCollection(edm::poolNames::processHistoryBranchName(), std::move(procHistModel));
0122   RNTupleWriteOptions options;
0123   options.SetCompression(file.GetCompressionSettings());
0124   m_ntuple = MakeRNTupleWriter(std::move(model),
0125                                std::make_unique<RPageSinkFile>(edm::poolNames::metaDataTreeName(), file, options));
0126 }
0127
0128 void MetadataNTuple::fill(const edm::ProcessHistoryRegistry& procHist, TFile& file) {
0129   if (!m_ntuple) {
0130     createFields(file);
0131   }
0132   for (const auto& ph : procHist) {
0133     std::string phid;
0134     ph.second.id().toString(phid);
0135     m_phId.fill(phid);
0136     m_procHist->Fill();
0137   }
0138   m_ntuple->Fill();
0139 }
0140
0141 void MetadataNTuple::finalizeWrite() { m_ntuple.reset(); }