Storing Generator Filter Information in NanoAOD

PhysicsTools/NanoAOD/interface/SimpleFlatTableProducer.h

@@ -1,4 +1,5 @@
 #include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/one/EDProducer.h"
 #include "FWCore/ParameterSet/interface/ParameterSet.h"
 #include "FWCore/Framework/interface/ConsumesCollector.h"
 #include "DataFormats/Common/interface/View.h"
@@ -271,3 +272,134 @@ class FirstObjectSimpleFlatTableProducer : public SimpleFlatTableProducerBase<T,
     return out;
   }
 };
+
+template <typename T, typename TProd>
+class SimpleFlatTableProducerBaseLumi : public edm::one::EDProducer<edm::EndLuminosityBlockProducer> {
+public:
+  SimpleFlatTableProducerBaseLumi(edm::ParameterSet const &params)
+      : name_(params.getParameter<std::string>("name")),
+        doc_(params.existsAs<std::string>("doc") ? params.getParameter<std::string>("doc") : ""),
+        extension_(params.existsAs<bool>("extension") ? params.getParameter<bool>("extension") : false),
+        skipNonExistingSrc_(
+            params.existsAs<bool>("skipNonExistingSrc") ? params.getParameter<bool>("skipNonExistingSrc") : false),
+        src_(skipNonExistingSrc_ ? mayConsume<TProd>(params.getParameter<edm::InputTag>("src"))
+                                 : consumes<TProd,edm::InLumi>(params.getParameter<edm::InputTag>("src"))) {
+    edm::ParameterSet const &varsPSet = params.getParameter<edm::ParameterSet>("variables");
+    for (const std::string &vname : varsPSet.getParameterNamesForType<edm::ParameterSet>()) {
+      const auto &varPSet = varsPSet.getParameter<edm::ParameterSet>(vname);
+      const std::string &type = varPSet.getParameter<std::string>("type");
+      if (type == "int")
+        vars_.push_back(std::make_unique<IntVar>(vname, varPSet));
+      else if (type == "float")
+        vars_.push_back(std::make_unique<FloatVar>(vname, varPSet));
+      else if (type == "uint8")
+        vars_.push_back(std::make_unique<UInt8Var>(vname, varPSet));
+      else if (type == "bool")
+        vars_.push_back(std::make_unique<BoolVar>(vname, varPSet));
+      else
+        throw cms::Exception("Configuration", "unsupported type " + type + " for variable " + vname);
+    }
+
+    produces<nanoaod::FlatTable,edm::Transition::EndLuminosityBlock>();
+  }
+
+  ~SimpleFlatTableProducerBaseLumi() override {}
+
+  // this is to be overriden by the child class
+  virtual std::unique_ptr<nanoaod::FlatTable> fillTable(const edm::LuminosityBlock &iLumi,
+                                                        const edm::Handle<TProd> &prod) const = 0;
+
+  void produce(edm::Event &iEvent, const edm::EventSetup &iSetup) override {
+    // do nothing 
+  }
+
+  void endLuminosityBlockProduce(edm::LuminosityBlock& iLumi, const edm::EventSetup &iSetup) final {
+    edm::Handle<TProd> src;
+    iLumi.getByToken(src_, src);
+
+    std::unique_ptr<nanoaod::FlatTable> out = fillTable(iLumi, src);
+    out->setDoc(doc_);
+
+    iLumi.put(std::move(out));
+  }
+
+protected:
+  const std::string name_;
+  const std::string doc_;
+  const bool extension_;
+  const bool skipNonExistingSrc_;
+  const edm::EDGetTokenT<TProd> src_;
+
+  class VariableBase {
+  public:
+    VariableBase(const std::string &aname, const edm::ParameterSet &cfg)
+        : name_(aname),
+          doc_(cfg.getParameter<std::string>("doc")),
+          precision_(cfg.existsAs<int>("precision") ? cfg.getParameter<int>("precision")
+                                                    : (cfg.existsAs<std::string>("precision") ? -2 : -1)) {}
+    virtual ~VariableBase() {}
+    const std::string &name() const { return name_; }
+
+  protected:
+    std::string name_, doc_;
+    int precision_;
+  };
+
+  class Variable : public VariableBase {
+  public:
+    Variable(const std::string &aname, const edm::ParameterSet &cfg) : VariableBase(aname, cfg) {}
+    virtual void fill(std::vector<const T *> selobjs, nanoaod::FlatTable &out) const = 0;
+  };
+
+  template <typename StringFunctor, typename ValType>
+  class FuncVariable : public Variable {
+  public:
+    FuncVariable(const std::string &aname, const edm::ParameterSet &cfg)
+        : Variable(aname, cfg),
+          func_(cfg.getParameter<std::string>("expr"), true),
+          precisionFunc_(cfg.existsAs<std::string>("precision") ? cfg.getParameter<std::string>("precision") : "23",
+                         true) {}
+    ~FuncVariable() override {}
+    void fill(std::vector<const T *> selobjs, nanoaod::FlatTable &out) const override {
+      std::vector<ValType> vals(selobjs.size());
+      for (unsigned int i = 0, n = vals.size(); i < n; ++i) {
+        if constexpr (std::is_same<ValType, float>()) {
+          if (this->precision_ == -2) {
+            vals[i] =
+                MiniFloatConverter::reduceMantissaToNbitsRounding(func_(*selobjs[i]), precisionFunc_(*selobjs[i]));
+          } else {
+            vals[i] = func_(*selobjs[i]);
+          }
+        } else {
+          vals[i] = func_(*selobjs[i]);
+        }
+      }
+      out.template addColumn<ValType>(this->name_, vals, this->doc_, this->precision_);
+    }
+
+  protected:
+    StringFunctor func_;
+    StringFunctor precisionFunc_;
+  };
+  typedef FuncVariable<StringObjectFunction<T>, int> IntVar;
+  typedef FuncVariable<StringObjectFunction<T>, float> FloatVar;
+  typedef FuncVariable<StringObjectFunction<T>, uint8_t> UInt8Var;
+  typedef FuncVariable<StringCutObjectSelector<T>, bool> BoolVar;
+  std::vector<std::unique_ptr<Variable>> vars_;
+};
+
+template <typename T>
+class LumiSingletonSimpleFlatTableProducer : public SimpleFlatTableProducerBaseLumi<T, T> {
+public:
+  LumiSingletonSimpleFlatTableProducer(edm::ParameterSet const &params) : SimpleFlatTableProducerBaseLumi<T, T>(params) {}
+
+  ~LumiSingletonSimpleFlatTableProducer() override {}
+
+  std::unique_ptr<nanoaod::FlatTable> fillTable(const edm::LuminosityBlock &, const edm::Handle<T> &prod) const override {
+    auto out = std::make_unique<nanoaod::FlatTable>(1, this->name_, true, this->extension_);
+    std::vector<const T *> selobjs(1, prod.product());
+    for (const auto &var : this->vars_)
+      var->fill(selobjs, *out);
+    return out;
+  }
+};

