Merge pull request #32019 from kpedro88/SLHATree_94X

Read SLHA from TTree [9_4_X]

GeneratorInterface/Pythia8Interface/BuildFile.xml

@@ -10,6 +10,7 @@
 <use name="vincia"/>
 <use name="hepmc"/>
 <use name="clhep"/>
+<use name="root"/>
 <export>
   <lib   name="1"/>
 </export>

GeneratorInterface/Pythia8Interface/interface/Py8InterfaceBase.h

@@ -31,6 +31,7 @@ namespace gen {
          virtual bool generatePartonsAndHadronize() = 0;
          bool decay() { return true; } // NOT used - let's call it "design imperfection"
          bool readSettings( int ); // common func
+         void makeTmpSLHA(const std::string&); //helper for above
          virtual bool initializeForInternalPartons() = 0;
          bool declareStableParticles( const std::vector<int>& ); // common func
          bool declareSpecialSettings( const std::vector<std::string>& ); // common func

GeneratorInterface/Pythia8Interface/interface/SLHAReaderBase.h

@@ -0,0 +1,32 @@
+#ifndef GeneratorInterface_Pythia8Interface_SLHAReaderBase
+#define GeneratorInterface_Pythia8Interface_SLHAReaderBase
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include <string>
+#include <vector>
+
+class TFile;
+class TTree;
+
+class SLHAReaderBase {
+public:
+  SLHAReaderBase(const edm::ParameterSet& conf);
+  virtual ~SLHAReaderBase();
+
+  //this function should parse the config description (e.g. with splitline() below)
+  //then use the information to get the SLHA info out of the tree and return it
+  virtual std::string getSLHA(const std::string& configDesc) = 0;
+
+  static std::vector<std::string> splitline(const std::string& line, char delim);
+
+protected:
+  //members
+  TFile* file_;
+  TTree* tree_;
+};
+
+#include "FWCore/PluginManager/interface/PluginFactory.h"
+typedef edmplugin::PluginFactory<SLHAReaderBase*(const edm::ParameterSet&)> SLHAReaderFactory;
+
+#endif

GeneratorInterface/Pythia8Interface/plugins/BuildFile.xml

@@ -8,3 +8,8 @@
   <use   name="GeneratorInterface/ExternalDecays"/>
   <flags   EDM_PLUGIN="1"/>
 </library>
+
+<library file="SLHA*.cc" name="GeneratorInterfacePythia8SLHAReader">
+  <use name="GeneratorInterface/Pythia8Interface"/>
+  <use name="root"/>
+</library>

GeneratorInterface/Pythia8Interface/plugins/SLHAReaderpMSSM.cc

@@ -0,0 +1,29 @@
+#include "GeneratorInterface/Pythia8Interface/interface/SLHAReaderBase.h"
+
+#include "TTree.h"
+#include "TString.h"
+
+#include <memory>
+
+class SLHAReaderpMSSM : public SLHAReaderBase {
+public:
+  SLHAReaderpMSSM(const edm::ParameterSet& conf) : SLHAReaderBase(conf) {}
+  ~SLHAReaderpMSSM() override {}
+
+  std::string getSLHA(const std::string& configDesc) override;
+};
+
+DEFINE_EDM_PLUGIN(SLHAReaderFactory, SLHAReaderpMSSM, "SLHAReaderpMSSM");
+
+std::string SLHAReaderpMSSM::getSLHA(const std::string& configDesc) {
+  const auto& config_fields = splitline(configDesc, '_');
+  int chain = std::stoi(config_fields.at(2));
+  int iteration = std::stoi(config_fields.at(3));
+
+  auto slhabranch = std::make_unique<TString>();
+  auto slhabranch_ptr = slhabranch.get();
+  tree_->SetBranchAddress("slhacontent", &slhabranch_ptr);
+  tree_->GetEntryWithIndex(chain, iteration);
+
+  return std::string(*slhabranch);
+}

GeneratorInterface/Pythia8Interface/src/Py8InterfaceBase.cc

@@ -1,4 +1,8 @@
+#include <memory>
+#include <string>
+
 #include "GeneratorInterface/Pythia8Interface/interface/Py8InterfaceBase.h"
+#include "GeneratorInterface/Pythia8Interface/interface/SLHAReaderBase.h"
 
 #include "FWCore/Utilities/interface/Exception.h"
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
@@ -130,21 +134,29 @@ bool Py8InterfaceBase::readSettings( int )
    else if( currentParameters.exists( "SLHATableForPythia8" ) ) {
      std::string slhatable = currentParameters.getParameter<std::string>("SLHATableForPythia8");
 
-     char tempslhaname[] = "pythia8SLHAtableXXXXXX";
-     int fd = mkstemp(tempslhaname);
-     write(fd,slhatable.c_str(),slhatable.size());
-     close(fd);
-
-     slhafile_ = tempslhaname;
-
-     fMasterGen->settings.mode("SLHA:readFrom", 2);
-     fMasterGen->settings.word("SLHA:file", slhafile_);
-   }   
+      makeTmpSLHA(slhatable);
+    } else if (currentParameters.exists("SLHATreeForPythia8")) {
+      auto slhaReaderParams = currentParameters.getParameter<edm::ParameterSet>("SLHATreeForPythia8");
+      std::unique_ptr<SLHAReaderBase> reader(SLHAReaderFactory::get()->create(slhaReaderParams.getParameter<std::string>("name"), slhaReaderParams));
+      makeTmpSLHA(reader->getSLHA(currentParameters.getParameter<std::string>("ConfigDescription")));
+    }
 
    return true;
 
 }
 
+void Py8InterfaceBase::makeTmpSLHA(const std::string& slhatable) {
+  char tempslhaname[] = "pythia8SLHAtableXXXXXX";
+  int fd = mkstemp(tempslhaname);
+  write(fd, slhatable.c_str(), slhatable.size());
+  close(fd);
+
+  slhafile_ = tempslhaname;
+
+  fMasterGen->settings.mode("SLHA:readFrom", 2);
+  fMasterGen->settings.word("SLHA:file", slhafile_);
+}
+
 bool Py8InterfaceBase::declareStableParticles( const std::vector<int>& pdgIds )
 {
 

GeneratorInterface/Pythia8Interface/src/SLHAReaderBase.cc

@@ -0,0 +1,33 @@
+#include "GeneratorInterface/Pythia8Interface/interface/SLHAReaderBase.h"
+#include "FWCore/Utilities/interface/Exception.h"
+
+#include <sstream>
+
+#include "TFile.h"
+#include "TTree.h"
+
+SLHAReaderBase::SLHAReaderBase(const edm::ParameterSet& conf) {
+  auto filename = conf.getParameter<std::string>("file");
+  file_ = TFile::Open(filename.c_str());
+  if (!file_)
+    throw cms::Exception("MissingFile") << "Could not open file " << filename;
+
+  auto treename = conf.getParameter<std::string>("tree");
+  tree_ = (TTree*)file_->Get(treename.c_str());
+  if (!tree_)
+    throw cms::Exception("MissingTree") << "Could not get tree " << treename << " from file " << filename;
+}
+
+SLHAReaderBase::~SLHAReaderBase() { file_->Close(); }
+
+std::vector<std::string> SLHAReaderBase::splitline(const std::string& line, char delim) {
+  std::stringstream ss(line);
+  std::string field;
+  std::vector<std::string> fields;
+  while (getline(ss, field, delim)) {
+    fields.push_back(field);
+  }
+  return fields;
+}
+
+EDM_REGISTER_PLUGINFACTORY(SLHAReaderFactory, "SLHAReaderFactory");