Merge pull request #40133 from Dr15Jones/reviveModulesInsideMixing

Revived being able to run modules inside a mixing module

FWCore/ServiceRegistry/interface/ServiceRegistry.h

@@ -109,11 +109,12 @@ namespace edm {
     template <typename T>
     static ServiceToken createContaining(std::unique_ptr<T> iService,
                                          ServiceToken iToken,
-                                         serviceregistry::ServiceLegacy iLegacy) {
+                                         bool iOverrideIfServiceInToken) {
       std::vector<edm::ParameterSet> config;
-      auto manager = std::make_shared<serviceregistry::ServicesManager>(iToken, iLegacy, config);
+      auto manager =
+          std::make_shared<serviceregistry::ServicesManager>(iToken, serviceregistry::kOverlapIsError, config);
       auto wrapper = std::make_shared<serviceregistry::ServiceWrapper<T> >(std::move(iService));
-      manager->put(wrapper);
+      manager->put(wrapper, iOverrideIfServiceInToken);
       return manager;
     }
     /// create a service token that holds the service held by iWrapper

FWCore/ServiceRegistry/interface/ServicesManager.h

@@ -126,9 +126,9 @@ namespace edm {
       // ---------- member functions ---------------------------
       ///returns false if put fails because a service of this type already exists
       template <typename T>
-      bool put(std::shared_ptr<ServiceWrapper<T>> iPtr) {
+      bool put(std::shared_ptr<ServiceWrapper<T>> iPtr, bool iOverride = false) {
         Type2Service::const_iterator itFound = type2Service_.find(TypeIDBase(typeid(T)));
-        if (itFound != type2Service_.end()) {
+        if (itFound != type2Service_.end() and not iOverride) {
           return false;
         }
         type2Service_[TypeIDBase(typeid(T))] = iPtr;

IOMC/RandomEngine/interface/cloneEngine.h

@@ -0,0 +1,30 @@
+#ifndef IOMC_RandomEngine_cloneEngine_h
+#define IOMC_RandomEngine_cloneEngine_h
+// -*- C++ -*-
+//
+// Package:     IOMC/RandomEngine
+// Class  :     cloneEngine
+//
+/**\function cloneEngine cloneEngine.h "IOMC/RandomEngine/interface/cloneEngine.h"
+
+ Description: Function used to clone a random number engine
+
+ Usage:
+    <usage>
+
+*/
+//
+// Original Author:  Christopher Jones
+//         Created:  Fri, 02 Dec 2022 19:32:10 GMT
+//
+#include <memory>
+
+namespace CLHEP {
+  class HepRandomEngine;
+}
+
+namespace edm {
+  std::unique_ptr<CLHEP::HepRandomEngine> cloneEngine(CLHEP::HepRandomEngine const&);
+};
+
+#endif

IOMC/RandomEngine/plugins/BuildFile.xml

@@ -1,6 +1,6 @@
 <use name="FWCore/ServiceRegistry"/>
 <use name="FWCore/Framework"/>
 <use name="IOMC/RandomEngine"/>
-<library file="Module.cc" name="IOMCRandomEnginePlugins">
+<library file="*.cc" name="IOMCRandomEnginePlugins">
   <flags EDM_PLUGIN="1"/>
 </library>

IOMC/RandomEngine/plugins/Module.cc

@@ -2,9 +2,9 @@
 #include "FWCore/Framework/interface/MakerMacros.h"
 #include "FWCore/ServiceRegistry/interface/ServiceMaker.h"
 
-#include "IOMC/RandomEngine/src/RandomNumberGeneratorService.h"
-#include "IOMC/RandomEngine/src/RandomEngineStateProducer.h"
-#include "IOMC/RandomEngine/src/RandomFilter.h"
+#include "RandomNumberGeneratorService.h"
+#include "RandomEngineStateProducer.h"
+#include "RandomFilter.h"
 
 using edm::RandomFilter;
 using edm::service::RandomNumberGeneratorService;

IOMC/RandomEngine/src/RandomEngineStateProducer.cc → IOMC/RandomEngine/plugins/RandomEngineStateProducer.cc

@@ -1,5 +1,5 @@
 
-#include "IOMC/RandomEngine/src/RandomEngineStateProducer.h"
+#include "RandomEngineStateProducer.h"
 
 #include "FWCore/Framework/interface/Event.h"
 #include "FWCore/Framework/interface/LuminosityBlock.h"

IOMC/RandomEngine/src/RandomFilter.cc → IOMC/RandomEngine/plugins/RandomFilter.cc

@@ -1,5 +1,5 @@
 
-#include "IOMC/RandomEngine/src/RandomFilter.h"
+#include "RandomFilter.h"
 
 #include "FWCore/Framework/interface/Event.h"
 #include "FWCore/ParameterSet/interface/ParameterSet.h"

IOMC/RandomEngine/src/RandomNumberGeneratorService.cc → IOMC/RandomEngine/plugins/RandomNumberGeneratorService.cc

@@ -10,7 +10,7 @@
 //   Created:  Tue Mar  7 09:43:46 EST 2006 (originally in FWCore/Services)
 //
 
-#include "IOMC/RandomEngine/src/RandomNumberGeneratorService.h"
+#include "RandomNumberGeneratorService.h"
 
 #include "DataFormats/Common/interface/Handle.h"
 #include "DataFormats/Provenance/interface/ModuleDescription.h"
@@ -35,6 +35,7 @@
 #include "FWCore/Utilities/interface/LuminosityBlockIndex.h"
 #include "FWCore/Utilities/interface/StreamID.h"
 #include "IOMC/RandomEngine/interface/TRandomAdaptor.h"
+#include "IOMC/RandomEngine/interface/cloneEngine.h"
 #include "SimDataFormats/RandomEngine/interface/RandomEngineState.h"
 #include "SimDataFormats/RandomEngine/interface/RandomEngineStates.h"
 
@@ -289,25 +290,7 @@ namespace edm {
 
     std::unique_ptr<CLHEP::HepRandomEngine> RandomNumberGeneratorService::cloneEngine(
         LuminosityBlockIndex const& lumiIndex) {
-      CLHEP::HepRandomEngine& existingEngine = getEngine(lumiIndex);
-
-      std::vector<unsigned long> stateL = existingEngine.put();
-      long seedL = existingEngine.getSeed();
-      std::unique_ptr<CLHEP::HepRandomEngine> newEngine;
-      if (stateL[0] == CLHEP::engineIDulong<CLHEP::HepJamesRandom>()) {
-        newEngine = std::make_unique<CLHEP::HepJamesRandom>(seedL);
-      } else if (stateL[0] == CLHEP::engineIDulong<CLHEP::RanecuEngine>()) {
-        newEngine = std::make_unique<CLHEP::RanecuEngine>();
-      } else if (stateL[0] == CLHEP::engineIDulong<CLHEP::MixMaxRng>()) {
-        newEngine = std::make_unique<CLHEP::MixMaxRng>(seedL);
-      } else if (stateL[0] == CLHEP::engineIDulong<TRandomAdaptor>()) {
-        newEngine = std::make_unique<TRandomAdaptor>(seedL);
-      } else {
-        // Sanity check, it should not be possible for this to happen.
-        throw Exception(errors::Unknown) << "The RandomNumberGeneratorService is trying to clone unknown engine type\n";
-      }
-      newEngine->get(stateL);
-      return newEngine;
+      return edm::cloneEngine(getEngine(lumiIndex));
     }
 
     // PROBABLY TO BE DELETED, This returns the configured seed without

IOMC/RandomEngine/src/cloneEngine.cc

@@ -0,0 +1,46 @@
+// -*- C++ -*-
+//
+// Package:     IOMC/RandomEngine
+// Class  :     cloneEngine
+//
+// Implementation:
+//     [Notes on implementation]
+//
+// Original Author:  Christopher Jones
+//         Created:  Fri, 02 Dec 2022 19:34:37 GMT
+//
+
+// system include files
+
+// user include files
+#include "IOMC/RandomEngine/interface/cloneEngine.h"
+#include "IOMC/RandomEngine/interface/TRandomAdaptor.h"
+
+#include "FWCore/Utilities/interface/EDMException.h"
+
+#include "CLHEP/Random/engineIDulong.h"
+#include "CLHEP/Random/JamesRandom.h"
+#include "CLHEP/Random/RanecuEngine.h"
+#include "CLHEP/Random/MixMaxRng.h"
+
+namespace edm {
+  std::unique_ptr<CLHEP::HepRandomEngine> cloneEngine(CLHEP::HepRandomEngine const& existingEngine) {
+    std::vector<unsigned long> stateL = existingEngine.put();
+    long seedL = existingEngine.getSeed();
+    std::unique_ptr<CLHEP::HepRandomEngine> newEngine;
+    if (stateL[0] == CLHEP::engineIDulong<CLHEP::HepJamesRandom>()) {
+      newEngine = std::make_unique<CLHEP::HepJamesRandom>(seedL);
+    } else if (stateL[0] == CLHEP::engineIDulong<CLHEP::RanecuEngine>()) {
+      newEngine = std::make_unique<CLHEP::RanecuEngine>();
+    } else if (stateL[0] == CLHEP::engineIDulong<CLHEP::MixMaxRng>()) {
+      newEngine = std::make_unique<CLHEP::MixMaxRng>(seedL);
+    } else if (stateL[0] == CLHEP::engineIDulong<TRandomAdaptor>()) {
+      newEngine = std::make_unique<TRandomAdaptor>(seedL);
+    } else {
+      // Sanity check, it should not be possible for this to happen.
+      throw Exception(errors::Unknown) << "The RandomNumberGeneratorService is trying to clone unknown engine type\n";
+    }
+    newEngine->get(stateL);
+    return newEngine;
+  }
+};  // namespace edm

Mixing/Base/BuildFile.xml

@@ -7,6 +7,8 @@
 <use name="FWCore/Sources"/>
 <use name="FWCore/Utilities"/>
 <use name="FWCore/Version"/>
+<use name="IOMC/RandomEngine"/>
+<use name="SimDataFormats/RandomEngine"/>
 <use name="clhep"/>
 <use name="roothistmatrix"/>
 <use name="CondFormats/RunInfo"/>

Mixing/Base/interface/BMixingModule.h

@@ -46,6 +46,8 @@ namespace edm {
     /**Default destructor*/
     ~BMixingModule() override;
 
+    void registerLateConsumes(eventsetup::ESRecordsToProxyIndices const&) override;
+
     /**Cumulates the pileup events onto this event*/
     void produce(edm::Event& e1, const edm::EventSetup& c) override;
 

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_;

Mixing/Base/src/BMixingModule.cc

@@ -201,6 +201,13 @@ namespace edm {
   // Virtual destructor needed.
   BMixingModule::~BMixingModule() { ; }
 
+  void BMixingModule::registerLateConsumes(eventsetup::ESRecordsToProxyIndices const& iES) {
+    for (size_t endIdx = 0; endIdx < maxNbSources_; ++endIdx) {
+      if (inputSources_[endIdx])
+        inputSources_[endIdx]->beginJob(iES);
+    }
+  }
+
   namespace MixingCache {
     Config::Config(edm::ParameterSet const& pset, unsigned int maxNbSources)
         : bunchSpace_(pset.getParameter<int>("bunchspace")),

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;