PhysicsTools/NanoAOD/plugins/LumiOutputBranches.cc

@@ -0,0 +1,94 @@
+#include "PhysicsTools/NanoAOD/plugins/LumiOutputBranches.h"
+
+#include <iostream>
+
+namespace {
+  std::string makeBranchName(const std::string &baseName, const std::string &leafName) {
+    return baseName.empty() ? leafName : (leafName.empty() ? baseName : baseName + "_" + leafName);
+  }
+}  // namespace
+
+void LumiOutputBranches::defineBranchesFromFirstEvent(const nanoaod::FlatTable &tab) {
+  m_baseName = tab.name();
+  for (size_t i = 0; i < tab.nColumns(); i++) {
+    const std::string &var = tab.columnName(i);
+    switch (tab.columnType(i)) {
+      case nanoaod::FlatTable::ColumnType::Float:
+        m_floatBranches.emplace_back(var, tab.columnDoc(i), "F");
+        break;
+      case nanoaod::FlatTable::ColumnType::Int:
+        m_intBranches.emplace_back(var, tab.columnDoc(i), "I");
+        break;
+      case nanoaod::FlatTable::ColumnType::UInt8:
+        m_uint8Branches.emplace_back(var, tab.columnDoc(i), "b");
+        break;
+      case nanoaod::FlatTable::ColumnType::Bool:
+        m_uint8Branches.emplace_back(var, tab.columnDoc(i), "O");
+        break;
+      default:
+        throw cms::Exception("LogicError", "Unsupported type");
+    }
+  }
+}
+
+void LumiOutputBranches::branch(TTree &tree) {
+  if (!m_singleton) {
+    if (m_extension == IsExtension) {
+      m_counterBranch = tree.FindBranch(("n" + m_baseName).c_str());
+      if (!m_counterBranch) {
+        throw cms::Exception("LogicError",
+                             "Trying to save an extension table for " + m_baseName +
+                                 " before having saved the corresponding main table\n");
+      }
+    } else {
+      if (tree.FindBranch(("n" + m_baseName).c_str()) != nullptr) {
+        throw cms::Exception("LogicError", "Trying to save multiple main tables for " + m_baseName + "\n");
+      }
+      m_counterBranch = tree.Branch(("n" + m_baseName).c_str(), &m_counter, ("n" + m_baseName + "/i").c_str());
+      m_counterBranch->SetTitle(m_doc.c_str());
+    }
+  }
+  std::string varsize = m_singleton ? "" : "[n" + m_baseName + "]";
+  for (std::vector<NamedBranchPtr> *branches : {&m_floatBranches, &m_intBranches, &m_uint8Branches}) {
+    for (auto &pair : *branches) {
+      std::string branchName = makeBranchName(m_baseName, pair.name);
+      pair.branch =
+          tree.Branch(branchName.c_str(), (void *)nullptr, (branchName + varsize + "/" + pair.rootTypeCode).c_str());
+      pair.branch->SetTitle(pair.title.c_str());
+    }
+  }
+}
+
+void LumiOutputBranches::fill(const edm::LuminosityBlockForOutput &iLumi, TTree &tree, bool extensions) {
+  if (m_extension != DontKnowYetIfMainOrExtension) {
+    if (extensions != m_extension)
+      return;  // do nothing, wait to be called with the proper flag
+  }
+
+  edm::Handle<nanoaod::FlatTable> handle;
+  iLumi.getByToken(m_token, handle);
+  const nanoaod::FlatTable &tab = *handle;
+  m_counter = tab.size();
+  m_singleton = tab.singleton();
+  if (!m_branchesBooked) {
+    m_extension = tab.extension() ? IsExtension : IsMain;
+    if (extensions != m_extension)
+      return;  // do nothing, wait to be called with the proper flag
+    defineBranchesFromFirstEvent(tab);
+    m_doc = tab.doc();
+    m_branchesBooked = true;
+    branch(tree);
+  }
+  if (!m_singleton && m_extension == IsExtension) {
+    if (m_counter != *reinterpret_cast<UInt_t *>(m_counterBranch->GetAddress())) {
+      throw cms::Exception("LogicError",
+                           "Mismatch in number of entries between extension and main table for " + tab.name());
+    }
+  }
+  for (auto &pair : m_floatBranches)
+    fillColumn<float>(pair, tab);
+  for (auto &pair : m_intBranches)
+    fillColumn<int>(pair, tab);
+  for (auto &pair : m_uint8Branches)
+    fillColumn<uint8_t>(pair, tab);
+}

