Implementation of cocktail generation

Works for sequence generation. SubFractions JSON syntax is implemented,
but it does not work yet during the generation.

* Updated Hybrid example

Author: jackal1-66

Generators/include/Generators/GeneratorHybrid.h

@@ -98,6 +98,11 @@ class GeneratorHybrid : public Generator
   int mEventCounter = 0;
   int mTasksStarted = 0;
 
+  // Cocktail mode
+  bool mCocktailMode = false;
+  std::vector<std::vector<int>> mGroups;
+  std::vector<std::vector<int>> mCocktailFractions;
+
   // Create a task arena with a specified number of threads
   std::thread mTBBTaskPoolRunner;
   tbb::concurrent_bounded_queue<int> mInputTaskQueue;

Generators/src/GeneratorHybrid.cxx

@@ -12,7 +12,7 @@
 #include "Generators/GeneratorHybrid.h"
 #include <fairlogger/Logger.h>
 #include <algorithm>
-
+#include <unordered_set>
 #include <tbb/concurrent_queue.h>
 #include <tbb/task_arena.h>
 #include <tbb/parallel_for.h>
@@ -42,9 +42,16 @@ GeneratorHybrid::GeneratorHybrid(const std::string& inputgens)
   }
   int index = 0;
   if (!(mRandomize || mGenerationMode == GenMode::kParallel)) {
-    if (mFractions.size() != mInputGens.size()) {
-      LOG(fatal) << "Number of fractions does not match the number of generators";
-      return;
+    if (mCocktailMode) {
+      if (mGroups.size() != mFractions.size()) {
+        LOG(fatal) << "Number of groups does not match the number of fractions";
+        return;
+      }
+    } else {
+      if (mFractions.size() != mInputGens.size()) {
+        LOG(fatal) << "Number of fractions does not match the number of generators";
+        return;
+      }
     }
     // Check if all elements of mFractions are 0
     if (std::all_of(mFractions.begin(), mFractions.end(), [](int i) { return i == 0; })) {
@@ -303,7 +310,7 @@ bool GeneratorHybrid::generateEvent()
           }
         }
       } else {
-        mIndex = gRandom->Integer(mGens.size());
+        mIndex = gRandom->Integer(mFractions.size());
       }
     } else {
       while (mFractions[mCurrentFraction] == 0 || mseqCounter == mFractions[mCurrentFraction]) {
@@ -322,32 +329,55 @@ bool GeneratorHybrid::generateEvent()
 bool GeneratorHybrid::importParticles()
 {
   int genIndex = -1;
+  std::vector<int> subGenIndex = {};
   if (mIndex == -1) {
     // this means parallel mode ---> we have a common queue
     mResultQueue[0].pop(genIndex);
   } else {
     // need to pop from a particular queue
-    mResultQueue[mIndex].pop(genIndex);
+    if (!mCocktailMode) {
+      mResultQueue[mIndex].pop(genIndex);
+    } else {
+      // in cocktail mode we need to pop from the group queue
+      subGenIndex.resize(mGroups[mIndex].size());
+      for (size_t pos = 0; pos < mGroups[mIndex].size(); ++pos) {
+        int subIndex = mGroups[mIndex][pos];
+        LOG(info) << "Getting generator " << mGens[subIndex] << " from cocktail group " << mIndex;
+        mResultQueue[subIndex].pop(subGenIndex[pos]);
+      }
+    }
   }
-  LOG(info) << "Importing particles for task " << genIndex;
-
-  // at this moment the mIndex-th generator is ready to be used
+  // Clear particles and event header
   mParticles.clear();
-  mParticles = gens[genIndex]->getParticles();
-
-  // fetch the event Header information from the underlying generator
   mMCEventHeader.clearInfo();
-  gens[genIndex]->updateHeader(&mMCEventHeader);
-
-  mInputTaskQueue.push(genIndex);
-  mTasksStarted++;
+  if (mCocktailMode) {
+    // in cocktail mode we need to merge the particles from the different generators
+    for (auto subIndex : subGenIndex) {
+      LOG(info) << "Importing particles for task " << subIndex;
+      auto subParticles = gens[subIndex]->getParticles();
+      mParticles.insert(mParticles.end(), subParticles.begin(), subParticles.end());
+      // fetch the event Header information from the underlying generator
+      gens[subIndex]->updateHeader(&mMCEventHeader);
+      mInputTaskQueue.push(subIndex);
+      mTasksStarted++;
+    }
+  } else {
+    LOG(info) << "Importing particles for task " << genIndex;
+    // at this moment the mIndex-th generator is ready to be used
+    mParticles = gens[genIndex]->getParticles();
+    // fetch the event Header information from the underlying generator
+    gens[genIndex]->updateHeader(&mMCEventHeader);
+    mInputTaskQueue.push(genIndex);
+    mTasksStarted++;
+  }
 
   mseqCounter++;
   mEventCounter++;
   if (mEventCounter == mNEvents) {
     LOG(info) << "HybridGen: Stopping TBB task pool";
     mStopFlag = true;
   }
+
   return true;
 }
 
@@ -465,6 +495,63 @@ Bool_t GeneratorHybrid::parseJSON(const std::string& path)
     }
   }
 
