Implement pools chaining via TChain + example and test macro

Author: Marco Giacalone

Generators/CMakeLists.txt

@@ -139,6 +139,12 @@ if(doBuildSimulation)
               LABELS generator
               PUBLIC_LINK_LIBRARIES O2::Generators)
 
+  o2_add_test(EventPoolChaining NAME test_Generator_test_EventPoolChaining
+              SOURCES test/test_EventPoolChaining.cxx
+              COMPONENT_NAME Generator
+              LABELS generator
+              PUBLIC_LINK_LIBRARIES O2::Generators)
+
   # o2_add_test(GeneratorPythia8Param NAME test_Generator_test_GeneratorPythia8Param
   #             SOURCES test/test_GeneratorPythia8Param.cxx
   #             COMPONENT_NAME Generator

Generators/include/Generators/GeneratorFromFile.h

@@ -21,6 +21,8 @@
 #include <TRandom3.h>
 #include <TGrid.h>
 #include <random>
+#include <string>
+#include <vector>
 
 class TBranch;
 class TFile;
@@ -66,13 +68,18 @@ class GeneratorFromFile : public FairGenerator
 };
 
 /// This class implements a generic FairGenerator which
-/// reads the particles from an external O2 sim kinematics file.
+/// reads the particles from one or more external O2 sim kinematics files.
 class GeneratorFromO2Kine : public o2::eventgen::Generator
 {
  public:
   GeneratorFromO2Kine() = default;
+  /// name may be a single file or a comma-separated list of files to be read one after the other
   GeneratorFromO2Kine(const char* name);
+  GeneratorFromO2Kine(std::vector<std::string> const& filenames);
   GeneratorFromO2Kine(O2KineGenConfig const& pars);
+  /// same as above but with an explicit list of files. Used for event pools
+  GeneratorFromO2Kine(O2KineGenConfig const& pars, std::vector<std::string> const& filenames);
+  ~GeneratorFromO2Kine() override;
 
   bool Init() override;
 
@@ -91,24 +98,58 @@ class GeneratorFromO2Kine : public o2::eventgen::Generator
   void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override;
   const o2::dataformats::MCEventHeader* getOrigMCEventHeader() const { return mOrigMCEventHeader.get(); }
 
+  /// number of events available in the file that is currently open
+  int getEventsAvailable() const { return mEventsAvailable; }
+  /// number of input files this generator can go through
+  int getNumberOfFiles() const { return (int)mFileNames.size(); }
+  /// index (within the file list) of the file currently open, -1 if none
+  int getCurrentFileIndex() const { return mCurrentFileIndex; }
+  /// name of the file currently open, empty if none
+  std::string getCurrentFileName() const;
+  /// number of opened files so far (including the current one)
+  int getNumberOfFilesUsed() const { return mFilesUsed; }
+  /// total number of events delivered so far
+  int getEventsServed() const { return mEventsServed; }
+
+  /// helper splitting a comma-separated list of file names into its components
+  static std::vector<std::string> splitFileNames(std::string const& filenames);
+
  private:
-  TFile* mEventFile = nullptr;     //! the file containing the persistent events
-  TBranch* mEventBranch = nullptr; //! the branch containing the persistent events
-  TBranch* mMCHeaderBranch = nullptr; //! branch containing MC event headers
-  int mEventCounter = 0;
-  int mEventsAvailable = 0;
-  bool mSkipNonTrackable = true; //! whether to pass non-trackable (decayed particles) to the MC stack
-  bool mContinueMode = false;    //! whether we want to continue simulation of previously inhibited tracks
-  bool mRoundRobin = false;      //! whether we want to take events from file in a round robin fashion
-  bool mRandomize = false;       //! whether we want to randomize the order of events in the input file
-  unsigned int mRngSeed = 0;     //! randomizer seed, 0 for random value
-  bool mRandomPhi = false;       //! whether we want to randomize the phi angle of the particles
-  TGrid* mAlienInstance = nullptr; // a cached connection to TGrid (needed for Alien locations)
+  /// closes the currently opened file currently
+  void closeCurrentFile();
+  /// fixes the order in which the events of the current file are served
+  void establishEventOrder();
+  /// opens the file at the given index of the file list.
+  /// Returns false in case the file cannot be used
+  bool openFile(int index);
+  /// moves on to the next usable file of the list; wraps around in round robin mode.
+  /// Returns false when no further file is available
+  bool openNextFile(bool wrapAround);
+
+  std::vector<std::string> mFileNames;      //! the list of input files, read one after the other
+  int mCurrentFileIndex = -1;               //! index of the file currently open
+  int mFilesUsed = 0;                       //! how many files have been opened so far
+  TFile* mCurrentFile = nullptr;            //! the file currently open
+  TBranch* mEventBranch = nullptr;          //! the branch containing the persistent events
+  TBranch* mMCHeaderBranch = nullptr;       //! branch containing MC event headers
+  std::vector<int> mEventOrder;             //! order in which the entries of the current file are served
+  int mEventCounter = 0;                    //! events already delivered from the current file
+  int mEventsServed = 0;                    //! events delivered in total, across all files
+  int mEventsAvailable = 0;                 //! events contained in the current file
+  int mStartEvent = 0;                      //! event to start from in the very first file
+  int mLastEntryRead = -1;                  //! entry of the current event within the current file
+  bool mSkipNonTrackable = true;            //! whether to pass non-trackable (decayed particles) to the MC stack
+  bool mContinueMode = false;               //! whether we want to continue simulation of previously inhibited tracks
+  bool mRoundRobin = false;                 //! whether we want to take events from file in a round robin fashion
+  bool mRandomize = false;                  //! whether we want to randomize the order of events in the input file
+  unsigned int mRngSeed = 0;                //! randomizer seed, 0 for random value
+  bool mRandomPhi = false;                  //! whether we want to randomize the phi angle of the particles
+  TGrid* mAlienInstance = nullptr;          // a cached connection to TGrid (needed for Alien locations)
   std::unique_ptr<O2KineGenConfig> mConfig; //! Configuration object
 
   std::unique_ptr<o2::dataformats::MCEventHeader> mOrigMCEventHeader; //! the MC event header of the original file
 
-  ClassDefOverride(GeneratorFromO2Kine, 2);
+  ClassDefOverride(GeneratorFromO2Kine, 3);
 };
 
 /// Special generator for event pools.
