Merge pull request #24 from kjvbrt/event-filter

Event filtering algorithm

k4Gen/options/filterRule.hxx

@@ -0,0 +1,5 @@
+#include "edm4hep/MCParticleCollection.h"
+
+bool filterRule(const edm4hep::MCParticleCollection* inColl) {
+  return inColl->size() > 1000;
+}

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]

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<int> 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<bool (*)(const edm4hep::MCParticleCollection*)>(
+            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)

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<edm4hep::MCParticleCollection> m_inColl{
+      "particles", Gaudi::DataHandle::Reader, this};
+
+  /// Rule to filter the events with
+  Gaudi::Property<std::string> m_filterRuleStr{
+      this, "filterRule", "", "Filter rule to apply on the events"};
+
+  /// Path of the filter rule file
+  Gaudi::Property<std::string> m_filterRulePath{
+      this, "filterRulePath", "", "Path to the filter rule file"};
+
+  /// Targeted number of events.
+  mutable std::atomic<unsigned int> m_nEventsTarget;
+  /// Keep track of how many events were already accepted.
+  mutable std::atomic<unsigned int> m_nEventsAccepted;
+  /// Keep track of how many events we went through.
+  mutable std::atomic<unsigned int> m_nEventsSeen;
+  /// Pointer to the property.
+  SmartIF<IProperty> m_property;
+  /// Pointer to the incident service.
+  SmartIF<IIncidentSvc> m_incidentSvc;
+  /// Pointer to the event processor.
+  SmartIF<IEventProcessor> m_eventProcessor;
+  /// Filter rule pointer.
+  bool (*m_filterRulePtr)(const edm4hep::MCParticleCollection*);
+};
+
+#endif  // GENERATION_GENEVENTFILTER_H

k4Gen/src/components/HepMCToEDMConverter.cpp

@@ -34,6 +34,12 @@ edm4hep::MutableMCParticle HepMCToEDMConverter::convert(std::shared_ptr<const He
     edm_particle.setVertex( {float(pos.x()), float(pos.y()), float(pos.z())} );
   }
 
+  auto endVtx = hepmcParticle->end_vertex();
+  if ( endVtx != nullptr ) {
+    auto& pos = endVtx->position();
+    edm_particle.setEndpoint( {pos.x(), pos.y(), pos.z()} );
+  }
+
   return edm_particle;
 }