NanoAOD/plugins/TableOutputBranches.cc

0001 #include "PhysicsTools/NanoAOD/plugins/TableOutputBranches.h"
0002
0003 #include <iostream>
0004
0005 namespace {
0006   std::string makeBranchName(const std::string &baseName, const std::string &leafName) {
0007     return baseName.empty() ? leafName : (leafName.empty() ? baseName : baseName + "_" + leafName);
0008   }
0009 }  // namespace
0010
0011 void TableOutputBranches::defineBranchesFromFirstEvent(const nanoaod::FlatTable &tab) {
0012   m_baseName = tab.name();
0013   for (size_t i = 0; i < tab.nColumns(); i++) {
0014     const std::string &var = tab.columnName(i);
0015     switch (tab.columnType(i)) {
0016       case nanoaod::FlatTable::ColumnType::Float:
0017         m_floatBranches.emplace_back(var, tab.columnDoc(i), "F");
0018         break;
0019       case nanoaod::FlatTable::ColumnType::Int:
0020         m_intBranches.emplace_back(var, tab.columnDoc(i), "I");
0021         break;
0022       case nanoaod::FlatTable::ColumnType::Int8:
0023         m_int8Branches.emplace_back(var, tab.columnDoc(i), "B");
0024         break;
0025       case nanoaod::FlatTable::ColumnType::UInt8:
0026         m_uint8Branches.emplace_back(var, tab.columnDoc(i), "b");
0027         break;
0028       case nanoaod::FlatTable::ColumnType::Bool:
0029         m_uint8Branches.emplace_back(var, tab.columnDoc(i), "O");
0030         break;
0031       case nanoaod::FlatTable::ColumnType::UInt32:
0032         m_uint32Branches.emplace_back(var, tab.columnDoc(i), "i");
0033         break;
0034       case nanoaod::FlatTable::ColumnType::Double:
0035         m_doubleBranches.emplace_back(var, tab.columnDoc(i), "D");
0036         break;
0037       default:
0038         throw cms::Exception("LogicError", "Unsupported type");
0039     }
0040   }
0041 }
0042
0043 void TableOutputBranches::branch(TTree &tree) {
0044   if (!m_singleton) {
0045     if (m_extension == IsExtension) {
0046       m_counterBranch = tree.FindBranch(("n" + m_baseName).c_str());
0047       if (!m_counterBranch) {
0048         throw cms::Exception("LogicError",
0049                              "Trying to save an extension table for " + m_baseName +
0050                                  " before having saved the corresponding main table\n");
0051       }
0052     } else {
0053       if (tree.FindBranch(("n" + m_baseName).c_str()) != nullptr) {
0054         throw cms::Exception("LogicError", "Trying to save multiple main tables for " + m_baseName + "\n");
0055       }
0056       m_counterBranch = tree.Branch(("n" + m_baseName).c_str(), &m_counter, ("n" + m_baseName + "/i").c_str());
0057       m_counterBranch->SetTitle(m_doc.c_str());
0058     }
0059   }
0060   std::string varsize = m_singleton ? "" : "[n" + m_baseName + "]";
0061   for (std::vector<NamedBranchPtr> *branches :
0062        {&m_floatBranches, &m_intBranches, &m_int8Branches, &m_uint8Branches, &m_uint32Branches, &m_doubleBranches}) {
0063     for (auto &pair : *branches) {
0064       std::string branchName = makeBranchName(m_baseName, pair.name);
0065       pair.branch =
0066           tree.Branch(branchName.c_str(), (void *)nullptr, (branchName + varsize + "/" + pair.rootTypeCode).c_str());
0067       pair.branch->SetTitle(pair.title.c_str());
0068     }
0069   }
0070 }
0071
0072 void TableOutputBranches::fill(const edm::OccurrenceForOutput &iWhatever, TTree &tree, bool extensions) {
0073   if (m_extension != DontKnowYetIfMainOrExtension) {
0074     if (extensions != m_extension)
0075       return;  // do nothing, wait to be called with the proper flag
0076   }
0077
0078   edm::Handle<nanoaod::FlatTable> handle;
0079   iWhatever.getByToken(m_token, handle);
0080   const nanoaod::FlatTable &tab = *handle;
0081   m_counter = tab.size();
0082   m_singleton = tab.singleton();
0083   if (!m_branchesBooked) {
0084     m_extension = tab.extension() ? IsExtension : IsMain;
0085     if (extensions != m_extension)
0086       return;  // do nothing, wait to be called with the proper flag
0087     defineBranchesFromFirstEvent(tab);
0088     m_doc = tab.doc();
0089     m_branchesBooked = true;
0090     branch(tree);
0091   }
0092   if (!m_singleton && m_extension == IsExtension) {
0093     if (m_counter != *reinterpret_cast<UInt_t *>(m_counterBranch->GetAddress())) {
0094       throw cms::Exception("LogicError",
0095                            "Mismatch in number of entries between extension and main table for " + tab.name());
0096     }
0097   }
0098   for (auto &pair : m_floatBranches)
0099     fillColumn<float>(pair, tab);
0100   for (auto &pair : m_intBranches)
0101     fillColumn<int>(pair, tab);
0102   for (auto &pair : m_int8Branches)
0103     fillColumn<int8_t>(pair, tab);
0104   for (auto &pair : m_uint8Branches)
0105     fillColumn<uint8_t>(pair, tab);
0106   for (auto &pair : m_uint32Branches)
0107     fillColumn<uint32_t>(pair, tab);
0108   for (auto &pair : m_doubleBranches)
0109     fillColumn<double>(pair, tab);
0110 }