PhysicsTools/NanoAOD/plugins/LumiOutputBranches.h

@@ -0,0 +1,58 @@
+#ifndef PhysicsTools_NanoAOD_LumiOutputBranches_h
+#define PhysicsTools_NanoAOD_LumiOutputBranches_h
+
+#include <string>
+#include <vector>
+#include <TTree.h>
+#include "FWCore/Framework/interface/LuminosityBlockForOutput.h"
+#include "DataFormats/NanoAOD/interface/FlatTable.h"
+#include "DataFormats/Provenance/interface/BranchDescription.h"
+#include "FWCore/Utilities/interface/EDGetToken.h"
+
+class LumiOutputBranches {
+public:
+  LumiOutputBranches(const edm::BranchDescription *desc, const edm::EDGetToken &token)
+      : m_token(token), m_extension(DontKnowYetIfMainOrExtension), m_branchesBooked(false) {
+    if (desc->className() != "nanoaod::FlatTable")
+      throw cms::Exception("Configuration", "NanoAODOutputModule can only write out nanoaod::FlatTable objects");
+  }
+
+  void defineBranchesFromFirstEvent(const nanoaod::FlatTable &tab);
+  void branch(TTree &tree);
+
+  /// Fill the current table, if extensions == table.extension().
+  /// This parameter is used so that the fill is called first for non-extensions and then for extensions
+  void fill(const edm::LuminosityBlockForOutput &iLumi, TTree &tree, bool extensions);
+
+private:
+  edm::EDGetToken m_token;
+  std::string m_baseName;
+  bool m_singleton;
+  enum { IsMain = 0, IsExtension = 1, DontKnowYetIfMainOrExtension = 2 } m_extension;
+  std::string m_doc;
+  UInt_t m_counter;
+  struct NamedBranchPtr {
+    std::string name, title, rootTypeCode;
+    TBranch *branch;
+    NamedBranchPtr(const std::string &aname,
+                   const std::string &atitle,
+                   const std::string &rootType,
+                   TBranch *branchptr = nullptr)
+        : name(aname), title(atitle), rootTypeCode(rootType), branch(branchptr) {}
+  };
+  TBranch *m_counterBranch;
+  std::vector<NamedBranchPtr> m_floatBranches;
+  std::vector<NamedBranchPtr> m_intBranches;
+  std::vector<NamedBranchPtr> m_uint8Branches;
+  bool m_branchesBooked;
+
+  template <typename T>
+  void fillColumn(NamedBranchPtr &pair, const nanoaod::FlatTable &tab) {
+    int idx = tab.columnIndex(pair.name);
+    if (idx == -1)
+      throw cms::Exception("LogicError", "Missing column in input for " + m_baseName + "_" + pair.name);
+    pair.branch->SetAddress(const_cast<T *>(&tab.columnData<T>(idx).front()));  // SetAddress should take a const * !
+  }
+};
+
+#endif

PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc

@@ -39,6 +39,7 @@
 #include "DataFormats/NanoAOD/interface/FlatTable.h"
 #include "DataFormats/NanoAOD/interface/UniqueString.h"
 #include "PhysicsTools/NanoAOD/plugins/TableOutputBranches.h"
+#include "PhysicsTools/NanoAOD/plugins/LumiOutputBranches.h"
 #include "PhysicsTools/NanoAOD/plugins/TriggerOutputBranches.h"
 #include "PhysicsTools/NanoAOD/plugins/EventStringOutputBranches.h"
 #include "PhysicsTools/NanoAOD/plugins/SummaryTableOutputBranches.h"
@@ -128,6 +129,7 @@ class NanoAODOutputModule : public edm::one::OutputModule<> {
 
   std::vector<SummaryTableOutputBranches> m_runTables;
   std::vector<SummaryTableOutputBranches> m_lumiTables;
+  std::vector<LumiOutputBranches> m_lumiTables2;
   std::vector<TableOutputBranches> m_runFlatTables;
 
   std::vector<std::pair<std::string, edm::EDGetToken>> m_nanoMetadata;
@@ -231,6 +233,11 @@ void NanoAODOutputModule::writeLuminosityBlock(edm::LuminosityBlockForOutput con
 
   for (auto& t : m_lumiTables)
     t.fill(iLumi, *m_lumiTree);
+
+  for (unsigned int extensions = 0; extensions <= 1; ++extensions) {
+    for (auto& t : m_lumiTables2)
+      t.fill(iLumi, *m_lumiTree, extensions);
+  }
 
   tbb::this_task_arena::isolate([&] { m_lumiTree->Fill(); });
 
@@ -301,6 +308,7 @@ void NanoAODOutputModule::openFile(edm::FileBlock const&) {
   m_evstrings.clear();
   m_runTables.clear();
   m_lumiTables.clear();
+  m_lumiTables2.clear();
   m_runFlatTables.clear();
   const auto& keeps = keptProducts();
   for (const auto& keep : keeps[edm::InEvent]) {
@@ -320,6 +328,8 @@ void NanoAODOutputModule::openFile(edm::FileBlock const&) {
       m_lumiTables.push_back(SummaryTableOutputBranches(keep.first, keep.second));
     else if (keep.first->className() == "nanoaod::UniqueString" && keep.first->moduleLabel() == "nanoMetadata")
       m_nanoMetadata.emplace_back(keep.first->productInstanceName(), keep.second);
+    else if (keep.first->className() == "nanoaod::FlatTable")
+      m_lumiTables2.push_back(LumiOutputBranches(keep.first, keep.second));
     else
       throw cms::Exception(
           "Configuration",

PhysicsTools/NanoAOD/plugins/SimpleFlatTableProducerPlugins.cc

@@ -6,6 +6,9 @@ typedef SimpleFlatTableProducer<reco::Candidate> SimpleCandidateFlatTableProduce
 #include "SimDataFormats/GeneratorProducts/interface/GenEventInfoProduct.h"
 typedef EventSingletonSimpleFlatTableProducer<GenEventInfoProduct> SimpleGenEventFlatTableProducer;
 
+#include "SimDataFormats/GeneratorProducts/interface/GenFilterInfo.h"
+typedef LumiSingletonSimpleFlatTableProducer<GenFilterInfo> SimpleGenFilterFlatTableProducerLumi;
+
 #include "SimDataFormats/HTXS/interface/HiggsTemplateCrossSections.h"
 typedef EventSingletonSimpleFlatTableProducer<HTXS::HiggsClassification> SimpleHTXSFlatTableProducer;
 
@@ -21,7 +24,8 @@ typedef EventSingletonSimpleFlatTableProducer<math::XYZPointF> SimpleXYZPointFla
 #include "FWCore/Framework/interface/MakerMacros.h"
 DEFINE_FWK_MODULE(SimpleCandidateFlatTableProducer);
 DEFINE_FWK_MODULE(SimpleGenEventFlatTableProducer);
+DEFINE_FWK_MODULE(SimpleGenFilterFlatTableProducerLumi);
 DEFINE_FWK_MODULE(SimpleHTXSFlatTableProducer);
 DEFINE_FWK_MODULE(SimpleProtonTrackFlatTableProducer);
 DEFINE_FWK_MODULE(SimpleLocalTrackFlatTableProducer);
-DEFINE_FWK_MODULE(SimpleXYZPointFlatTableProducer);
+DEFINE_FWK_MODULE(SimpleXYZPointFlatTableProducer);

PhysicsTools/NanoAOD/python/globals_cff.py

@@ -38,5 +38,20 @@
         ),
 )
 
+genFilterTable = cms.EDProducer("SimpleGenFilterFlatTableProducerLumi",
+        src = cms.InputTag("genFilterEfficiencyProducer"),
+        cut = cms.string(""), 
+        name= cms.string("GenFilter"),
+        doc = cms.string("Generator filter information"),
+        singleton = cms.bool(True), 
+        extension = cms.bool(False),
+    variables = cms.PSet(
+        numEventsTotal        = Var("numEventsTotal()",        int,   doc="generator filter: total number of events",  precision=6),
+        numEventsPassed       = Var("numEventsPassed()",       int,   doc="generator filter: passed number of events", precision=6),
+        filterEfficiency      = Var("filterEfficiency()",      float, doc="generator filter: efficiency",              precision=14),
+        filterEfficiencyError = Var("filterEfficiencyError()", float, doc="generator filter: efficiency error",        precision=14),
+        ),
+)
+
 globalTables = cms.Sequence(rhoTable)
-globalTablesMC = cms.Sequence(puTable+genTable)
+globalTablesMC = cms.Sequence(puTable+genTable+genFilterTable)