GeneratorInterface/LHEInterface/plugins/LH5Source.cc - master - githubmirror/cmssw

GeneratorInterface/LHEInterface/plugins/LH5Source.cc
245 строк9 KB
W. David Dagenhart
Reduce the memory usage of the InputFileCatalog
18 май 2026, 23:53
18 май 2026, 23:537e6433c
#include <algorithm>
#include <cstddef>
#include <functional>
#include <iostream>
#include <string>
#include <memory>

#include "FWCore/Framework/interface/InputSourceMacros.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/Framework/interface/EventPrincipal.h"
#include "FWCore/Framework/interface/LuminosityBlockPrincipal.h"
#include "FWCore/Framework/interface/RunPrincipal.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/Sources/interface/ProducerSourceBase.h"
#include "FWCore/Utilities/interface/TypeID.h"
#include "FWStorage/Catalog/interface/InputFileCatalog.h"

#include "DataFormats/Common/interface/OrphanHandle.h"
#include "DataFormats/Provenance/interface/LuminosityBlockAuxiliary.h"
#include "DataFormats/Provenance/interface/ProcessHistoryID.h"
#include "DataFormats/Provenance/interface/RunAuxiliary.h"
#include "DataFormats/Provenance/interface/Timestamp.h"

#include "SimDataFormats/GeneratorProducts/interface/LesHouches.h"
#include "SimDataFormats/GeneratorProducts/interface/LHERunInfoProduct.h"
#include "SimDataFormats/GeneratorProducts/interface/LHEEventProduct.h"

#include "GeneratorInterface/LHEInterface/interface/LHEEvent.h"
#include "GeneratorInterface/LHEInterface/plugins/LHEProvenanceHelper.h"
#include "GeneratorInterface/LHEInterface/interface/LH5Reader.h"
#include "GeneratorInterface/LHEInterface/interface/LHERunInfo.h"

using namespace lhef;

class LH5Source : public edm::ProducerSourceBase {
public:
  explicit LH5Source(const edm::ParameterSet& params, const edm::InputSourceDescription& desc);
  ~LH5Source() override;

  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

private:
  void endJob() override;
  bool setRunAndEventInfo(edm::EventID&, edm::TimeValue_t&, edm::EventAuxiliary::ExperimentType&) override;
  void readRun_(edm::RunPrincipal& runPrincipal) override;
  void readLuminosityBlock_(edm::LuminosityBlockPrincipal& lumiPrincipal) override;
  void readEvent_(edm::EventPrincipal& eventPrincipal) override;
  void produce(edm::Event&) override {}
  std::shared_ptr<edm::RunAuxiliary> readRunAuxiliary_() override;
  std::shared_ptr<edm::LuminosityBlockAuxiliary> readLuminosityBlockAuxiliary_() override;
  size_t fileIndex() const override;

  void nextEvent();

  void putRunInfoProduct(edm::RunPrincipal&);
  void fillRunInfoProduct(lhef::LHERunInfo const&, LHERunInfoProduct&);

  edm::InputFileCatalog inputFileCatalog_;
  size_t fileIndex_ = 0;
  std::unique_ptr<lhef::LH5Reader> reader_;

  std::shared_ptr<lhef::LHERunInfo> runInfoLast_;
  std::shared_ptr<lhef::LHEEvent> partonLevel_;

  std::unique_ptr<LHERunInfoProduct> runInfoProductLast_;
  edm::LHEProvenanceHelper lheProvenanceHelper_;
  edm::ProcessHistoryID phid_;
};

LH5Source::LH5Source(const edm::ParameterSet& params, const edm::InputSourceDescription& desc)
    : ProducerSourceBase(params, desc, false),
      inputFileCatalog_(params),
      //      reader_(new LH5Reader(params)),
      reader_(new LH5Reader(inputFileCatalog_.allPFNsFromFirstCatalog(),
                            params.getUntrackedParameter<unsigned int>("skipEvents", 0))),
      lheProvenanceHelper_(edm::TypeID(typeid(LHEEventProduct)),
                           edm::TypeID(typeid(LHERunInfoProduct)),
                           productRegistryUpdate(),
                           *branchIDListHelper()),
      phid_() {
  nextEvent();
  lheProvenanceHelper_.lheAugment(nullptr);
  // Initialize metadata, and save the process history ID for use every event.
  phid_ = lheProvenanceHelper_.lheInit(processHistoryRegistryForUpdate());

  // These calls are not wanted, because the principals are used for putting the products.
  //produces<LHEEventProduct>();
  //produces<LHERunInfoProduct, edm::InRun>();
}

LH5Source::~LH5Source() {}

void LH5Source::endJob() { reader_.reset(); }

void LH5Source::nextEvent() {
  if (partonLevel_) {
    return;
  }

  bool newFileOpened;
  do {
    newFileOpened = false;
    partonLevel_ = reader_->next(&newFileOpened);
    if (newFileOpened) {
      ++fileIndex_;
    }
  } while (newFileOpened && !partonLevel_);

  if (!partonLevel_) {
    return;
  }

  auto runInfoThis = partonLevel_->getRunInfo();
  if (runInfoThis != runInfoLast_) {
    runInfoLast_ = runInfoThis;
    std::unique_ptr<LHERunInfoProduct> product = std::make_unique<LHERunInfoProduct>(*runInfoThis->getHEPRUP());
    fillRunInfoProduct(*runInfoThis, *product);

    if (runInfoProductLast_) {
      if (!runInfoProductLast_->mergeProduct(*product)) {
        //cannot be merged so must start new Run
        runInfoProductLast_ = std::move(product);
        lheProvenanceHelper_.lheAugment(runInfoThis.get());
        // Initialize metadata, and save the process history ID for use every event.
        phid_ = lheProvenanceHelper_.lheInit(processHistoryRegistryForUpdate());
        resetRunAuxiliary();
      }
    } else {
      runInfoProductLast_ = std::move(product);
    }
  }
}

