NanoAOD/plugins/TableOutputBranches.cc

0001 #include "PhysicsTools/NanoAOD/plugins/TableOutputBranches.h"
0002
0003 #include <iostream>
0004 #include <limits>
0005
0006 namespace {
0007   std::string makeBranchName(const std::string &baseName, const std::string &leafName) {
0008     return baseName.empty() ? leafName : (leafName.empty() ? baseName : baseName + "_" + leafName);
0009   }
0010 }  // namespace
0011
0012 void TableOutputBranches::defineBranchesFromFirstEvent(const nanoaod::FlatTable &tab) {
0013   m_baseName = tab.name();
0014   for (size_t i = 0; i < tab.nColumns(); i++) {
0015     const std::string &var = tab.columnName(i);
0016     switch (tab.columnType(i)) {
0017       case nanoaod::FlatTable::ColumnType::UInt8:
0018         m_uint8Branches.emplace_back(var, tab.columnDoc(i), "b");
0019         break;
0020       case nanoaod::FlatTable::ColumnType::Int16:
0021         m_int16Branches.emplace_back(var, tab.columnDoc(i), "S");
0022         break;
0023       case nanoaod::FlatTable::ColumnType::UInt16:
0024         m_uint16Branches.emplace_back(var, tab.columnDoc(i), "s");
0025         break;
0026       case nanoaod::FlatTable::ColumnType::Int32:
0027         m_int32Branches.emplace_back(var, tab.columnDoc(i), "I");
0028         break;
0029       case nanoaod::FlatTable::ColumnType::UInt32:
0030         m_uint32Branches.emplace_back(var, tab.columnDoc(i), "i");
0031         break;
0032       case nanoaod::FlatTable::ColumnType::Int64:
0033         m_int64Branches.emplace_back(var, tab.columnDoc(i), "L");
0034         break;
0035       case nanoaod::FlatTable::ColumnType::UInt64:
0036         m_uint64Branches.emplace_back(var, tab.columnDoc(i), "l");
0037         break;
0038       case nanoaod::FlatTable::ColumnType::Bool:
0039         m_uint8Branches.emplace_back(var, tab.columnDoc(i), "O");
0040         break;
0041       case nanoaod::FlatTable::ColumnType::Float:
0042         m_floatBranches.emplace_back(var, tab.columnDoc(i), "F");
0043         break;
0044       case nanoaod::FlatTable::ColumnType::Double:
0045         m_doubleBranches.emplace_back(var, tab.columnDoc(i), "D");
0046         break;
0047       default:
0048         throw cms::Exception("LogicError", "Unsupported type");
0049     }
0050   }
0051 }
0052
0053 void TableOutputBranches::branch(TTree &tree) {
0054   if (!m_singleton) {
0055     if (m_extension == IsExtension) {
0056       m_counterBranch = tree.FindBranch(("n" + m_baseName).c_str());
0057       if (!m_counterBranch) {
0058         throw cms::Exception("LogicError",
0059                              "Trying to save an extension table for " + m_baseName +
0060                                  " before having saved the corresponding main table\n");
0061       }
0062     } else {
0063       if (tree.FindBranch(("n" + m_baseName).c_str()) != nullptr) {
0064         throw cms::Exception("LogicError", "Trying to save multiple main tables for " + m_baseName + "\n");
0065       }
0066       m_counterBranch = tree.Branch(("n" + m_baseName).c_str(), &m_counter, ("n" + m_baseName + "/I").c_str());
0067       m_counterBranch->SetTitle(m_doc.c_str());
0068     }
0069   }
0070   std::string varsize = m_singleton ? "" : "[n" + m_baseName + "]";
0071   for (std::vector<NamedBranchPtr> *branches : {&m_uint8Branches,
0072                                                 &m_int16Branches,
0073                                                 &m_uint16Branches,
0074                                                 &m_int32Branches,
0075                                                 &m_uint32Branches,
0076                                                 &m_int64Branches,
0077                                                 &m_uint64Branches,
0078                                                 &m_floatBranches,
0079                                                 &m_doubleBranches}) {
0080     for (auto &pair : *branches) {
0081       std::string branchName = makeBranchName(m_baseName, pair.name);
0082       pair.branch =
0083           tree.Branch(branchName.c_str(), (void *)nullptr, (branchName + varsize + "/" + pair.rootTypeCode).c_str());
0084       pair.branch->SetTitle(pair.title.c_str());
0085     }
0086   }
0087 }
0088
0089 void TableOutputBranches::fill(const edm::OccurrenceForOutput &iWhatever, TTree &tree, bool extensions) {
0090   if (m_extension != DontKnowYetIfMainOrExtension) {
0091     if (extensions != m_extension)
0092       return;  // do nothing, wait to be called with the proper flag
0093   }
0094
0095   edm::Handle<nanoaod::FlatTable> handle;
0096   iWhatever.getByToken(m_token, handle);
0097   const nanoaod::FlatTable &tab = *handle;
0098   auto size = tab.size();
0099   // ROOT native array size branches may only be signed integers,
0100   // until this is changed we need to make sure the vector sizes do not exceed that
0101   if (size > std::numeric_limits<CounterType>::max()) {
0102     throw cms::Exception("Table " + tab.name() + " size is " + std::to_string(size) +
0103                          ", is too large for ROOT native array branch");
0104   }
0105   m_counter = size;
0106   m_singleton = tab.singleton();
0107   if (!m_branchesBooked) {
0108     m_extension = tab.extension() ? IsExtension : IsMain;
0109     if (extensions != m_extension)
0110       return;  // do nothing, wait to be called with the proper flag
0111     defineBranchesFromFirstEvent(tab);
0112     m_doc = tab.doc();
0113     m_branchesBooked = true;
0114     branch(tree);
0115   }
0116   if (!m_singleton && m_extension == IsExtension) {
0117     if (m_counter != *reinterpret_cast<CounterType *>(m_counterBranch->GetAddress())) {
0118       throw cms::Exception("LogicError",
0119                            "Mismatch in number of entries between extension and main table for " + tab.name());
0120     }
0121   }
0122   for (auto &pair : m_uint8Branches)
0123     fillColumn<uint8_t>(pair, tab);
0124   for (auto &pair : m_int16Branches)
0125     fillColumn<int16_t>(pair, tab);
0126   for (auto &pair : m_uint16Branches)
0127     fillColumn<uint16_t>(pair, tab);
0128   for (auto &pair : m_int32Branches)
0129     fillColumn<int32_t>(pair, tab);
0130   for (auto &pair : m_uint32Branches)
0131     fillColumn<uint32_t>(pair, tab);
0132   for (auto &pair : m_int64Branches)
0133     fillColumn<int64_t>(pair, tab);
0134   for (auto &pair : m_uint64Branches)
0135     fillColumn<uint64_t>(pair, tab);
0136   for (auto &pair : m_floatBranches)
0137     fillColumn<float>(pair, tab);
0138   for (auto &pair : m_doubleBranches)
0139     fillColumn<double>(pair, tab);
0140 }