Update weighttype, read storage info from config

GeneratorInterface/Core/plugins/GenWeightsTestAnalyzer.cc

@@ -422,7 +422,7 @@ GenWeightsTestAnalyzer::beginLuminosityBlock(edm::LuminosityBlock const& iLumi,
     iLumi.getByToken(lheWeightInfoToken_, lheWeightInfoHandle);
 
     // Should add a search by name function
-    auto allScaleWeights = lheWeightInfoHandle->weightGroupIndicesByType(gen::kScaleWeights);
+    auto allScaleWeights = lheWeightInfoHandle->weightGroupIndicesByType(gen::WeightType::kScaleWeights);
     if (allScaleWeights.size() > 0)
         scaleWeightsIndex_ = allScaleWeights.front();
 
@@ -441,7 +441,7 @@ GenWeightsTestAnalyzer::beginLuminosityBlock(edm::LuminosityBlock const& iLumi,
     scaleWeightsOrder_.push_back(scaleWeights->muR2muF1Index());
     scaleWeightsOrder_.push_back(scaleWeights->muR2muF2Index());
 
-    auto pdfGroups = lheWeightInfoHandle->weightGroupsByType(gen::kPdfWeights);
+    auto pdfGroups = lheWeightInfoHandle->weightGroupsByType(gen::WeightType::kPdfWeights);
     auto ct14Set = std::find_if(pdfGroups.begin(), pdfGroups.end(), 
             [] (gen::WeightGroupInfo* group) { 
                     auto pdfGroup = dynamic_cast<gen::PdfWeightGroupInfo*>(group); 

GeneratorInterface/Core/src/LHEWeightHelper.cc

@@ -155,7 +155,7 @@ namespace gen {
                     weightGroups_.push_back(std::make_unique<gen::UnknownWeightGroupInfo>("Unknown"));
                 }
                 auto& group = weightGroups_.back();
-                if (group.weightType() == gen::kScaleWeights) {
+                if (group.weightType() == gen::WeightType::kScaleWeights) {
                     if (currWeightAttributeMap_["mur"].empty()  || currWeightAttributeMap_["muf"].empty()) {
                         group.setIsWellFormed(false);
                     }

GeneratorInterface/Core/src/WeightHelper.cc

@@ -53,7 +53,7 @@ namespace gen {
     void WeightHelper::setGroupInfo() {
         auto& group = weightGroups_.back();
         const std::string& name = group.name();
-        if (group.weightType() == kPdfWeights) {
+        if (group.weightType() == WeightType::kPdfWeights) {
             PdfWeightGroupInfo* pdfGroup = dynamic_cast<PdfWeightGroupInfo*>(&group);
             auto pdfInfo = std::find_if(pdfSetsInfo.begin(), pdfSetsInfo.end(),
                 [name] (const PdfSetInfo& setInfo) { return setInfo.name == name; });

GeneratorInterface/LHEInterface/interface/TestWeightInfo.h

@@ -25,7 +25,7 @@ gen::WeightGroupInfo getExampleScaleWeights() {
         R"(<weight MUF="2" MUR="0.5" PDF="263400" id="8"> mur=0.5 muf=2 </weight>)",
         R"(<weight MUF="0.5" MUR="0.5" PDF="263400" id="9"> mur=0.5 muf=0.5 </weight>)",
     };
-    scaleInfo.setWeightType(gen::kScaleWeights);
+    scaleInfo.setWeightType(gen::WeightType::kScaleWeights);
 
     for (size_t i = 0; i < entries.size(); i++) {
         scaleInfo.addContainedId(i, parseId(entries[i]), entries[i]); 
@@ -49,7 +49,7 @@ gen::WeightGroupInfo getExampleScaleWeightsOutOfOrder() {
         R"(<weight MUF="1" MUR="0.5" PDF="263400" id="7"> mur=0.5 muf=1 </weight>)",
         R"(<weight MUF="0.5" MUR="0.5" PDF="263400" id="9"> mur=0.5 muf=0.5 </weight>)",
     };
-    scaleInfo.setWeightType(gen::kScaleWeights);
+    scaleInfo.setWeightType(gen::WeightType::kScaleWeights);
 
     for (size_t i = 0; i < entries.size(); i++) {
         scaleInfo.addContainedId(i, parseId(entries[i]), entries[i]); 

GeneratorInterface/LHEInterface/plugins/LHESource.cc

@@ -128,12 +128,12 @@ void LHESource::putWeightInfoProduct(edm::RunPrincipal& iRunPrincipal) {
     gen::WeightGroupInfo scaleInfo(
         "<weightgroup name=\"Central scale variation\" combine=\"envelope\">"
     );
-    scaleInfo.setWeightType(gen::kScaleWeights);
+    scaleInfo.setWeightType(gen::WeightType::kScaleWeights);
 
     gen::WeightGroupInfo cenPdfInfo(
         "<weightgroup name=\"NNPDF31_nnlo_hessian_pdfas\" combine=\"hessian\">"
     );
-    cenPdfInfo.setWeightType(gen::kPdfWeights);
+    cenPdfInfo.setWeightType(gen::WeightType::kPdfWeights);
 
     product->addWeightGroupInfo(&scaleInfo);
     product->addWeightGroupInfo(&cenPdfInfo);

PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc

@@ -13,6 +13,7 @@
 #include "SimDataFormats/GeneratorProducts/interface/WeightGroupInfo.h"
 #include "SimDataFormats/GeneratorProducts/interface/ScaleWeightGroupInfo.h"
 #include "SimDataFormats/GeneratorProducts/interface/PdfWeightGroupInfo.h"
+#include "FWCore/Utilities/interface/transform.h"
 
 
 #include <optional>
@@ -21,25 +22,23 @@
 
 namespace {
   typedef std::pair<std::vector<gen::WeightGroupData>, std::vector<gen::WeightGroupData>> WeightGroupsToStore;
-  //struct WeightGroupToStore {
-  //    gen::WeightType type;
-  //    std::string name;
-  //    bool inLHE;
-  //    gen::WeightGroupData groupInfo;
-  //}
 }  // namespace
 
 class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBlockCache<WeightGroupsToStore>> {
 public:
-  LHEWeightsTableProducer(edm::ParameterSet const& params)
-      : lheInputTag_(params.getParameter<edm::InputTag>("lheInfo")),
+  LHEWeightsTableProducer(edm::ParameterSet const& params) :
         lheToken_(consumes<LHEEventProduct>(params.getParameter<edm::InputTag>("lheInfo"))),
         lheWeightToken_(consumes<GenWeightProduct>(params.getParameter<edm::InputTag>("lheWeights"))),
         lheWeightInfoToken_(consumes<GenWeightInfoProduct, edm::InLumi>(params.getParameter<edm::InputTag>("lheWeights"))),
         genWeightToken_(consumes<GenWeightProduct>(params.getParameter<edm::InputTag>("genWeights"))),
         genWeightInfoToken_(consumes<GenWeightInfoProduct, edm::InLumi>(params.getParameter<edm::InputTag>("genWeights"))),
-        //weightgroups_(params.getParameter<std::vector<int>>("weightgroups")),
+        weightgroups_(edm::vector_transform(params.getParameter<std::vector<std::string>>("weightgroups"),
+                    [](auto& c) { return gen::WeightType(c.at(0)); } )),
+        maxGroupsPerType_(params.getParameter<std::vector<int>>("maxGroupsPerType")),
+        pdfIds_(params.getParameter<std::vector<int>>("pdfIds")),
         lheWeightPrecision_(params.getParameter<int32_t>("lheWeightPrecision")) {
+    if (weightgroups_.size() != maxGroupsPerType_.size())
+       throw std::invalid_argument("Inputs 'weightgroups' and 'weightgroupNums' must have equal size");
     produces<std::vector<nanoaod::FlatTable>>();
   }
 
@@ -61,35 +60,39 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
     double w0 = lheInfo.originalXWGTUP();
     auto const& weightInfos = *luminosityBlockCache(iEvent.getLuminosityBlock().index());
 
-    addWeightGroupToTable(lheWeightTables, "Lhe", weightInfos.first, lheWeights, w0);
-    addWeightGroupToTable(lheWeightTables, "Gen", weightInfos.second, genWeights, w0);
+    addWeightGroupToTable(lheWeightTables, "LHE", weightInfos.first, lheWeights, w0);
+    addWeightGroupToTable(lheWeightTables, "GEN", weightInfos.second, genWeights, w0);
 
     iEvent.put(std::move(lheWeightTables));
   }
 
-  void addWeightGroupToTable(std::unique_ptr<std::vector<nanoaod::FlatTable>>& lheWeightTables, std::string entryName,
+  void addWeightGroupToTable(std::unique_ptr<std::vector<nanoaod::FlatTable>>& lheWeightTables, const char* typeName,
           const std::vector<gen::WeightGroupData>& weightInfos, WeightsContainer& lheWeights, double w0) const {
     size_t typeCount = 0;
-    gen::WeightType currentGroup = gen::kUnknownWeights;
+    gen::WeightType previousType = gen::WeightType::kUnknownWeights;
 
+    size_t index = 0;
     for (const auto& groupInfo : weightInfos) {
+        std::string entryName = typeName;
         gen::WeightType weightType = groupInfo.group->weightType();
+        if (previousType != weightType)
+            typeCount = 0;
         std::string name = weightTypeNames_.at(weightType);
-        auto weights = weightType != gen::kScaleWeights ? normalizedWeights(lheWeights, groupInfo, w0) :
+        auto weights = weightType != gen::WeightType::kScaleWeights ? normalizedWeights(lheWeights, groupInfo, w0) :
                             orderedScaleWeights(lheWeights, groupInfo, w0);
-        if (typeCount > 0)
-            entryName.append(std::to_string(typeCount));
         entryName.append(name);
         entryName.append("Weight");
+        if (typeCount > 0) {
+            entryName.append("AltSet");
+            entryName.append(std::to_string(typeCount));
+        }
 
         lheWeightTables->emplace_back(weights.size(), entryName, false);
         lheWeightTables->back().addColumn<float>(
             "", weights, groupInfo.group->name(), nanoaod::FlatTable::FloatColumn, lheWeightPrecision_);
 
         typeCount++;
-        if (currentGroup != weightType) 
-            typeCount = 0;
-        currentGroup = weightType;
+        previousType = weightType;
     }
   }
 
@@ -105,26 +108,48 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
     edm::Handle<GenWeightInfoProduct> genWeightInfoHandle;
     iLumi.getByToken(genWeightInfoToken_, genWeightInfoHandle);
 
-    auto scaleGroups = lheWeightInfoHandle->weightGroupsAndIndicesByType(gen::kScaleWeights);
-    auto meGroups = lheWeightInfoHandle->weightGroupsAndIndicesByType(gen::kMEParamWeights);
+    std::unordered_map<gen::WeightType, int> storePerType;
+    for (size_t i = 0; i < weightgroups_.size(); i++)
+        storePerType[weightgroups_.at(i)] = maxGroupsPerType_.at(i);
 
     WeightGroupsToStore weightsToStore;
-    weightsToStore.first.insert(weightsToStore.first.end(), scaleGroups.begin(), scaleGroups.end());
-    weightsToStore.first.insert(weightsToStore.first.end(), meGroups.begin(), meGroups.end());
+    for (auto weightType : gen::allGenWeightTypes) {
+        auto lheWeights = weightDataPerType(lheWeightInfoHandle, weightType, storePerType[weightType]);
+        weightsToStore.first.insert(weightsToStore.first.end(), lheWeights.begin(), lheWeights.end());
 
-    for (auto lhaid : {306000, 91400, 260000}) {
-        if (auto pdfGroup = lheWeightInfoHandle->pdfGroupWithIndexByLHAID(lhaid)) {
-            weightsToStore.first.push_back(pdfGroup.value());
-            break;
-        }
+        auto genWeights = weightDataPerType(genWeightInfoHandle, weightType, storePerType[weightType]);
+        weightsToStore.second.insert(weightsToStore.second.end(), genWeights.begin(), genWeights.end());
     }
 
-    auto psGroups = genWeightInfoHandle->weightGroupsAndIndicesByType(gen::kPartonShowerWeights);
-    weightsToStore.second.insert(weightsToStore.second.end(), psGroups.begin(), psGroups.end());
-
     return std::make_shared<WeightGroupsToStore>(weightsToStore);
   }
 
+  std::vector<gen::WeightGroupData> weightDataPerType(edm::Handle<GenWeightInfoProduct>& weightsHandle, 
+                                                gen::WeightType weightType, int& maxStore) const {
+      std::vector<gen::WeightGroupData> group;
+      if (weightType == gen::WeightType::kPdfWeights) {
+        for (auto lhaid : pdfIds_) {
+            if (auto pdfGroup = weightsHandle->pdfGroupWithIndexByLHAID(lhaid)) {
+                group.push_back(pdfGroup.value());
+                maxStore--;
+                if (maxStore == 0)
+                    break;
+            }
+        }
+        return group;
+      }
+
+      group = weightsHandle->weightGroupsAndIndicesByType(weightType);
+
+      if (maxStore < 0 || static_cast<int>(group.size()) < maxStore) {
+          // Modify size in case one type of weight is present in multiple products
+          maxStore -= group.size();
+          return group;
+      }
+      return std::vector(group.begin(), group.begin()+maxStore);
+  }
+
+
   std::vector<float> normalizedWeights(WeightsContainer& lheWeights, const gen::WeightGroupData& meGroupInfo, double w0) const {
     std::vector<float> normalizedWeights;
     for (const auto& weight : lheWeights.at(meGroupInfo.index))
@@ -161,26 +186,30 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
     edm::ParameterSetDescription desc;
     desc.add<edm::InputTag>("lheInfo", {"externalLHEProducer"})
         ->setComment("tag(s) for the LHE information (LHEEventProduct and LHERunInfoProduct)");
-    //desc.add<std::vector<int>>("weightgroups");
     desc.add<edm::InputTag>("lheWeights");
     desc.add<edm::InputTag>("genWeights");
+    desc.add<std::vector<std::string>>("weightgroups");
+    desc.add<std::vector<int>>("maxGroupsPerType");
+    desc.add<std::vector<int>>("pdfIds");
     desc.add<int32_t>("lheWeightPrecision", -1)->setComment("Number of bits in the mantissa for LHE weights");
     descriptions.addDefault(desc);
   }
 
 protected:
-  const edm::InputTag lheInputTag_;
   const edm::EDGetTokenT<LHEEventProduct> lheToken_;
   const edm::EDGetTokenT<GenWeightProduct> lheWeightToken_;
   const edm::EDGetTokenT<GenWeightInfoProduct> lheWeightInfoToken_;
   const edm::EDGetTokenT<GenWeightProduct> genWeightToken_;
   const edm::EDGetTokenT<GenWeightInfoProduct> genWeightInfoToken_;
-  //const std::vector<gen::WeightType> weightgroups_;
+  const std::vector<gen::WeightType> weightgroups_;
+  const std::vector<int> maxGroupsPerType_;
+  const std::vector<int> pdfIds_;
   const std::unordered_map<gen::WeightType, std::string> weightTypeNames_ = {
-      {gen::kScaleWeights, "Scale"},
-      {gen::kPdfWeights, "Pdf"},
-      {gen::kMEParamWeights, "MEParam"},
-      {gen::kUnknownWeights, "Unknown"},
+      {gen::WeightType::kScaleWeights, "Scale"},
+      {gen::WeightType::kPdfWeights, "Pdf"},
+      {gen::WeightType::kMEParamWeights, "MEParam"},
+      {gen::WeightType::kPartonShowerWeights, "PartonShower"},
+      {gen::WeightType::kUnknownWeights, "Unknown"},
   };
 
   //std::unordered_map<std::string, int> weightGroupIndices_;

PhysicsTools/NanoAOD/python/nanogen_cff.py

@@ -5,7 +5,6 @@
 from PhysicsTools.NanoAOD.particlelevel_cff import *
 from PhysicsTools.NanoAOD.lheInfoTable_cfi import *
 from PhysicsTools.NanoAOD.genWeightsTable_cfi import *
-import ROOT
 
 genWeights = cms.EDProducer("GenWeightProductProducer")
 
@@ -14,10 +13,12 @@
     lheInfo = cms.InputTag("externalLHEProducer"),
     lheWeights = cms.InputTag("externalLHEProducer"),
     genWeights = cms.InputTag("genWeights"),
-    #weightgroups = cms.vint32([ROOT.gen.kScaleWeights, ROOT.gen.kMEParamWeights, ROOT.gen.kPdfWeights, ROOT.UnknownWeights]),
-    #weightgroups = cms.vint32([0,1,2,3,4]),
-    #numWeightgroups = cms.vint([1, -1, 1, 2, 1]),
-    #pdfs = cms.vint([91400, 306000, 260000]),
+    # Warning: you can use a full string, but only the first character is read.
+    # Note also that the capitalization is important! For example, 'parton shower' 
+    # must be lower case and 'PDF' must be capital
+    weightgroups = cms.vstring(['scale', 'PDF', 'matrix element', 'unknown', 'shower']),
+    maxGroupsPerType = cms.vint32([1, -1, 1, 2, 1]),
+    pdfIds = cms.vint32([91400, 306000, 260000]),
     lheWeightPrecision = cms.int32(14),
 )
 

SimDataFormats/GeneratorProducts/interface/MEParamWeightGroupInfo.h

@@ -6,9 +6,9 @@
 namespace gen {
     class MEParamWeightGroupInfo : public WeightGroupInfo {
         public:
-            MEParamWeightGroupInfo() : WeightGroupInfo() { weightType_ = kMEParamWeights; }
+            MEParamWeightGroupInfo() : WeightGroupInfo() { weightType_ = WeightType::kMEParamWeights; }
 	        MEParamWeightGroupInfo(std::string header, std::string name) : 
-                WeightGroupInfo(header, name) { weightType_ = kMEParamWeights; } 
+                WeightGroupInfo(header, name) { weightType_ = WeightType::kMEParamWeights; } 
 	        MEParamWeightGroupInfo(std::string header) : 
                 MEParamWeightGroupInfo(header, header) {} 
             virtual ~MEParamWeightGroupInfo() override {}

SimDataFormats/GeneratorProducts/interface/PartonShowerWeightGroupInfo.h

@@ -9,7 +9,7 @@ namespace gen {
             PartonShowerWeightGroupInfo() : PartonShowerWeightGroupInfo("") {}
 	        PartonShowerWeightGroupInfo(std::string header, std::string name) : 
                 WeightGroupInfo(header, name) { 
-                weightType_ = kPartonShowerWeights;
+                weightType_ = WeightType::kPartonShowerWeights;
             }
 	        PartonShowerWeightGroupInfo(std::string header) : 
                 PartonShowerWeightGroupInfo(header, header) { } 

SimDataFormats/GeneratorProducts/interface/PdfWeightGroupInfo.h

@@ -20,11 +20,11 @@ namespace gen {
             int alphasDownIndex_;
             std::vector<std::pair<int, int>> lhapdfIdsContained_;
         public:
-            PdfWeightGroupInfo() : WeightGroupInfo() { weightType_ = kPdfWeights; }
+            PdfWeightGroupInfo() : WeightGroupInfo() { weightType_ = WeightType::kPdfWeights; }
 	        PdfWeightGroupInfo(std::string header, std::string name) : 
-                WeightGroupInfo(header, name) { weightType_ = kPdfWeights; } 
+                WeightGroupInfo(header, name) { weightType_ = WeightType::kPdfWeights; } 
 	        PdfWeightGroupInfo(std::string header) : 
-                WeightGroupInfo(header) { weightType_ = kPdfWeights; } 
+                WeightGroupInfo(header) { weightType_ = WeightType::kPdfWeights; } 
             PdfWeightGroupInfo(const PdfWeightGroupInfo &other) {
                 copy(other);
             }

SimDataFormats/GeneratorProducts/interface/ScaleWeightGroupInfo.h

@@ -20,7 +20,7 @@ namespace gen {
             ScaleWeightGroupInfo() : ScaleWeightGroupInfo("") {}
 	        ScaleWeightGroupInfo(std::string header, std::string name) : 
                 WeightGroupInfo(header, name) { 
-                weightType_ = kScaleWeights;
+                weightType_ = WeightType::kScaleWeights;
                 isFunctionalFormVar_ = false;
                 centralIndex_ = 0;
                 muR1muF2Index_ = 0;

SimDataFormats/GeneratorProducts/interface/UnknownWeightGroupInfo.h

@@ -6,11 +6,11 @@
 namespace gen {
     class UnknownWeightGroupInfo : public WeightGroupInfo {
         public:
-            UnknownWeightGroupInfo() : WeightGroupInfo() { weightType_ = kUnknownWeights; }
+            UnknownWeightGroupInfo() : WeightGroupInfo() { weightType_ = WeightType::kUnknownWeights; }
 	        UnknownWeightGroupInfo(std::string header, std::string name) : 
-                WeightGroupInfo(header, name) { weightType_ = kUnknownWeights; isWellFormed_ = false;} 
+                WeightGroupInfo(header, name) { weightType_ = WeightType::kUnknownWeights; isWellFormed_ = false;} 
 	        UnknownWeightGroupInfo(std::string header) : 
-                WeightGroupInfo(header) { weightType_ = kUnknownWeights; isWellFormed_ = false;} 
+                WeightGroupInfo(header) { weightType_ = WeightType::kUnknownWeights; isWellFormed_ = false;} 
             virtual ~UnknownWeightGroupInfo() override {}
             void copy(const UnknownWeightGroupInfo &other);
             virtual UnknownWeightGroupInfo* clone() const override;

SimDataFormats/GeneratorProducts/interface/WeightGroupInfo.h

@@ -7,6 +7,7 @@
 #include <string>
 #include <algorithm>
 #include <vector>
+#include <array>
 
 namespace gen {
     struct WeightMetaInfo {
@@ -16,14 +17,19 @@ namespace gen {
         std::string label;
     };
 
-    enum WeightType {
-        kPdfWeights,
-        kScaleWeights,
-        kMEParamWeights,
-        kUnknownWeights,
-        kPartonShowerWeights,
+    enum class WeightType : char {
+        kPdfWeights = 'P',
+        kScaleWeights = 's',
+        kMEParamWeights = 'm',
+        kPartonShowerWeights = 'p',
+        kUnknownWeights = 'u',
     };
 
+    const std::array<WeightType, 5> allGenWeightTypes = {{WeightType::kPdfWeights, 
+        WeightType::kScaleWeights, WeightType::kMEParamWeights, 
+        WeightType::kPartonShowerWeights, WeightType::kUnknownWeights, 
+    }};
+
 	class WeightGroupInfo {
         public:
             WeightGroupInfo() : WeightGroupInfo("") {}

SimDataFormats/GeneratorProducts/src/GenWeightInfoProduct.cc

@@ -55,7 +55,7 @@ std::vector<gen::WeightGroupInfo*> GenWeightInfoProduct::weightGroupsByType(gen:
 }
 
 std::optional<gen::WeightGroupData> GenWeightInfoProduct::pdfGroupWithIndexByLHAID(size_t lhaid) const {
-    auto pdfGroups = weightGroupsAndIndicesByType(gen::kPdfWeights);
+    auto pdfGroups = weightGroupsAndIndicesByType(gen::WeightType::kPdfWeights);
 
     auto matchingPdfSet = std::find_if(pdfGroups.begin(), pdfGroups.end(), 
             [lhaid] (gen::WeightGroupData& data) {