+  if (doc.HasMember("cocktail")) {
+    mCocktailMode = true;
+    const auto& groups = doc["cocktail"];
+    for (const auto& group : groups.GetArray()) {
+      if (!group.HasMember("group")) {
+        LOG(fatal) << "Group not provided for cocktail";
+        return false;
+      }
+      const auto& sets = group["group"];
+      mGroups.push_back({});
+      for (const auto& subset : sets.GetArray()) {
+        mGroups.back().push_back(subset.GetInt());
+      }
+      if (group.HasMember("fractions")) {
+        const auto& subfrac = group["fractions"];
+        mCocktailFractions.push_back({});
+        for (const auto& frac : subfrac.GetArray()) {
+          mCocktailFractions.back().push_back(frac.GetInt());
+        }
+        if (mGroups.back().size() != mCocktailFractions.back().size()) {
+          LOG(fatal) << "Number of fractions does not match the number of generators in the cocktail";
+          return false;
+        }
+      } else {
+        LOG(warn) << "Fractions not provided for cocktail group " << mGroups.size() - 1 << " setting all fractions to 1";
+        mCocktailFractions.push_back({});
+        for (int i = 0; i < mGroups.back().size(); i++) {
+          mCocktailFractions.back().push_back(1);
+        }
+      }
+    }
+    // Check if mGroups items contain the same number twice and ensure all items are unique across groups
+    std::unordered_set<int> allItems;
+    for (const auto& group : mGroups) {
+      std::unordered_set<int> uniqueItems;
+      for (const auto& item : group) {
+        if (!uniqueItems.insert(item).second) {
+          LOG(fatal) << "Duplicate generator index " << item << " found in a cocktail group";
+          return false;
+        }
+        if (!allItems.insert(item).second) {
+          LOG(fatal) << "Duplicate generator index " << item << " found in the cocktails";
+          return false;
+        }
+      }
+    }
+    // Check if allItems is missing any generator index and create single gen groups for them
+    // Fractions for these generators are the back ones in mFractions
+    for (int i = 0; i < mInputGens.size(); i++) {
+      if (allItems.find(i) == allItems.end()) {
+        LOG(info) << "Generator index " << i << " is missing in the cocktails";
+        LOG(info) << "Setting up a group for generator " << mInputGens[i];
+        mGroups.push_back({i});
+      }
+    }
+  }
+
   // Get fractions and put them in mFractions
   if (doc.HasMember("fractions")) {
     const auto& fractions = doc["fractions"];

run/SimExamples/Hybrid/README.md

@@ -2,7 +2,7 @@
 \page refrunSimExamplesHybrid Example Hybrid
 /doxy -->
 
-The usage of the Hybrid generator with the o2-sim is presented in this short manual.
+The usage of the Hybrid generator with o2-sim is presented in this short manual.
 All the other generators are implemented as sub-generators and they can be called thanks to a
 JSON file, fed to o2-sim via the GeneratorHybrid.configFile parameter. The O2sim package needs to be loaded in order to use this example.
 
@@ -13,6 +13,4 @@ available generators in O2. The JSON template can be generated using the ${O2DPG
 
 - **runo2sim.sh** → allows to use the hybrid generator example
 - **hybridconfig.json** → example JSON file for the hybrid generator configuration
-- **example.optns** → options file to be used in EPOS4 implemented as subgenerator in this example (the .optns must be available in the current working directory)
-- **evtpool.root** → cached events to be used with the extkinO2 generator
-- **epos4.hepmc** → EPOS4 events stored as hepmc file
+- **example.optns** → options file to be used in EPOS4 implemented as subgenerator in this example (the .optns must be available in the current working directory)

run/SimExamples/Hybrid/hybridconfig.json

@@ -53,7 +53,8 @@
       "name": "external",
       "config": {
         "fileName": "${O2DPG_ROOT}/MC/config/PWGDQ/external/generator/GeneratorParamPromptJpsiToElectronEvtGen_pp13TeV.C",
-        "funcName": "GeneratorParamPromptJpsiToElectronEvtGen_pp13TeV()"
+        "funcName": "GeneratorParamPromptJpsiToElectronEvtGen_pp13TeV()",
+        "iniFile": ""
       }
     },
     {

run/SimExamples/Hybrid_cocktail/README.md

@@ -0,0 +1,13 @@
+<!-- doxy
+\page refrunSimExamplesHybrid Example Hybrid_cocktail
+/doxy -->
+
+The usage of the Hybrid generator using cocktails is presented in this example.
+The syntax of the JSON file shows how the cocktails can be grouped. Subfractions are still not implemented, but the JSON syntax to be used is already in place. For now each generator will be pulled exactly once in order for each cocktail event generation.
+The basic Hybrid Generator mechanisms are shown in the Hybrid and Hybrid_parallel example folders.
+The cocktail configuration can not be setup with the template script for now.
+
+# Files description
+
+- **runo2sim.sh** &rarr; allows to use the hybrid generator example
+- **hybridcocktail.json** &rarr; example JSON file for the hybrid generator configuration using cocktails

run/SimExamples/Hybrid_cocktail/hybridcocktail.json

@@ -0,0 +1,69 @@
+{
+  "generators": [
+    {
+      "name": "pythia8",
+      "config": {
+        "config": "$O2_ROOT/share/Generators/egconfig/pythia8_inel.cfg",
+        "hooksFileName": "",
+        "hooksFuncName": "",
+        "includePartonEvent": false,
+        "particleFilter": "",
+        "verbose": 0
+      }
+    },
+    {
+      "name": "external",
+      "config": {
+        "fileName": "${O2DPG_ROOT}/MC/config/PWGDQ/external/generator/GeneratorParamPromptJpsiToElectronEvtGen_pp13TeV.C",
+        "funcName": "GeneratorParamPromptJpsiToElectronEvtGen_pp13TeV()",
+        "iniFile": ""
+      }
+    },
+    {
+      "name": "pythia8pp"
+    },
+    {
+      "name": "extkinO2",
+      "config": {
+        "skipNonTrackable": true,
+        "continueMode": false,
+        "roundRobin": false,
+        "randomize": false,
+        "rngseed": 0,
+        "randomphi": false,
+        "fileName": "${PWD}/evtpool.root"
+      }
+    },
+    {
+      "name": "pythia8hf",
+      "config": ""
+    }
+  ],
+  "cocktail": [
+    {
+      "group": [
+        0,
+        1
+      ],
+      "fractions": [
+        1,
+        1
+      ]
+    },
+    {
+      "group": [
+        2,
+        3
+      ],
+      "fractions": [
+        1,
+        1
+      ]
+    }
+  ],
+  "fractions": [
+    1,
+    1,
+    1
+  ]
+}

run/SimExamples/Hybrid_cocktail/runo2sim.sh

@@ -0,0 +1,68 @@
+#!/usr/bin/env bash
+#
+# Hybrid generator simulation example using cocktails:
+# the simulation is configured using a JSON file (hybridcocktail.json in this folder)
+set -x
+if [ ! "${O2DPG_ROOT}" ]; then
+    echo "This needs O2DPG loaded; alienv enter ..."
+    exit 1
+fi
+
+[ ! "${O2_ROOT}" ] && echo "Error: This needs O2 loaded" && exit 2
+
+NEV=-1
+more=""
+JOBS=2
+
+usage()
+{
+    cat <<EOF
+Usage: $0 [OPTIONS]
+
+Options:
+
+  -m,--more    CONFIG      More configurations ($more)
+  -n,--nevents EVENTS      Number of events ($NEV)
+  -j,--jobs    JOBS        Number of jobs ($JOBS)
+  -h,--help                Print these instructions
+  --                       Rest of command line sent to o2-sim
+
+COMMAND must be quoted if it contains spaces or other special
+characters
+
+Below follows the help output of o2-sim
+
+EOF
+}
+
+if [ "$#" -lt 2 ]; then
+    echo "Running with default values"
+fi
+
+while test $# -gt 0 ; do
+    case $1 in
+        -m|--more)    more="$2" ; shift ;;
+        -n|--nevents) NEV=$2 ; shift ;;
+        -j|--jobs)    JOBS=$2 ; shift ;;
+        -h|--help) usage; o2-sim --help full ; exit 0 ;;
+        --)           shift ; break ;;
+        *) echo "Unknown option '$1', did you forget '--'?" >/dev/stderr
+           exit 3
+           ;;
+    esac
+    shift
+done
+
+# Set number of events in optns file
+if [ ! $NEV -eq -1 ]; then
+    echo "Setting number of events to $NEV"
+else
+    echo "Number of events not set, defaulting to 10..."
+    NEV=10
+fi
+
+# Generation of event pool with pythia8 (10000 events) in a evtpool.root file
+${O2DPG_ROOT}/MC/run/examples/event_pool.sh --make
+
+# Starting simulation with Hybrid generator
+${O2_ROOT}/bin/o2-sim --noGeant -j $JOBS --field ccdb --vertexMode kCCDB --run 300000 --configKeyValues "MFTBase.buildAlignment=true;GeneratorHybrid.configFile=$PWD/hybridcocktail.json;GeneratorHybrid.randomize=false;${more}" -g hybrid -o genevents --timestamp 1546300800000 --seed 836302859 -n $NEV