Added alien path compatibility, Phi rotation and event pool randomisation to GenO2Kine

Common/SimConfig/include/SimConfig/SimConfig.h

@@ -55,6 +55,7 @@ struct SimConfigData {
   std::string mTrigger;                               // chosen VMC generator trigger
   unsigned int mNEvents;                              // number of events to be simulated
   std::string mExtKinFileName;                        // file name of external kinematics file (needed for ext kinematics generator)
+  std::string mKineInput;                             // file name of external kinematics file used as input for the generation (events redirection)
   std::string mEmbedIntoFileName;                     // filename containing the reference events to be used for the embedding
   unsigned int mStartEvent;                           // index of first event to be taken
   float mBMax;                                        // maximum for impact parameter sampling
@@ -154,6 +155,7 @@ class SimConfig
   unsigned int getNEvents() const { return mConfigData.mNEvents; }
 
   std::string getExtKinematicsFileName() const { return mConfigData.mExtKinFileName; }
+  std::string getKineInput() const { return mConfigData.mKineInput; }
   std::string getEmbedIntoFileName() const { return mConfigData.mEmbedIntoFileName; }
   unsigned int getStartEvent() const { return mConfigData.mStartEvent; }
   float getBMax() const { return mConfigData.mBMax; }

Common/SimConfig/src/SimConfig.cxx

@@ -77,6 +77,7 @@ void SimConfig::initOptions(boost::program_options::options_description& options
     "forwardKine", bpo::bool_switch(), "forward kinematics on a FairMQ channel")(
     "noDiscOutput", bpo::bool_switch(), "switch off writing sim results to disc (useful in combination with forwardKine)");
   options.add_options()("fromCollContext", bpo::value<std::string>()->default_value(""), "Use a pregenerated collision context to infer number of events to simulate, how to embedd them, the vertex position etc. Takes precedence of other options such as \"--nEvents\".");
+  options.add_options()("kine-input", bpo::value<std::string>()->default_value(""), "Use a pregenerated kinematic .root file. The optional \':\' character with a number after it can be used to infer number of events to skip on the input file.");
 }
 
 void SimConfig::determineActiveModules(std::vector<std::string> const& inputargs, std::vector<std::string> const& skippedModules, std::vector<std::string>& activeModules, bool isUpgrade)
@@ -304,6 +305,7 @@ bool SimConfig::resetFromParsedMap(boost::program_options::variables_map const&
   mConfigData.mTrigger = vm["trigger"].as<std::string>();
   mConfigData.mNEvents = vm["nEvents"].as<unsigned int>();
   mConfigData.mExtKinFileName = vm["extKinFile"].as<std::string>();
+  mConfigData.mKineInput = vm["kine-input"].as<std::string>();
   mConfigData.mEmbedIntoFileName = vm["embedIntoFile"].as<std::string>();
   mConfigData.mStartEvent = vm["startEvent"].as<unsigned int>();
   mConfigData.mBMax = vm["bMax"].as<float>();

Generators/include/Generators/GeneratorFromFile.h

@@ -16,6 +16,8 @@
 
 #include "FairGenerator.h"
 #include "Generators/Generator.h"
+#include <TRandom3.h>
+#include <TGrid.h>
 
 class TBranch;
 class TFile;
@@ -57,6 +59,7 @@ class GeneratorFromFile : public FairGenerator
   int mEventsAvailable = 0;
   bool mSkipNonTrackable = true; //! whether to pass non-trackable (decayed particles) to the MC stack
   bool mFixOffShell = true;      // fix particles with M_assigned != M_calculated
+  TGrid* mAlienInstance = nullptr; // a cached connection to TGrid (needed for Alien locations)
   ClassDefOverride(GeneratorFromFile, 1);
 };
 
@@ -94,6 +97,10 @@ class GeneratorFromO2Kine : public o2::eventgen::Generator
   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<o2::dataformats::MCEventHeader> mOrigMCEventHeader; //! the MC event header of the original file
 

Generators/include/Generators/GeneratorFromO2KineParam.h

@@ -25,13 +25,16 @@
 /**
  ** a parameter class/struct to keep the settings of
  ** the FromO2Kine event generator and
  ** allow the user to modify them 
  **/
 
 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
+  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
   O2ParamDef(GeneratorFromO2KineParam, "GeneratorFromO2Kine");
 };

Generators/src/GeneratorFromFile.cxx