@@ -167,15 +208,24 @@ class GeneratorFromEventPool : public o2::eventgen::Generator
 
   std::vector<std::string> const& getFileUniverse() const { return mPoolFilesAvailable; }
 
+  /// the file universe, in the order this generator instance will go through it
+  std::vector<std::string> const& getChosenFiles() const { return mFilesChosen; }
+
+  /// shuffles the given universe of pool files into the order this instance will use
+  std::vector<std::string> selectFiles(std::vector<std::string> const& universe);
+
+  /// access to the underlying kinematics generator
+  o2::eventgen::GeneratorFromO2Kine const* getO2KineGenerator() const { return mO2KineGenerator.get(); }
+
  private:
   EventPoolGenConfig mConfig;                                                    //! Configuration object
   std::unique_ptr<o2::eventgen::GeneratorFromO2Kine> mO2KineGenerator = nullptr; //! actual generator doing the work
   std::vector<std::string> mPoolFilesAvailable;                                  //! container keeping the collection of files in the event pool
-  std::string mFileChosen;                                                       //! the file chosen for the pool
+  std::vector<std::string> mFilesChosen;                                         //! the file(s) chosen from the pool
   // random number generator to determine a concrete file name
   std::mt19937 mRandomEngine; //!
 
-  ClassDefOverride(GeneratorFromEventPool, 1);
+  ClassDefOverride(GeneratorFromEventPool, 2);
 };
 
 } // end namespace eventgen

Generators/include/Generators/GeneratorFromO2KineParam.h

@@ -31,31 +31,33 @@ namespace eventgen
 struct GeneratorFromO2KineParam : public o2::conf::ConfigurableParamHelper<GeneratorFromO2KineParam> {
   bool skipNonTrackable = true;
   bool continueMode = false;
-  bool roundRobin = false;   // read events with period boundary conditions
-  bool randomize = false;    // randomize the order of events
+  bool roundRobin = false;   // start over from the first file/event once all events have been used
+  bool randomize = false;    // serve the events of each file in random order (each one exactly once)
   unsigned int rngseed = 0;  // randomizer seed, 0 for random value
   bool randomphi = false;    // randomize phi angle
-  std::string fileName = ""; // filename to read from - takes precedence over SimConfig if given
+  std::string fileName = ""; // filename(s) to read from - takes precedence over SimConfig if given;
+                             // a comma-separated list of files is read one file after the other
   O2ParamDef(GeneratorFromO2KineParam, "GeneratorFromO2Kine");
 };
 
 struct O2KineGenConfig {
   bool skipNonTrackable = true;
   bool continueMode = false;
-  bool roundRobin = false;   // read events with period boundary conditions
-  bool randomize = false;    // randomize the order of events
+  bool roundRobin = false;   // start over from the first file/event once all events have been used
+  bool randomize = false;    // serve the events of each file in random order (each one exactly once)
   unsigned int rngseed = 0;  // randomizer seed, 0 for random value
   bool randomphi = false;    // randomize phi angle
-  std::string fileName = ""; // filename to read from - takes precedence over SimConfig if given
+  std::string fileName = ""; // filename(s) to read from - takes precedence over SimConfig if given;
+                             // a comma-separated list of files is read one file after the other
 };
 
 struct EventPoolGenConfig {
   std::string eventPoolPath = ""; // In that order: The path where an event pool can be found ;
                                   // or .. a local file containing a list of files to use
                                   // or .. a concrete file path to a kinematics file
   bool skipNonTrackable = true;   // <--- do we need this?
-  bool roundRobin = false;        // read events with period boundary conditions
-  bool randomize = true;          // randomize the order of events
+  bool roundRobin = false;        // start over from the first file/event once all events have been used
+  bool randomize = true;          // serve the events of each file in random order (each one exactly once)
   unsigned int rngseed = 0;       // randomizer seed, 0 for random value
   bool randomphi = false;         // randomize phi angle; rotates tracks in events by some phi-angle
 };