feat: Add random sampling to HepMC3 reader

Examples/Io/HepMC3/include/ActsExamples/Io/HepMC3/HepMC3Reader.hpp

@@ -15,6 +15,7 @@
 #include "ActsExamples/Utilities/VertexGenerators.hpp"
 
 #include <filesystem>
+#include <memory>
 #include <mutex>
 #include <string>
 
@@ -25,20 +26,97 @@ class Reader;
 
 namespace ActsExamples {
 
+struct MultiplicityGenerator;
+
 /// HepMC3 event reader.
+///
+/// This reader supports reading events from one or more HepMC3 files, with
+/// flexible control over how many physical events are merged into each logical
+/// event.
+///
+/// ## Multiplicity Generators
+///
+/// Each input file must have an associated multiplicity generator that
+/// determines how many physical events to read from that file for each logical
+/// event:
+/// - FixedMultiplicityGenerator: Always reads a fixed number of events
+/// (deterministic)
+/// - PoissonMultiplicityGenerator: Reads a Poisson-distributed number of events
+/// (stochastic)
+///
+/// ## Usage Patterns
+///
+/// 1. Single file with default multiplicity:
+///    ```cpp
+///    HepMC3Reader::Config cfg;
+///    cfg.inputPath = "events.hepmc3";
+///    cfg.outputEvent = "hepmc3_event";
+///    HepMC3Reader reader(cfg, Acts::Logging::INFO);
+///    ```
+///    Automatically uses FixedMultiplicityGenerator(n=1).
+///
+/// 2. Multiple files with fixed multiplicities (e.g., hard-scatter + pileup):
+///    ```cpp
+///    HepMC3Reader::Config cfg;
+///    cfg.inputs = {
+///      {.path = "signal.hepmc3",
+///       .multiplicityGenerator =
+///       std::make_shared<FixedMultiplicityGenerator>(1)},
+///      {.path = "pileup.hepmc3",
+///       .multiplicityGenerator =
+///       std::make_shared<FixedMultiplicityGenerator>(50)}
+///    };
+///    cfg.outputEvent = "hepmc3_event";
+///    HepMC3Reader reader(cfg, Acts::Logging::INFO);
+///    ```
+///
+/// 3. With stochastic pileup multiplicity:
+///    ```cpp
+///    HepMC3Reader::Config cfg;
+///    cfg.inputs = {
+///      {.path = "signal.hepmc3",
+///       .multiplicityGenerator =
+///       std::make_shared<FixedMultiplicityGenerator>(1)},
+///      {.path = "pileup.hepmc3",
+///       .multiplicityGenerator =
+///       std::make_shared<PoissonMultiplicityGenerator>(50.0)}
+///    };
+///    cfg.outputEvent = "hepmc3_event";
+///    cfg.randomNumbers = rng; // Required for stochastic generators
+///    HepMC3Reader reader(cfg, Acts::Logging::INFO);
+///    ```
+///
+/// ## Event Skipping
+///
+/// Event skipping (via Sequencer's skip parameter) is only supported when ALL
+/// multiplicity generators are Fixed. This is because skipping requires knowing
+/// exactly how many physical events to skip, which is only deterministic with
+/// FixedMultiplicityGenerator.
+///
 class HepMC3Reader final : public IReader {
  public:
+  /// Input specification per file
+  struct Input {
+    /// Path to the HepMC3 file
+    std::filesystem::path path;
+    /// Multiplicity generator determining how many events to read per logical
+    /// event. Must always be set. Use FixedMultiplicityGenerator for
+    /// deterministic behavior, or PoissonMultiplicityGenerator for stochastic
+    /// pileup simulation.
+    std::shared_ptr<const MultiplicityGenerator> multiplicityGenerator;
+  };
+
   struct Config {
-    /// The input file path for reading HepMC3 events.
-    /// This is a helper to simplify the most basic configuration, which is
-    /// reading single events from a single input file
-    std::filesystem::path inputPath;
+    /// Input files to read. For each file, the multiplicity generator
+    /// determines how many events are read per logical event. This can be used
+    /// to read e.g. hard-scatter events from one file and pileup events from
+    /// another. Mutually exclusive with inputPath.
+    std::vector<Input> inputs;
 
-    /// This configuration option is used to read multiple files in a single
-    /// run. For each file, a specific number of events is read per requested
-    /// event. This can be used to read e.g. hard-scatter events from one file
-    /// and pileup events from another.
-    std::vector<std::pair<std::filesystem::path, std::size_t>> inputPaths;
+    /// Convenience parameter for single-file reading. Mutually exclusive with
+    /// inputs. If set, creates a single input with
+    /// FixedMultiplicityGenerator(n=1).
+    std::optional<std::filesystem::path> inputPath;
 
     /// The output collection
     std::string outputEvent;
@@ -62,7 +140,11 @@ class HepMC3Reader final : public IReader {
     /// used. If this number is exceeded the reader will error out.
     std::size_t maxEventBufferSize = 128;
 
-    /// The random number service. Required if vertexGenerator is set.
+    /// The random number service. Required if:
+    /// - vertexGenerator is set, OR
+    /// - any non-Fixed multiplicityGenerator is used (e.g.,
+    /// PoissonMultiplicityGenerator) Not required when using only
+    /// FixedMultiplicityGenerator with no vertexGenerator.
     std::shared_ptr<const RandomNumbers> randomNumbers;
     /// Position generator that will be used to shift read events
     std::shared_ptr<PrimaryVertexPositionGenerator> vertexGenerator;
@@ -134,10 +216,18 @@ class HepMC3Reader final : public IReader {
 
   std::mutex m_queueMutex;
 
+  /// Precomputed flag: true if RNG is needed (vertex generator or non-Fixed
+  /// multiplicity)
+  bool m_needsRng = false;
+
+  /// Precomputed flag: true if any input uses a non-Fixed multiplicity
+  /// generator
+  bool m_hasNonFixedMultiplicity = false;
+
   struct InputConfig {
     std::shared_ptr<HepMC3::Reader> reader;
-    std::size_t numEvents;
     std::filesystem::path path;
+    std::shared_ptr<const MultiplicityGenerator> multiplicityGenerator;
   };
 
   std::vector<InputConfig> m_inputs;

Examples/Io/HepMC3/src/HepMC3Reader.cpp

@@ -15,6 +15,7 @@
 #include "Acts/Utilities/ThrowAssert.hpp"
 #include "ActsExamples/Framework/ProcessCode.hpp"
 #include "ActsExamples/Io/HepMC3/HepMC3Util.hpp"
+#include "ActsExamples/Utilities/MultiplicityGenerators.hpp"
 
 #include <filesystem>
 #include <memory>
@@ -38,23 +39,41 @@ HepMC3Reader::HepMC3Reader(const HepMC3Reader::Config& cfg,
 
   m_outputEvent.initialize(m_cfg.outputEvent);
 
-  if (!m_cfg.inputPaths.empty() && !m_cfg.inputPath.empty()) {
-    throw std::invalid_argument(
-        "inputPath and inputPaths are mutually exclusive");
-  }
+  // Validate: exactly one of inputPath or inputs must be set
+  bool hasInputPath = m_cfg.inputPath.has_value();
+  bool hasInputs = !m_cfg.inputs.empty();
 
-  if (!m_cfg.inputPath.empty()) {
-    m_cfg.inputPaths.emplace_back(m_cfg.inputPath, 1);
+  if (!hasInputPath && !hasInputs) {
+    throw std::invalid_argument(
+        "HepMC3 reader requires either 'inputPath' or 'inputs' to be set");
   }
 
-  if (m_cfg.inputPaths.empty()) {
+  if (hasInputPath && hasInputs) {
     throw std::invalid_argument(
-        "HepMC3 reader was not configured with any input files");
+        "HepMC3 reader: 'inputPath' and 'inputs' are mutually exclusive. "
+        "Use 'inputPath' for single file or 'inputs' for multiple files.");
   }
 
-  for (const auto& [path, numEvents] : m_cfg.inputPaths) {
-    auto reader = HepMC3Util::deduceReader(path);
-    m_inputs.emplace_back(reader, numEvents, path);
+  // If inputPath is set, create a single input with default multiplicity
+  // generator
+  if (hasInputPath) {
+    Input input;
+    input.path = m_cfg.inputPath.value();
+    input.multiplicityGenerator =
+        std::make_shared<FixedMultiplicityGenerator>(1);
+
+    auto reader = HepMC3Util::deduceReader(input.path);
+    m_inputs.emplace_back(reader, input.path, input.multiplicityGenerator);
+  } else {
+    // Use the provided inputs
+    for (const auto& input : m_cfg.inputs) {
+      if (!input.multiplicityGenerator) {
+        throw std::invalid_argument(
+            "All Input objects must have a multiplicityGenerator set");
+      }
+      auto reader = HepMC3Util::deduceReader(input.path);
+      m_inputs.emplace_back(reader, input.path, input.multiplicityGenerator);
+    }
   }
 
   if (m_cfg.numEvents.has_value()) {
@@ -69,9 +88,24 @@ HepMC3Reader::HepMC3Reader(const HepMC3Reader::Config& cfg,
     m_eventsRange = {0, determineNumEvents(*reader)};
   }
 
-  if (m_cfg.vertexGenerator != nullptr && m_cfg.randomNumbers == nullptr) {
+  // Check if any input uses a non-Fixed multiplicity generator
+  m_hasNonFixedMultiplicity =
+      std::ranges::any_of(m_inputs, [](const auto& input) {
+        const auto& [reader, path, multiplicityGenerator] = input;
+        // Check if this is NOT a FixedMultiplicityGenerator
+        return dynamic_cast<const FixedMultiplicityGenerator*>(
+                   multiplicityGenerator.get()) == nullptr;
+      });
+
+  // Check if randomNumbers is required
+  // RNG is needed if we have a vertex generator or any non-Fixed multiplicity
+  // generator
+  m_needsRng = (m_cfg.vertexGenerator != nullptr) || m_hasNonFixedMultiplicity;
+
+  if (m_needsRng && m_cfg.randomNumbers == nullptr) {
     throw std::invalid_argument(
-        "randomNumbers must be set if vertexGenerator is set");
+        "randomNumbers must be set if vertexGenerator or any non-Fixed "
+        "multiplicityGenerator is used");
   }
 
   ACTS_DEBUG("HepMC3Reader: " << m_eventsRange.first << " - "
@@ -116,18 +150,39 @@ ProcessCode HepMC3Reader::skip(std::size_t events) {
     return SUCCESS;
   }
 
-  ACTS_DEBUG("Skipping " << events << " events");
+  ACTS_DEBUG("Skipping " << events << " logical events");
 
-  for (const auto& [reader, numEvents, path] : m_inputs) {
-    ACTS_VERBOSE("Skipping " << events << "*" << numEvents << "="
-                             << events * numEvents << " events from " << path);
-    if (!reader->skip(static_cast<int>(events * numEvents))) {
-      ACTS_ERROR("Error skipping events " << events << " " << path);
-      return ABORT;
+  // Check if all multiplicity generators are Fixed (deterministic)
+  // Use the precomputed flag
+  if (m_hasNonFixedMultiplicity) {
+    ACTS_ERROR(
+        "Cannot skip events with non-Fixed multiplicityGenerator (e.g., "
+        "PoissonMultiplicityGenerator). Skipping requires knowing the exact "
+        "number of physical events to skip from each input file, which is only "
+        "possible with deterministic (Fixed) multiplicity generators.");
+    return ABORT;
+  }
+
+  // For each logical event to skip, evaluate the Fixed multiplicity generators
+  // to determine how many physical events to skip from each input file
+  for (std::size_t logicalEvent = 0; logicalEvent < events; ++logicalEvent) {
+    for (const auto& [reader, path, multiplicityGenerator] : m_inputs) {
+      // Must be FixedMultiplicityGenerator (checked above), so downcast is safe
+      const auto* fixedGen = static_cast<const FixedMultiplicityGenerator*>(
+          multiplicityGenerator.get());
+      std::size_t count = fixedGen->n;
+
+      ACTS_VERBOSE("Skipping " << count << " events from " << path
+                               << " for logical event "
+                               << (m_nextEvent + logicalEvent));
+      if (!reader->skip(static_cast<int>(count))) {
+        ACTS_ERROR("Error skipping " << count << " events from " << path);
+        return ABORT;
+      }
     }
   }
 
-  m_nextEvent = events;
+  m_nextEvent += events;
 
   return SUCCESS;
 }
@@ -334,14 +389,43 @@ ProcessCode HepMC3Reader::readLogicalEvent(
 
   // @TODO: Add the index as an attribute to the event and it's content
 
-  for (const auto& [reader, numEvents, path] : m_inputs) {
-    ACTS_VERBOSE("Reading " << numEvents << " events from " << path);
-    for (std::size_t i = 0; i < numEvents; ++i) {
+  // Spawn RNG for multiplicity generators if needed
+  std::optional<RandomEngine> rng;
+  if (m_needsRng) {
+    rng = m_cfg.randomNumbers->spawnGenerator(ctx);
+  }
+
+  for (std::size_t inputIndex = 0; inputIndex < m_inputs.size(); ++inputIndex) {
+    auto& reader = m_inputs[inputIndex].reader;
+    auto& path = m_inputs[inputIndex].path;
+    auto& multiplicityGenerator = m_inputs[inputIndex].multiplicityGenerator;
+
+    // Use multiplicityGenerator to determine count
+    std::size_t count = 0;
+    if (rng.has_value()) {
+      count = (*multiplicityGenerator)(*rng);
+    } else {
+      // Must be FixedMultiplicityGenerator if no RNG, so downcast is safe
+      const auto& fixedGen = static_cast<const FixedMultiplicityGenerator&>(
+          *multiplicityGenerator);
+      count = fixedGen.n;
+    }
+
+    ACTS_VERBOSE("Reading " << count << " events from " << path);
+    for (std::size_t i = 0; i < count; ++i) {
       auto event = makeEvent();
 
       reader->read_event(*event);
       if (reader->failed()) {
-        ACTS_ERROR("Error reading event " << i << " from " << path);
+        ACTS_ERROR("Error reading event " << i << " (input index = "
+                                          << inputIndex << ") from " << path);
+        if (inputIndex > 0) {
+          ACTS_ERROR("-> since this is input file index "
+                     << inputIndex
+                     << ", this probably means that the "
+                        "input file has "
+                        "fewer events than expected.");
+        }
         return ABORT;
       }
       events.push_back(std::move(event));
@@ -355,7 +439,7 @@ ProcessCode HepMC3Reader::readLogicalEvent(
 
 ProcessCode HepMC3Reader::finalize() {
   ACTS_VERBOSE("Closing " << m_inputs.size() << " input files");
-  for (const auto& [reader, numEvents, path] : m_inputs) {
+  for (const auto& [reader, path, multiplicityGenerator] : m_inputs) {
     reader->close();
   }