Setup esConsumes and random service within PileUp handling

cms-sw · Nov 22, 2022 · 6971e19 · 6971e19
1 parent 24b5ecc
commit 6971e19
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 8 deletions.
diff --git a/Mixing/Base/interface/PileUp.h b/Mixing/Base/interface/PileUp.h
@@ -4,21 +4,18 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include <optional>
 #include "FWCore/Framework/interface/Frameworkfwd.h"
 #include "FWCore/Framework/interface/ConsumesCollector.h"
 #include "FWCore/Sources/interface/VectorInputSource.h"
 #include "FWCore/Utilities/interface/ESGetToken.h"
 #include "DataFormats/Provenance/interface/EventID.h"
 #include "FWCore/Framework/interface/EventPrincipal.h"
+#include "FWCore/ServiceRegistry/interface/ServiceToken.h"
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
 
-#include "TRandom.h"
-#include "TFile.h"
 #include "TH1F.h"
 
-class TFile;
-class TH1F;
-
 namespace CLHEP {
   class RandPoissonQ;
   class RandPoisson;
@@ -27,6 +24,7 @@ namespace CLHEP {
 
 class MixingModuleConfig;
 class MixingRcd;
+class PileupRandomNumberGenerator;
 
 namespace edm {
   class SecondaryEventProvider;
@@ -81,6 +79,7 @@ namespace edm {
     void dropUnwantedBranches(std::vector<std::string> const& wantedBranches) {
       input_->dropUnwantedBranches(wantedBranches);
     }
+    void beginJob(eventsetup::ESRecordsToProxyIndices const&);
     void beginStream(edm::StreamID);
     void endStream();
 
@@ -111,6 +110,8 @@ namespace edm {
     std::unique_ptr<CLHEP::RandPoissonQ> const& poissonDistribution(StreamID const& streamID);
     std::unique_ptr<CLHEP::RandPoisson> const& poissonDistr_OOT(StreamID const& streamID);
     CLHEP::HepRandomEngine* randomEngine(StreamID const& streamID);
+    void setRandomEngine(StreamID);
+    void setRandomEngine(LuminosityBlock const&);
 
     unsigned int inputType_;
     std::string type_;
@@ -146,6 +147,8 @@ namespace edm {
     std::unique_ptr<EventPrincipal> eventPrincipal_;
     std::shared_ptr<LuminosityBlockPrincipal> lumiPrincipal_;
     std::shared_ptr<RunPrincipal> runPrincipal_;
+    PileupRandomNumberGenerator* randomGenerator_ = nullptr;
+    std::optional<ServiceToken> serviceToken_;
     std::unique_ptr<SecondaryEventProvider> provider_;
     std::unique_ptr<CLHEP::RandPoissonQ> PoissonDistribution_;
     std::unique_ptr<CLHEP::RandPoisson> PoissonDistr_OOT_;

diff --git a/Mixing/Base/src/PileUp.cc b/Mixing/Base/src/PileUp.cc
@@ -30,6 +30,8 @@
 #include "CLHEP/Random/RandPoissonQ.h"
 #include "CLHEP/Random/RandPoisson.h"
 
+#include "PileupRandomNumberGenerator.h"
+
 #include <algorithm>
 #include <memory>
 #include "TMath.h"
@@ -105,6 +107,18 @@ namespace edm {
 
     if (pset.existsAs<std::vector<ParameterSet> >("producers", true)) {
       std::vector<ParameterSet> producers = pset.getParameter<std::vector<ParameterSet> >("producers");
+
+      std::vector<std::string> names;
+      names.reserve(producers.size());
+      std::transform(producers.begin(), producers.end(), std::back_inserter(names), [](edm::ParameterSet const& iPSet) {
+        return iPSet.getParameter<std::string>("@module_label");
+      });
+      auto randomGenerator = std::make_unique<PileupRandomNumberGenerator>(names);
+      randomGenerator_ = randomGenerator.get();
+      std::unique_ptr<edm::RandomNumberGenerator> baseGen = std::move(randomGenerator);
+      serviceToken_ = edm::ServiceRegistry::createContaining(
+          std::move(baseGen), edm::ServiceRegistry::instance().presentToken(), true);
+
       provider_ = std::make_unique<SecondaryEventProvider>(producers, *productRegistry_, processConfiguration_);
     }
 
@@ -184,19 +198,26 @@ namespace edm {
     }
   }  // end of constructor
 
+  void PileUp::beginJob(eventsetup::ESRecordsToProxyIndices const& iES) {
+    input_->doBeginJob();
+    if (provider_.get() != nullptr) {
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
+      provider_->beginJob(*productRegistry_, iES);
+    }
+  }
+
   void PileUp::beginStream(edm::StreamID) {
     auto iID = eventPrincipal_->streamID();  // each producer has its own workermanager, so use default streamid
     streamContext_.reset(new StreamContext(iID, processContext_.get()));
-    input_->doBeginJob();
     if (provider_.get() != nullptr) {
-      //TODO for now, we do not support consumes from EventSetup
-      provider_->beginJob(*productRegistry_, eventsetup::ESRecordsToProxyIndices{{}});
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->beginStream(iID, *streamContext_);
     }
   }
 
   void PileUp::endStream() {
     if (provider_.get() != nullptr) {
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->endStream(streamContext_->streamID(), *streamContext_);
       provider_->endJob();
     }
@@ -207,6 +228,7 @@ namespace edm {
     if (provider_.get() != nullptr) {
       runPrincipal_.reset(new RunPrincipal(productRegistry_, *processConfiguration_, nullptr, 0));
       runPrincipal_->setAux(run.runAuxiliary());
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->beginRun(*runPrincipal_, setup.impl(), run.moduleCallingContext(), *streamContext_);
     }
   }
@@ -215,17 +237,21 @@ namespace edm {
       lumiPrincipal_.reset(new LuminosityBlockPrincipal(productRegistry_, *processConfiguration_, nullptr, 0));
       lumiPrincipal_->setAux(lumi.luminosityBlockAuxiliary());
       lumiPrincipal_->setRunPrincipal(runPrincipal_);
+      setRandomEngine(lumi);
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->beginLuminosityBlock(*lumiPrincipal_, setup.impl(), lumi.moduleCallingContext(), *streamContext_);
     }
   }
 
   void PileUp::endRun(const edm::Run& run, const edm::EventSetup& setup) {
     if (provider_.get() != nullptr) {
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->endRun(*runPrincipal_, setup.impl(), run.moduleCallingContext(), *streamContext_);
     }
   }
   void PileUp::endLuminosityBlock(const edm::LuminosityBlock& lumi, const edm::EventSetup& setup) {
     if (provider_.get() != nullptr) {
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->endLuminosityBlock(*lumiPrincipal_, setup.impl(), lumi.moduleCallingContext(), *streamContext_);
     }
   }
@@ -235,6 +261,7 @@ namespace edm {
       // note:  run and lumi numbers must be modified to match lumiPrincipal_
       eventPrincipal_->setLuminosityBlockPrincipal(lumiPrincipal_.get());
       eventPrincipal_->setRunAndLumiNumber(lumiPrincipal_->run(), lumiPrincipal_->luminosityBlock());
+      edm::ServiceRegistry::Operate guard(*serviceToken_);
       provider_->setupPileUpEvent(*eventPrincipal_, setup.impl(), *streamContext_);
     }
   }
@@ -353,6 +380,13 @@ namespace edm {
     return randomEngine_;
   }
 
+  void PileUp::setRandomEngine(StreamID streamID) { randomGenerator_->setEngine(*randomEngine(streamID)); }
+  void PileUp::setRandomEngine(LuminosityBlock const& iLumi) {
+    Service<RandomNumberGenerator> rng;
+    randomGenerator_->setSeed(rng->mySeed());
+    randomGenerator_->setEngine(rng->getEngine(iLumi.index()));
+  }
+
   void PileUp::CalculatePileup(int MinBunch,
                                int MaxBunch,
                                std::vector<int>& PileupSelection,
@@ -364,6 +398,10 @@ namespace edm {
     int nzero_crossing = -1;
     double Fnzero_crossing = -1;
 
+    if (provider_) {
+      setRandomEngine(streamID);
+    }
+
     if (manage_OOT_) {
       if (none_) {
         nzero_crossing = 0;