1212#include " Generators/GeneratorHybrid.h"
1313#include < fairlogger/Logger.h>
1414#include < algorithm>
15-
1615#include < tbb/concurrent_queue.h>
1716#include < tbb/task_arena.h>
1817#include < tbb/parallel_for.h>
@@ -42,9 +41,16 @@ GeneratorHybrid::GeneratorHybrid(const std::string& inputgens)
4241 }
4342 int index = 0 ;
4443 if (!(mRandomize || mGenerationMode == GenMode::kParallel )) {
45- if (mFractions .size () != mInputGens .size ()) {
46- LOG (fatal) << " Number of fractions does not match the number of generators" ;
47- return ;
44+ if (mCocktailMode ) {
45+ if (mGroups .size () != mFractions .size ()) {
46+ LOG (fatal) << " Number of groups does not match the number of fractions" ;
47+ return ;
48+ }
49+ } else {
50+ if (mFractions .size () != mInputGens .size ()) {
51+ LOG (fatal) << " Number of fractions does not match the number of generators" ;
52+ return ;
53+ }
4854 }
4955 // Check if all elements of mFractions are 0
5056 if (std::all_of (mFractions .begin (), mFractions .end (), [](int i) { return i == 0 ; })) {
@@ -303,7 +309,7 @@ bool GeneratorHybrid::generateEvent()
303309 }
304310 }
305311 } else {
306- mIndex = gRandom ->Integer (mGens .size ());
312+ mIndex = gRandom ->Integer (mFractions .size ());
307313 }
308314 } else {
309315 while (mFractions [mCurrentFraction ] == 0 || mseqCounter == mFractions [mCurrentFraction ]) {
@@ -322,32 +328,55 @@ bool GeneratorHybrid::generateEvent()
322328bool GeneratorHybrid::importParticles ()
323329{
324330 int genIndex = -1 ;
331+ std::vector<int > subGenIndex = {};
325332 if (mIndex == -1 ) {
326333 // this means parallel mode ---> we have a common queue
327334 mResultQueue [0 ].pop (genIndex);
328335 } else {
329336 // need to pop from a particular queue
330- mResultQueue [mIndex ].pop (genIndex);
337+ if (!mCocktailMode ) {
338+ mResultQueue [mIndex ].pop (genIndex);
339+ } else {
340+ // in cocktail mode we need to pop from the group queue
341+ subGenIndex.resize (mGroups [mIndex ].size ());
342+ for (size_t pos = 0 ; pos < mGroups [mIndex ].size (); ++pos) {
343+ int subIndex = mGroups [mIndex ][pos];
344+ LOG (info) << " Getting generator " << mGens [subIndex] << " from cocktail group " << mIndex ;
345+ mResultQueue [subIndex].pop (subGenIndex[pos]);
346+ }
347+ }
331348 }
332- LOG (info) << " Importing particles for task " << genIndex;
333-
334- // at this moment the mIndex-th generator is ready to be used
349+ // Clear particles and event header
335350 mParticles .clear ();
336- mParticles = gens[genIndex]->getParticles ();
337-
338- // fetch the event Header information from the underlying generator
339351 mMCEventHeader .clearInfo ();
340- gens[genIndex]->updateHeader (&mMCEventHeader );
341-
342- mInputTaskQueue .push (genIndex);
343- mTasksStarted ++;
352+ if (mCocktailMode ) {
353+ // in cocktail mode we need to merge the particles from the different generators
354+ for (auto subIndex : subGenIndex) {
355+ LOG (info) << " Importing particles for task " << subIndex;
356+ auto subParticles = gens[subIndex]->getParticles ();
357+ mParticles .insert (mParticles .end (), subParticles.begin (), subParticles.end ());
358+ // fetch the event Header information from the underlying generator
359+ gens[subIndex]->updateHeader (&mMCEventHeader );
360+ mInputTaskQueue .push (subIndex);
361+ mTasksStarted ++;
362+ }
363+ } else {
364+ LOG (info) << " Importing particles for task " << genIndex;
365+ // at this moment the mIndex-th generator is ready to be used
366+ mParticles = gens[genIndex]->getParticles ();
367+ // fetch the event Header information from the underlying generator
368+ gens[genIndex]->updateHeader (&mMCEventHeader );
369+ mInputTaskQueue .push (genIndex);
370+ mTasksStarted ++;
371+ }
344372
345373 mseqCounter++;
346374 mEventCounter ++;
347375 if (mEventCounter == mNEvents ) {
348376 LOG (info) << " HybridGen: Stopping TBB task pool" ;
349377 mStopFlag = true ;
350378 }
379+
351380 return true ;
352381}
353382
@@ -371,6 +400,59 @@ std::string GeneratorHybrid::jsonValueToString(const T& value)
371400 return buffer.GetString ();
372401}
373402
403+ Bool_t GeneratorHybrid::confSetter (const auto & gen)
404+ {
405+ std::string name = gen[" name" ].GetString ();
406+ mInputGens .push_back (name);
407+ if (gen.HasMember (" config" )) {
408+ if (name == " boxgen" ) {
409+ const auto & boxconf = gen[" config" ];
410+ auto boxConfig = TBufferJSON::FromJSON<o2::eventgen::BoxGenConfig>(jsonValueToString (boxconf).c_str ());
411+ mBoxGenConfigs .push_back (std::move (boxConfig));
412+ mConfigs .push_back (" boxgen_" + std::to_string (mBoxGenConfigs .size () - 1 ));
413+ } else if (name == " pythia8" ) {
414+ const auto & pythia8conf = gen[" config" ];
415+ auto pythia8Config = TBufferJSON::FromJSON<o2::eventgen::Pythia8GenConfig>(jsonValueToString (pythia8conf).c_str ());
416+ mPythia8GenConfigs .push_back (std::move (pythia8Config));
417+ mConfigs .push_back (" pythia8_" + std::to_string (mPythia8GenConfigs .size () - 1 ));
418+ } else if (name == " extkinO2" ) {
419+ const auto & o2kineconf = gen[" config" ];
420+ auto o2kineConfig = TBufferJSON::FromJSON<o2::eventgen::O2KineGenConfig>(jsonValueToString (o2kineconf).c_str ());
421+ mO2KineGenConfigs .push_back (std::move (o2kineConfig));
422+ mConfigs .push_back (" extkinO2_" + std::to_string (mO2KineGenConfigs .size () - 1 ));
423+ } else if (name == " evtpool" ) {
424+ const auto & o2kineconf = gen[" config" ];
425+ auto poolConfig = TBufferJSON::FromJSON<o2::eventgen::EventPoolGenConfig>(jsonValueToString (o2kineconf).c_str ());
426+ mEventPoolConfigs .push_back (*poolConfig);
427+ mConfigs .push_back (" evtpool_" + std::to_string (mEventPoolConfigs .size () - 1 ));
428+ } else if (name == " external" ) {
429+ const auto & extconf = gen[" config" ];
430+ auto extConfig = TBufferJSON::FromJSON<o2::eventgen::ExternalGenConfig>(jsonValueToString (extconf).c_str ());
431+ mExternalGenConfigs .push_back (std::move (extConfig));
432+ mConfigs .push_back (" external_" + std::to_string (mExternalGenConfigs .size () - 1 ));
433+ } else if (name == " hepmc" ) {
434+ const auto & genconf = gen[" config" ];
435+ const auto & cmdconf = genconf[" configcmd" ];
436+ const auto & hepmcconf = genconf[" confighepmc" ];
437+ auto cmdConfig = TBufferJSON::FromJSON<o2::eventgen::FileOrCmdGenConfig>(jsonValueToString (cmdconf).c_str ());
438+ auto hepmcConfig = TBufferJSON::FromJSON<o2::eventgen::HepMCGenConfig>(jsonValueToString (hepmcconf).c_str ());
439+ mFileOrCmdGenConfigs .push_back (std::move (cmdConfig));
440+ mHepMCGenConfigs .push_back (std::move (hepmcConfig));
441+ mConfigs .push_back (" hepmc_" + std::to_string (mFileOrCmdGenConfigs .size () - 1 ));
442+ } else {
443+ mConfigs .push_back (" " );
444+ }
445+ } else {
446+ if (name == " boxgen" || name == " pythia8" || name == " extkinO2" || name == " external" || name == " hepmc" ) {
447+ LOG (fatal) << " No configuration provided for generator " << name;
448+ return false ;
449+ } else {
450+ mConfigs .push_back (" " );
451+ }
452+ }
453+ return true ;
454+ }
455+
374456Bool_t GeneratorHybrid::parseJSON (const std::string& path)
375457{
376458 // Parse JSON file to build map
@@ -407,60 +489,26 @@ Bool_t GeneratorHybrid::parseJSON(const std::string& path)
407489 if (doc.HasMember (" generators" )) {
408490 const auto & gens = doc[" generators" ];
409491 for (const auto & gen : gens.GetArray ()) {
410- // push in mInputGens the "name" of the generator
411- std::string name = gen[" name" ].GetString ();
412- mInputGens .push_back (name);
413- if (gen.HasMember (" config" )) {
414- if (name == " boxgen" ) {
415- const auto & boxconf = gen[" config" ];
416- auto boxConfig = TBufferJSON::FromJSON<o2::eventgen::BoxGenConfig>(jsonValueToString (boxconf).c_str ());
417- mBoxGenConfigs .push_back (std::move (boxConfig));
418- mConfigs .push_back (" boxgen_" + std::to_string (mBoxGenConfigs .size () - 1 ));
419- continue ;
420- } else if (name == " pythia8" ) {
421- const auto & pythia8conf = gen[" config" ];
422- auto pythia8Config = TBufferJSON::FromJSON<o2::eventgen::Pythia8GenConfig>(jsonValueToString (pythia8conf).c_str ());
423- mPythia8GenConfigs .push_back (std::move (pythia8Config));
424- mConfigs .push_back (" pythia8_" + std::to_string (mPythia8GenConfigs .size () - 1 ));
425- continue ;
426- } else if (name == " extkinO2" ) {
427- const auto & o2kineconf = gen[" config" ];
428- auto o2kineConfig = TBufferJSON::FromJSON<o2::eventgen::O2KineGenConfig>(jsonValueToString (o2kineconf).c_str ());
429- mO2KineGenConfigs .push_back (std::move (o2kineConfig));
430- mConfigs .push_back (" extkinO2_" + std::to_string (mO2KineGenConfigs .size () - 1 ));
431- continue ;
432- } else if (name == " evtpool" ) {
433- const auto & o2kineconf = gen[" config" ];
434- auto poolConfig = TBufferJSON::FromJSON<o2::eventgen::EventPoolGenConfig>(jsonValueToString (o2kineconf).c_str ());
435- mEventPoolConfigs .push_back (*poolConfig);
436- mConfigs .push_back (" evtpool_" + std::to_string (mEventPoolConfigs .size () - 1 ));
437- continue ;
438- } else if (name == " external" ) {
439- const auto & extconf = gen[" config" ];
440- auto extConfig = TBufferJSON::FromJSON<o2::eventgen::ExternalGenConfig>(jsonValueToString (extconf).c_str ());
441- mExternalGenConfigs .push_back (std::move (extConfig));
442- mConfigs .push_back (" external_" + std::to_string (mExternalGenConfigs .size () - 1 ));
443- continue ;
444- } else if (name == " hepmc" ) {
445- const auto & genconf = gen[" config" ];
446- const auto & cmdconf = genconf[" configcmd" ];
447- const auto & hepmcconf = genconf[" confighepmc" ];
448- auto cmdConfig = TBufferJSON::FromJSON<o2::eventgen::FileOrCmdGenConfig>(jsonValueToString (cmdconf).c_str ());
449- auto hepmcConfig = TBufferJSON::FromJSON<o2::eventgen::HepMCGenConfig>(jsonValueToString (hepmcconf).c_str ());
450- mFileOrCmdGenConfigs .push_back (std::move (cmdConfig));
451- mHepMCGenConfigs .push_back (std::move (hepmcConfig));
452- mConfigs .push_back (" hepmc_" + std::to_string (mFileOrCmdGenConfigs .size () - 1 ));
453- continue ;
454- } else {
455- mConfigs .push_back (" " );
492+ mGroups .push_back ({});
493+ // Check if gen is an array (cocktail mode)
494+ if (gen.HasMember (" cocktail" )) {
495+ mCocktailMode = true ;
496+ for (const auto & subgen : gen[" cocktail" ].GetArray ()) {
497+ if (confSetter (subgen)) {
498+ mGroups .back ().push_back (mInputGens .size () - 1 );
499+ } else {
500+ return false ;
501+ }
456502 }
457503 } else {
458- if (name == " boxgen" || name == " pythia8" || name == " extkinO2" || name == " external" || name == " hepmc" ) {
459- LOG (fatal) << " No configuration provided for generator " << name;
504+ if (!confSetter (gen)) {
460505 return false ;
461- } else {
462- mConfigs .push_back (" " );
463506 }
507+ // Groups are created in case cocktail mode is activated, this way
508+ // cocktails can be declared anywhere in the JSON file, without the need
509+ // of grouping single generators. If no cocktail is defined
510+ // groups will be ignored nonetheless.
511+ mGroups .back ().push_back (mInputGens .size () - 1 );
464512 }
465513 }
466514 }
0 commit comments