@@ -29,6 +29,13 @@ namespace eventgen
 {
 GeneratorFromFile::GeneratorFromFile(const char* name)
 {
+  if (strncmp(name, "alien:/", 7) == 0) {
+    mAlienInstance = TGrid::Connect("alien");
+    if (!mAlienInstance) {
+      LOG(fatal) << "Could not connect to alien, did you check the alien token?";
+      return;
+    }
+  }
   mEventFile = TFile::Open(name);
   if (mEventFile == nullptr) {
     LOG(fatal) << "EventFile " << name << " not found \n";
@@ -175,6 +182,13 @@ GeneratorFromO2Kine::GeneratorFromO2Kine(const char* name)
   setPositionUnit(1.);
   setTimeUnit(1.);
 
+  if (strncmp(name, "alien:/", 7) == 0) {
+    mAlienInstance = TGrid::Connect("alien");
+    if (!mAlienInstance) {
+      LOG(fatal) << "Could not connect to alien, did you check the alien token?";
+      return;
+    }
+  }
   mEventFile = TFile::Open(name);
   if (mEventFile == nullptr) {
     LOG(fatal) << "EventFile " << name << " not found";
@@ -210,6 +224,12 @@ bool GeneratorFromO2Kine::Init()
   mSkipNonTrackable = param.skipNonTrackable;
   mContinueMode = param.continueMode;
   mRoundRobin = param.roundRobin;
+  mRandomize = param.randomize;
+  mRngSeed = param.rngseed;
+  mRandomPhi = param.randomphi;
+  if (mRandomize) {
+    gRandom->SetSeed(mRngSeed);
+  }
 
   return true;
 }
@@ -229,6 +249,18 @@ bool GeneratorFromO2Kine::importParticles()
   // It might need some adjustment to make it work with secondaries or to continue
   // from a kinematics snapshot
 
+  // Randomize the order of events in the input file
+  if (mRandomize) {
+    mEventCounter = gRandom->Integer(mEventsAvailable);
+  }
+
+  double dPhi = 0.;
+  // Phi rotation
+  if (mRandomPhi) {
+    dPhi = gRandom->Uniform(2 * TMath::Pi());
+    LOG(info) << "Rotating phi by " << dPhi;
+  }
+
   if (mEventCounter < mEventsAvailable) {
     int particlecounter = 0;
 
@@ -253,6 +285,13 @@ bool GeneratorFromO2Kine::importParticles()
       auto pdg = t.GetPdgCode();
       auto px = t.Px();
       auto py = t.Py();
+      if (mRandomPhi) {
+        auto phi = TMath::ATan2(py, px);
+        auto pt = TMath::Sqrt(px * px + py * py);
+        phi += dPhi;
+        px = pt * TMath::Cos(phi);
+        py = pt * TMath::Sin(phi);
+      }
       auto pz = t.Pz();
       auto vx = t.Vx();
       auto vy = t.Vy();

run/o2sim_parallel.cxx

@@ -512,6 +512,46 @@ int main(int argc, char* argv[])
     }
     return r;
   }
+  // check if Kinematics input file is provided
+  LOG(info) << "Kinematics input file: " << conf.getKineInput();
+  if (conf.getKineInput().size() > 0) {
+    std::string path = conf.getKineInput();
+    std::size_t pos = path.find(':');
+    std::string kinefilename;
+    int ntf = 1;
+    if (pos != std::string::npos) {
+      kinefilename = path.substr(0, pos);
+      ntf = std::stoi(path.substr(pos + 1));
+    } else {
+      kinefilename = path;
+    }
+    LOG(info) << "Using Kinematics input file: " << kinefilename << ", TimeFrame number: " << ntf;
+    TFile kinefile(kinefilename.c_str(), "READ");
+    if (kinefile.IsZombie()) {
+      LOG(fatal) << "Could not open Kinematics input file: " << kinefilename;
+      return 2;
+    }
+    TFile splitkinefile(Form("%s_Kine.root", conf.getOutPrefix().c_str()), "RECREATE");
+    TTree* kine = (TTree*)kinefile.Get("o2sim");
+    const int nentries = kine->GetEntries();
+    const int nevsim = conf.getNEvents();
+    TTree* kineout = kine->CloneTree(0);
+    ntf -= 1;
+    if ((ntf + 1) * nevsim > kine->GetEntries()) {
+      LOG(fatal) << "Not enough events in the input file to get the entries in the interval [" << ntf * nevsim << ',' << (ntf + 1) * nevsim << "]";
+      return 2;
+    }
+    for (int i = ntf * nevsim; i < (ntf + 1) * nevsim; ++i) {
+      kine->GetEntry(i);
+      kineout->Fill();
+    }
+    splitkinefile.Write();
+    kinefile.Close();
+    splitkinefile.Close();
+    LOG(info) << "Kinematics input file processed successfully";
+    LOG(info) << "Splitting time took " << timer.RealTime() << " s";
+    return 0;
+  }
 
   gAskedEvents = conf.getNEvents();
   if (conf.asService()) {