void LH5Source::fillRunInfoProduct(lhef::LHERunInfo const& iInfo, LHERunInfoProduct& oProduct) {
  for (auto const& h : iInfo.getHeaders()) {
    oProduct.addHeader(h);
  }
  for (auto const& c : iInfo.getComments()) {
    oProduct.addComment(c);
  }
}

void LH5Source::readRun_(edm::RunPrincipal& runPrincipal) {
  runAuxiliary()->setProcessHistoryID(phid_);
  runPrincipal.fillRunPrincipal(processHistoryRegistryForUpdate());

  putRunInfoProduct(runPrincipal);
}

void LH5Source::readLuminosityBlock_(edm::LuminosityBlockPrincipal& lumiPrincipal) {
  luminosityBlockAuxiliary()->setProcessHistoryID(phid_);
  lumiPrincipal.fillLuminosityBlockPrincipal(
      processHistoryRegistry().getMapped(lumiPrincipal.aux().processHistoryID()));
}

void LH5Source::putRunInfoProduct(edm::RunPrincipal& iRunPrincipal) {
  if (runInfoProductLast_) {
    auto product = std::make_unique<LHERunInfoProduct>(*runInfoProductLast_);
    std::unique_ptr<edm::WrapperBase> rdp = std::make_unique<edm::Wrapper<LHERunInfoProduct>>(std::move(product));
    iRunPrincipal.put(lheProvenanceHelper_.runProductProductDescription_, std::move(rdp));
  }
}

bool LH5Source::setRunAndEventInfo(edm::EventID&, edm::TimeValue_t&, edm::EventAuxiliary::ExperimentType&) {
  nextEvent();
  if (!partonLevel_) {
    // We just finished an input file. See if there is another.
    nextEvent();
    if (!partonLevel_) {
      // No more input files.
      return false;
    }
  }
  return true;
}

void LH5Source::readEvent_(edm::EventPrincipal& eventPrincipal) {
  assert(eventCached() || processingMode() != RunsLumisAndEvents);
  edm::EventAuxiliary aux(eventID(), processGUID(), edm::Timestamp(presentTime()), false);
  aux.setProcessHistoryID(phid_);
  eventPrincipal.fillEventPrincipal(aux,
                                    processHistoryRegistry().getMapped(aux.processHistoryID()),
                                    edm::EventSelectionIDVector(),
                                    lheProvenanceHelper_.branchListIndexes_);

  std::unique_ptr<LHEEventProduct> product =
      std::make_unique<LHEEventProduct>(*partonLevel_->getHEPEUP(), partonLevel_->originalXWGTUP());
  if (partonLevel_->getPDF()) {
    product->setPDF(*partonLevel_->getPDF());
  }
  std::for_each(partonLevel_->weights().begin(),
                partonLevel_->weights().end(),
                std::bind(&LHEEventProduct::addWeight, product.get(), std::placeholders::_1));
  product->setScales(partonLevel_->scales());
  product->setNpLO(partonLevel_->npLO());
  product->setNpNLO(partonLevel_->npNLO());
  std::for_each(partonLevel_->getComments().begin(),
                partonLevel_->getComments().end(),
                std::bind(&LHEEventProduct::addComment, product.get(), std::placeholders::_1));

  std::unique_ptr<edm::WrapperBase> edp = std::make_unique<edm::Wrapper<LHEEventProduct>>(std::move(product));
  eventPrincipal.put(lheProvenanceHelper_.eventProductProductDescription_,
                     std::move(edp),
                     lheProvenanceHelper_.eventProductProvenance_);

  partonLevel_.reset();

  resetEventCached();
}

std::shared_ptr<edm::RunAuxiliary> LH5Source::readRunAuxiliary_() {
  edm::Timestamp ts = edm::Timestamp(presentTime());
  resetNewRun();
  auto aux = std::make_shared<edm::RunAuxiliary>(eventID().run(), ts, edm::Timestamp::invalidTimestamp());
  aux->setProcessHistoryID(phid_);
  return aux;
}

std::shared_ptr<edm::LuminosityBlockAuxiliary> LH5Source::readLuminosityBlockAuxiliary_() {
  if (processingMode() == Runs)
    return std::shared_ptr<edm::LuminosityBlockAuxiliary>();
  edm::Timestamp ts = edm::Timestamp(presentTime());
  resetNewLumi();
  auto aux = std::make_shared<edm::LuminosityBlockAuxiliary>(
      eventID().run(), eventID().luminosityBlock(), ts, edm::Timestamp::invalidTimestamp());
  aux->setProcessHistoryID(phid_);
  return aux;
}

size_t LH5Source::fileIndex() const { return fileIndex_; }

void LH5Source::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  desc.setComment("A source which reads LHE files.");
  edm::ProducerSourceBase::fillDescription(desc);
  edm::InputFileCatalog::fillDescription(desc);

  desc.addUntracked<unsigned int>("skipEvents", 0U)->setComment("Skip the first 'skipEvents' events.");
  //  desc.addUntracked<int>("limitEvents", -1)->setComment("Limit the number of read events.");
  descriptions.add("source", desc);
}

DEFINE_FWK_INPUT_SOURCE(LH5Source);