Implementation of cocktail generation

Works for sequence generation

* Updated Hybrid example

Generators/include/Generators/GeneratorHybrid.h

@@ -98,6 +98,10 @@ class GeneratorHybrid : public Generator
   int mEventCounter = 0;
   int mTasksStarted = 0;
 
+  // Cocktail mode
+  bool mCocktailMode = false;
+  std::vector<std::vector<int>> mGroups;
+
   // 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,46 @@ 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());
+      }
+    }
+    // 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. 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,61 @@
+{
+  "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
+      ]
+    },
+    {
+      "group": [
+        2,
+        3
+      ]
+    }
+  ],
+  "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