diff --git a/k4Gen/options/filterRule.hxx b/k4Gen/options/filterRule.hxx new file mode 100644 index 0000000..01694d8 --- /dev/null +++ b/k4Gen/options/filterRule.hxx @@ -0,0 +1,5 @@ +#include "edm4hep/MCParticleCollection.h" + +bool filterRule(const edm4hep::MCParticleCollection* inColl) { + return inColl->size() > 1000; +} diff --git a/k4Gen/options/pythiaEventsFiltered.py b/k4Gen/options/pythiaEventsFiltered.py new file mode 100644 index 0000000..5bc6c45 --- /dev/null +++ b/k4Gen/options/pythiaEventsFiltered.py @@ -0,0 +1,73 @@ +''' +Pythia8, integrated in the Key4hep ecosystem. + +Generate events according to a Pythia .cmd file and save them in EDM4hep +format. +''' + +import os + +from GaudiKernel import SystemOfUnits as units +from Gaudi.Configuration import INFO, DEBUG + +from Configurables import ApplicationMgr, k4DataSvc, PodioOutput +from Configurables import GaussSmearVertex, PythiaInterface, GenAlg +from Configurables import HepMCToEDMConverter, GenParticleFilter +from Configurables import GenEventFilter + + +ApplicationMgr().EvtSel = 'NONE' +ApplicationMgr().EvtMax = 20 +ApplicationMgr().OutputLevel = INFO +ApplicationMgr().ExtSvc += ["RndmGenSvc"] + +# Data service +podioevent = k4DataSvc("EventDataSvc") +ApplicationMgr().ExtSvc += [podioevent] + +smeartool = GaussSmearVertex() +smeartool.xVertexSigma = 0.5 * units.mm +smeartool.yVertexSigma = 0.5 * units.mm +smeartool.zVertexSigma = 40.0 * units.mm +smeartool.tVertexSigma = 180.0 * units.picosecond + +pythia8gentool = PythiaInterface() +# Example of Pythia configuration file to generate events +# taken from $K4GEN if defined, locally otherwise +path_to_pythiafile = os.environ.get("K4GEN", "") +PYTHIA_FILENAME = "Pythia_standard.cmd" +pythiafile = os.path.join(path_to_pythiafile, PYTHIA_FILENAME) +# Example of pythia configuration file to read LH event file +# pythiafile="options/Pythia_LHEinput.cmd" +pythia8gentool.pythiacard = pythiafile +pythia8gentool.doEvtGenDecays = False +pythia8gentool.printPythiaStatistics = True +pythia8gentool.pythiaExtraSettings = [""] + +pythia8gen = GenAlg("Pythia8") +pythia8gen.SignalProvider = pythia8gentool +pythia8gen.VertexSmearingTool = smeartool +pythia8gen.hepmc.Path = "hepmc" +ApplicationMgr().TopAlg += [pythia8gen] + +# Reads an HepMC::GenEvent from the data service and writes a collection of +# EDM Particles +hepmc_converter = HepMCToEDMConverter() +hepmc_converter.hepmc.Path = "hepmc" +hepmc_converter.hepmcStatusList = [] # convert particles with all statuses +hepmc_converter.GenParticles.Path = "GenParticles" +ApplicationMgr().TopAlg += [hepmc_converter] + +# Filters events +eventfilter = GenEventFilter("EventFilter") +eventfilter.particles.Path = "GenParticles" +# eventfilter.filterRule = \ +# "bool filterRule(const edm4hep::MCParticleCollection* inColl){" \ +# " return inColl->size() > 1000;}" +eventfilter.filterRulePath = "k4Gen/options/filterRule.hxx" +eventfilter.OutputLevel = DEBUG +ApplicationMgr().TopAlg += [eventfilter] + +out = PodioOutput("out") +out.outputCommands = ["keep *"] +ApplicationMgr().TopAlg += [out] diff --git a/k4Gen/src/components/GenEventFilter.cpp b/k4Gen/src/components/GenEventFilter.cpp new file mode 100644 index 0000000..c0ddc70 --- /dev/null +++ b/k4Gen/src/components/GenEventFilter.cpp @@ -0,0 +1,185 @@ +#include "GenEventFilter.h" + +// Gaudi +#include "GaudiKernel/MsgStream.h" +#include "GaudiKernel/StatusCode.h" +#include "GaudiKernel/ISvcLocator.h" +#include "GaudiKernel/IEventProcessor.h" +#include "GaudiKernel/Incident.h" + +// Datamodel +#include "edm4hep/MCParticleCollection.h" + +// ROOT +#include "TROOT.h" +#include "TSystem.h" +#include "TInterpreter.h" +#include "TGlobal.h" + + +GenEventFilter::GenEventFilter( + const std::string& name, + ISvcLocator* svcLoc) : Gaudi::Algorithm(name, svcLoc) { + declareProperty("particles", m_inColl, + "Generated particles to decide on (input)"); +} + +StatusCode GenEventFilter::initialize() { + { + StatusCode sc = Gaudi::Algorithm::initialize(); + if (sc.isFailure()) { + return sc; + } + } + + m_property = service("ApplicationMgr", m_property); + Gaudi::Property evtMax; + evtMax.assign(m_property->getProperty("EvtMax")); + m_nEventsTarget = evtMax; + debug() << "Targeted number of events: " << m_nEventsTarget << endmsg; + m_nEventsAccepted = 0; + m_nEventsSeen = 0; + + m_incidentSvc = service("IncidentSvc"); + + m_eventProcessor = service("ApplicationMgr"); + + if (m_filterRuleStr.value().empty() && m_filterRulePath.value().empty()) { + error() << "Filter rule not found!" << endmsg; + error() << "Provide it as a string or in the cxx file." << endmsg; + return StatusCode::FAILURE; + } + + if (!m_filterRuleStr.value().empty() && !m_filterRulePath.value().empty()) { + error() << "Multiple ilter rules found!" << endmsg; + error() << "Provide either a string or the cxx file." << endmsg; + return StatusCode::FAILURE; + } + + if (!m_filterRuleStr.value().empty()) { + // Include(s) needed + { + bool success = gInterpreter->Declare( + "#include \"edm4hep/MCParticleCollection.h\""); + if (!success) { + error() << "Unable to find edm4hep::MCParticleCollection header file!" + << endmsg; + return StatusCode::FAILURE; + } + debug() << "Found edm4hep::MCParticleCollection header file." << endmsg; + } + + // Filter rule provided directly as a string + { + bool success = gInterpreter->Declare(m_filterRuleStr.value().c_str()); + if (!success) { + error() << "Unable to compile filter rule!" << endmsg; + return StatusCode::FAILURE; + } + debug() << "Filter rule compiled successfully." << endmsg; + } + } + + if (!m_filterRulePath.value().empty()) { + if (gSystem->AccessPathName(m_filterRulePath.value().c_str())) { + error() << "Unable to access filter rule file!" << endmsg; + error() << "Provided filter rule file path: " << m_filterRulePath.value() << endmsg; + return StatusCode::FAILURE; + } + // Include and compile the file + { + bool success = gInterpreter->Declare( + ("#include \"" + m_filterRulePath + "\"").c_str()); + if (!success) { + error() << "Unable to include filter rule file!" + << endmsg; + return StatusCode::FAILURE; + } + debug() << "Included filter rule file." << endmsg; + } + } + + // Get the address of the compiled filter rule from the interpreter + { + enum TInterpreter::EErrorCode err = TInterpreter::kProcessing; + m_filterRulePtr = + reinterpret_cast( + gInterpreter->ProcessLineSynch("&filterRule", &err)); + if (err != TInterpreter::kNoError) { + error() << "Unable to obtain filter rule pointer!" << endmsg; + return StatusCode::FAILURE; + } + debug() << "Filter rule pointer obtained successfully." << endmsg; + } + + // Check if the filter rule pointer has correct signature + { + auto success = gInterpreter->Declare("auto filterRulePtr = &filterRule;"); + if (!success) { + error() << "Unable to declare filter rule pointer in the interpreter!" + << endmsg; + return StatusCode::FAILURE; + } + auto global = gROOT->GetGlobal("filterRulePtr"); + if (!global) { + error() << "Unable to obtain filter rule pointer from the interpreter!" + << endmsg; + return StatusCode::FAILURE; + } + std::string filterRuleType = global->GetTypeName(); + if (filterRuleType != "bool(*)(const edm4hep::MCParticleCollection*)") { + error() << "Found filter rule pointer has wrong signature!" << endmsg; + error() << "Required: bool(*)(const edm4hep::MCParticleCollection*)" + << endmsg; + error() << "Found: " << filterRuleType << endmsg; + return StatusCode::FAILURE; + } + debug() << "Found filter rule pointer has correct signature." << endmsg; + } + + return StatusCode::SUCCESS; +} + +StatusCode GenEventFilter::execute(const EventContext& evtCtx) const { + const edm4hep::MCParticleCollection* inParticles = m_inColl.get(); + m_nEventsSeen++; + + if (!(*m_filterRulePtr)(inParticles)) { + debug() << "Skipping event..." << endmsg; + + { + StatusCode sc = m_eventProcessor->nextEvent(1); + if (sc.isFailure()) { + error() << "Error when attempting to skip the event! Aborting..." + << endmsg; + return sc; + } + } + + m_incidentSvc->fireIncident(Incident(name(), IncidentType::AbortEvent)); + + return StatusCode::SUCCESS; + } + + m_nEventsAccepted++; + + debug() << "Event contains " << inParticles->size() << " particles." + << endmsg; + debug() << "Targeted number of events to generate: " << m_nEventsTarget + << endmsg; + debug() << "Number of events already generated: " << m_nEventsAccepted + << endmsg; + debug() << "Remaining number of event to generate: " + << m_nEventsTarget - m_nEventsAccepted << endmsg; + debug() << "Number of events seen so far: " << m_nEventsSeen << endmsg; + + return StatusCode::SUCCESS; +} + +StatusCode GenEventFilter::finalize() { + debug() << "Number of events seen: " << m_nEventsSeen << endmsg; + + return Gaudi::Algorithm::finalize(); +} + +DECLARE_COMPONENT(GenEventFilter) diff --git a/k4Gen/src/components/GenEventFilter.h b/k4Gen/src/components/GenEventFilter.h new file mode 100644 index 0000000..7db6e34 --- /dev/null +++ b/k4Gen/src/components/GenEventFilter.h @@ -0,0 +1,67 @@ +#ifndef GENERATION_GENEVENTFILTER_H +#define GENERATION_GENEVENTFILTER_H + +// Gaudi +#include "GaudiKernel/Algorithm.h" +class IProperty; +class IIncidentSvc; +class IEventProcessor; + +// k4FWCore +#include "k4FWCore/DataHandle.h" + +// Datamodel +namespace edm4hep { + class MCParticleCollection; +} + +/** @class GenEventFilter Generation/src/components/GenEventFilter.h GenEventFilter.h + * + * Filters events based on the user defined filter rule applied on MCParticle + * collection. + * + * @author J. Smiesko +*/ + +class GenEventFilter : public Gaudi::Algorithm { + +public: + /// Constructor + GenEventFilter(const std::string& name, ISvcLocator* svcLoc); + /// Initialize + virtual StatusCode initialize(); + /// Execute: Applies the filter + virtual StatusCode execute(const EventContext& evtCtx) const; + /// Finalize + virtual StatusCode finalize(); + +private: + /// Handle for the MCParticle collection to be read + mutable DataHandle m_inColl{ + "particles", Gaudi::DataHandle::Reader, this}; + + /// Rule to filter the events with + Gaudi::Property m_filterRuleStr{ + this, "filterRule", "", "Filter rule to apply on the events"}; + + /// Path of the filter rule file + Gaudi::Property m_filterRulePath{ + this, "filterRulePath", "", "Path to the filter rule file"}; + + /// Targeted number of events. + mutable std::atomic m_nEventsTarget; + /// Keep track of how many events were already accepted. + mutable std::atomic m_nEventsAccepted; + /// Keep track of how many events we went through. + mutable std::atomic m_nEventsSeen; + /// Pointer to the property. + SmartIF m_property; + /// Pointer to the incident service. + SmartIF m_incidentSvc; + /// Pointer to the event processor. + SmartIF m_eventProcessor; + /// Filter rule pointer. + bool (*m_filterRulePtr)(const edm4hep::MCParticleCollection*); +}; + +#endif // GENERATION_GENEVENTFILTER_H diff --git a/k4Gen/src/components/HepMCToEDMConverter.cpp b/k4Gen/src/components/HepMCToEDMConverter.cpp index 06adc13..399a84d 100644 --- a/k4Gen/src/components/HepMCToEDMConverter.cpp +++ b/k4Gen/src/components/HepMCToEDMConverter.cpp @@ -34,6 +34,12 @@ edm4hep::MutableMCParticle HepMCToEDMConverter::convert(std::shared_ptrend_vertex(); + if ( endVtx != nullptr ) { + auto& pos = endVtx->position(); + edm_particle.setEndpoint( {pos.x(), pos.y(), pos.z()} ); + } + return edm_particle; }