Implement HepMC reading randomisation

Author: jackal1-66

Generators/CMakeLists.txt

@@ -139,6 +139,14 @@ if(doBuildSimulation)
               LABELS generator
               PUBLIC_LINK_LIBRARIES O2::Generators)
 
+  if(HepMC3_FOUND)
+    o2_add_test(GeneratorHepMCIndexed NAME test_Generator_test_GeneratorHepMCIndexed
+                SOURCES test/test_GeneratorHepMCIndexed.cxx
+                COMPONENT_NAME Generator
+                LABELS generator
+                PUBLIC_LINK_LIBRARIES O2::Generators)
+  endif()
+
   # o2_add_test(GeneratorPythia8Param NAME test_Generator_test_GeneratorPythia8Param
   #             SOURCES test/test_GeneratorPythia8Param.cxx
   #             COMPONENT_NAME Generator

Generators/include/Generators/GeneratorHepMC.h

@@ -18,6 +18,10 @@
 #include "Generators/GeneratorFileOrCmd.h"
 #include "Generators/GeneratorHepMCParam.h"
 #include "Generators/GeneratorFileOrCmdParam.h"
+#include <iosfwd>
+#include <memory>
+#include <string>
+#include <vector>
 
 #ifdef GENERATORS_WITH_HEPMC3_DEPRECATED
 namespace HepMC
@@ -68,7 +72,7 @@ class GeneratorHepMC : public Generator, public GeneratorFileOrCmd
    * simulation configuration.  This is implemented as a member
    * function so as to better facilitate changes. */
   void setup(const GeneratorFileOrCmdParam& param0,
-             const GeneratorHepMCParam& param,
+             const HepMCGenConfig& param,
              const conf::SimConfig& config);
   // Generator configuration from external local parameters
   void setup(const FileOrCmdGenConfig& param0,
@@ -108,8 +112,19 @@ class GeneratorHepMC : public Generator, public GeneratorFileOrCmd
 
   /** methods that can be overridded **/
   void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override;
-  /** Make our reader */
+  /** Apply the HepMC-specific configuration */
+  void setupHepMC(const HepMCGenConfig& param);
+  /** Make our reader, taking the next file off the list of file names */
   bool makeReader();
+  /** Fix the order in which the entries of the input file are served */
+  void establishEventOrder();
+  /** Index the events of a file by byte offset and open the reader on it, so
+   *  that any entry can later be reached with a single seek. Available only for ASCII format */
+  bool buildIndex(const std::string& filename);
+  /** Read the given entry of the indexed file */
+  bool readEntry(int entry);
+  /** Generate an event following the established event order in random mode */
+  Bool_t generateEventOrdered();
 
   /** Type of function to select particles to keep when pruning
    * events */
@@ -127,8 +142,34 @@ class GeneratorHepMC : public Generator, public GeneratorFileOrCmd
   HepMC3::GenEvent* mEvent = nullptr;
   /** Option whether to prune event */
   bool mPrune; //!
-
-  ClassDefOverride(GeneratorHepMC, 1);
+  /** Name of the file the reader is attached to, needed to re-open it */
+  std::string mCurrentFileName; //!
+  /** Order in which the entries of the input file are served */
+  std::vector<int> mEventOrder; //!
+  /** Events already delivered in the current pass over the file */
+  int mEventCounter = 0; //!
+  /** Events delivered in total */
+  int mEventsServed = 0; //!
+  /** Events contained in the input file */
+  int mEventsAvailable = 0; //!
+  /** Entry the reader is currently positioned on, -1 if none */
+  int mLastEntryRead = -1; //!
+  /** Option whether to serve the events in random order */
+  bool mRandomize = false; //!
+  /** Option whether to start over once all events have been used */
+  bool mRoundRobin = false; //!
+  /** Option to have a new order when round-robin enabled */
+  bool mReshuffleOnRepeat = true; //!
+  /** Randomizer seed, 0 to leave gRandom alone */
+  unsigned int mRngSeed = 0; //!
+  /** Whether the indexed input is HepMC2 rather than HepMC3 ASCII */
+  bool mIndexedHepMC2 = false; //!
+  /** The stream the reader is attached to, ours so that we may seek in it */
+  std::shared_ptr<std::istream> mIndexedStream; //!
+  /** Byte offset at which every entry of the input file starts */
+  std::vector<std::streamoff> mEventOffsets; //!
+
+  ClassDefOverride(GeneratorHepMC, 2);
 
 }; /** class GeneratorHepMC **/
 

Generators/include/Generators/GeneratorHepMCParam.h

@@ -29,7 +29,7 @@ namespace eventgen
  ** allow the user to modify them
  **/
 
-struct GeneratorHepMCParam : public o2::conf::ConfigurableParamHelper<GeneratorHepMCParam> {
+struct HepMCGenConfig {
   /** Version number of event structure to decode.  Note, when reading
    *  from a file, this key is ignored.  The interface will figure out
    *  the version automatically.  When reading from a pipe, and the
@@ -51,15 +51,19 @@ struct GeneratorHepMCParam : public o2::conf::ConfigurableParamHelper<GeneratorH
    * event generator producing the event.  Use with caution, as it may
    * corrupt the event record. */
   bool prune = false;
-  O2ParamDef(GeneratorHepMCParam, "HepMC");
+  /** Serve once the events of the input file in random order */
+  bool randomize = false;
+  /** Restart reading events */
+  bool roundRobin = false;
+  /** Draw a fresh random order on each new pass over the input file */
+  bool reshuffleOnRepeat = true;
+  /** Randomizer seed, 0 for random value. */
+  unsigned int rngseed = 0;
 };
 
-struct HepMCGenConfig {
-  // Same parameters as GeneratorHepMCParam
-  int version = 0;
-  uint64_t eventsToSkip = 0;
-  std::string fileName = "";
-  bool prune = false;
+// construct a configurable param singleton out of the HepMCGenConfig struct
+struct GeneratorHepMCParam : public o2::conf::ConfigurableParamPromoter<GeneratorHepMCParam, HepMCGenConfig> {
+  O2ParamDef(GeneratorHepMCParam, "HepMC");
 };
 
 } // end namespace eventgen