Skip to content

Commit

Permalink
New event generator/wrapper facilitating event-pools (#13766)
Browse files Browse the repository at this point in the history
A new event generator for event-pools.
This is a convenience wrapper on top of GeneratorO2Kine
which handles event pools better. It can talk to a pool as
a collection, instead of only to specific files.

In particular it offers functionality to
- self-pick a file from a pool
- discover available files in a pool be it on AliEn
  or local
- makes it easier to generate generic JSON configs, where users
  don't need to provide a full file path to use (which would
  be impractical for productions)
  • Loading branch information
sawenzel authored Dec 5, 2024
1 parent 5f1e50c commit 808730c
Show file tree
Hide file tree
Showing 9 changed files with 467 additions and 3 deletions.
50 changes: 49 additions & 1 deletion Generators/include/Generators/GeneratorFromFile.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@
#include "FairGenerator.h"
#include "Generators/Generator.h"
#include "Generators/GeneratorFromO2KineParam.h"
#include "SimulationDataFormat/MCEventHeader.h"
#include <TRandom3.h>
#include <TGrid.h>
#include <random>

class TBranch;
class TFile;
class TParticle;
class TGrid;

namespace o2
{
Expand Down Expand Up @@ -109,6 +111,52 @@ class GeneratorFromO2Kine : public o2::eventgen::Generator
ClassDefOverride(GeneratorFromO2Kine, 2);
};

/// Special generator for event pools.
/// What do we like to have:
/// - ability to give a file which contains the list of files to read
/// - ability to give directly a file to read the event from
/// - ability to give a pool path and to find the top N list of files closest to myself
/// - ability to select itself one file from the pool
class GeneratorFromEventPool : public o2::eventgen::Generator
{
public:
constexpr static std::string_view eventpool_filename = "evtpool.root";
constexpr static std::string_view alien_protocol_prefix = "alien://";

GeneratorFromEventPool() = default; // mainly for ROOT IO
GeneratorFromEventPool(EventPoolGenConfig const& pars);

bool Init() override;

// the o2 Generator interface methods
bool generateEvent() override
{ /* trivial - actual work in importParticles */
return mO2KineGenerator->generateEvent();
}
bool importParticles() override
{
auto import_good = mO2KineGenerator->importParticles();
// transfer the particles (could be avoided)
mParticles = mO2KineGenerator->getParticles();
return import_good;
}

// determine the collection of available files
std::vector<std::string> setupFileUniverse(std::string const& path) const;

std::vector<std::string> const& getFileUniverse() const { return mPoolFilesAvailable; }

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
// random number generator to determine a concrete file name
std::mt19937 mRandomEngine; //!

ClassDefOverride(GeneratorFromEventPool, 1);
};

} // end namespace eventgen
} // end namespace o2

Expand Down
16 changes: 16 additions & 0 deletions Generators/include/Generators/GeneratorFromO2KineParam.h
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,22 @@ struct O2KineGenConfig {
std::string fileName = ""; // filename to read from - takes precedence over SimConfig if given
};

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
unsigned int rngseed = 0; // randomizer seed, 0 for random value
bool randomphi = false; // randomize phi angle; rotates tracks in events by some phi-angle
};

// construct a configurable param singleton out of the
struct GeneratorEventPoolParam : public o2::conf::ConfigurableParamPromoter<GeneratorEventPoolParam, EventPoolGenConfig> {
O2ParamDef(GeneratorEventPoolParam, "GeneratorEventPool");
};

} // end namespace eventgen
} // end namespace o2

Expand Down
3 changes: 2 additions & 1 deletion Generators/include/Generators/GeneratorHybrid.h
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ class GeneratorHybrid : public Generator
private:
o2::eventgen::Generator* currentgen = nullptr;
std::vector<std::unique_ptr<o2::eventgen::Generator>> gens;
const std::vector<std::string> generatorNames = {"extkinO2", "boxgen", "external", "hepmc", "pythia8", "pythia8pp", "pythia8hi", "pythia8hf", "pythia8powheg"};
const std::vector<std::string> generatorNames = {"extkinO2", "evtpool", "boxgen", "external", "hepmc", "pythia8", "pythia8pp", "pythia8hi", "pythia8hf", "pythia8powheg"};
std::vector<std::string> mInputGens;
std::vector<std::string> mGens;
std::vector<std::string> mConfigs;
Expand All @@ -73,6 +73,7 @@ class GeneratorHybrid : public Generator
std::vector<std::unique_ptr<o2::eventgen::BoxGenConfig>> mBoxGenConfigs;
std::vector<std::unique_ptr<o2::eventgen::Pythia8GenConfig>> mPythia8GenConfigs;
std::vector<std::unique_ptr<o2::eventgen::O2KineGenConfig>> mO2KineGenConfigs;
std::vector<o2::eventgen::EventPoolGenConfig> mEventPoolConfigs;
std::vector<std::unique_ptr<o2::eventgen::ExternalGenConfig>> mExternalGenConfigs;
std::vector<std::unique_ptr<o2::eventgen::FileOrCmdGenConfig>> mFileOrCmdGenConfigs;
std::vector<std::unique_ptr<o2::eventgen::HepMCGenConfig>> mHepMCGenConfigs;
Expand Down
7 changes: 7 additions & 0 deletions Generators/src/GeneratorFactory.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,13 @@ void GeneratorFactory::setPrimaryGenerator(o2::conf::SimConfig const& conf, Fair
}
}
LOG(info) << "using external O2 kinematics";
} else if (genconfig.compare("evtpool") == 0) {
// case of an "event-pool" which is a specialization of extkinO2
// with some additional logic in file management and less configurability
// and not features such as "continue transport"
auto extGen = new o2::eventgen::GeneratorFromEventPool(o2::eventgen::GeneratorEventPoolParam::Instance().detach());
primGen->AddGenerator(extGen);
LOG(info) << "using the eventpool generator";
} else if (genconfig.compare("tparticle") == 0) {
// External ROOT file(s) with tree of TParticle in clones array,
// or external program generating such a file
Expand Down
232 changes: 231 additions & 1 deletion Generators/src/GeneratorFromFile.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
#include <TParticle.h>
#include <TTree.h>
#include <sstream>
#include <filesystem>
#include <TGrid.h>

namespace o2
{
Expand Down Expand Up @@ -249,6 +251,7 @@ bool GeneratorFromO2Kine::importParticles()
// Randomize the order of events in the input file
if (mRandomize) {
mEventCounter = gRandom->Integer(mEventsAvailable);
LOG(info) << "GeneratorFromO2Kine - Picking event " << mEventCounter;
}

double dPhi = 0.;
Expand Down Expand Up @@ -352,8 +355,235 @@ void GeneratorFromO2Kine::updateHeader(o2::dataformats::MCEventHeader* eventHead
eventHeader->putInfo<int>("forwarding-generator_inputEventNumber", mEventCounter - 1);
}

namespace
{
// some helper to execute a command and capture it's output in a vector
std::vector<std::string> executeCommand(const std::string& command)
{
std::vector<std::string> result;
std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(command.c_str(), "r"), pclose);
if (!pipe) {
throw std::runtime_error("Failed to open pipe");
}

char buffer[1024];
while (fgets(buffer, sizeof(buffer), pipe.get()) != nullptr) {
std::string line(buffer);
// Remove trailing newline character, if any
if (!line.empty() && line.back() == '\n') {
line.pop_back();
}
result.push_back(line);
}
return result;
}
} // namespace

GeneratorFromEventPool::GeneratorFromEventPool(EventPoolGenConfig const& pars) : mConfig{pars}
{
}

bool GeneratorFromEventPool::Init()
{
// initialize the event pool
if (mConfig.rngseed > 0) {
mRandomEngine.seed(mConfig.rngseed);
} else {
std::random_device rd;
mRandomEngine.seed(rd());
}
mPoolFilesAvailable = setupFileUniverse(mConfig.eventPoolPath);

if (mPoolFilesAvailable.size() == 0) {
LOG(error) << "No file found that can be used with EventPool generator";
return false;
}

// now choose the actual file
std::uniform_int_distribution<int> distribution(0, mPoolFilesAvailable.size());
mFileChosen = mPoolFilesAvailable[distribution(mRandomEngine)];
LOG(info) << "EventPool is using file " << mFileChosen;

// we bring up the internal mO2KineGenerator
auto kine_config = O2KineGenConfig{
.skipNonTrackable = mConfig.skipNonTrackable,
.continueMode = false,
.roundRobin = false,
.randomize = mConfig.randomize,
.rngseed = mConfig.rngseed,
.randomphi = mConfig.randomphi,
.fileName = mFileChosen};
mO2KineGenerator.reset(new GeneratorFromO2Kine(kine_config));
return mO2KineGenerator->Init();
}

namespace
{
namespace fs = std::filesystem;
// checks a single file name
bool checkFileName(std::string const& pathStr)
{
// LOG(info) << "Checking filename " << pathStr;
try {
// Remove optional protocol prefix "alien://"
const std::string protocol = "alien://";
std::string finalPathStr(pathStr);
if (pathStr.starts_with(protocol)) {
finalPathStr = pathStr.substr(protocol.size());
}
fs::path path(finalPathStr);

// Check if the filename is "eventpool.root"
return path.filename() == GeneratorFromEventPool::eventpool_filename;
} catch (const fs::filesystem_error& e) {
// Invalid path syntax will throw an exception
std::cerr << "Filesystem error: " << e.what() << '\n';
return false;
} catch (...) {
// Catch-all for other potential exceptions
std::cerr << "An unknown error occurred while checking the path.\n";
return false;
}
}

// checks a whole universe of file names
bool checkFileUniverse(std::vector<std::string> const& universe)
{
if (universe.size() == 0) {
return false;
}
for (auto& fn : universe) {
if (!checkFileName(fn)) {
return false;
}
}
// TODO: also check for a common path structure with maximally 00X as only difference

return true;
}

std::vector<std::string> readLines(const std::string& filePath)
{
std::vector<std::string> lines;

// Check if the file is a valid text file
fs::path path(filePath);

// Open the file
std::ifstream file(filePath);
if (!file.is_open()) {
throw std::ios_base::failure("Failed to open the file.");
}

// Read up to n lines
std::string line;
while (std::getline(file, line)) {
lines.push_back(line);
}
return lines;
}

// Function to find all files named eventpool_filename under a given path
std::vector<std::string> getLocalFileList(const fs::path& rootPath)
{
std::vector<std::string> result;

// Ensure the root path exists and is a directory
if (!fs::exists(rootPath) || !fs::is_directory(rootPath)) {
throw std::invalid_argument("The provided path is not a valid directory.");
}

// Iterate over the directory and subdirectories
for (const auto& entry : fs::recursive_directory_iterator(rootPath)) {
if (entry.is_regular_file() && entry.path().filename() == GeneratorFromEventPool::eventpool_filename) {
result.push_back(entry.path().string());
}
}
return result;
}

} // end anonymous namespace

/// A function determining the universe of event pool files, as determined by the path string
/// returns empty vector if it fails
std::vector<std::string> GeneratorFromEventPool::setupFileUniverse(std::string const& path) const
{
// the path could refer to a local or alien filesystem; find out first
bool onAliEn = strncmp(path.c_str(), std::string(alien_protocol_prefix).c_str(), alien_protocol_prefix.size()) == 0;
std::vector<std::string> result;

if (onAliEn) {
// AliEn case
// we support: (a) an actual evtgen file and (b) a path containing multiple eventfiles

auto alienStatTypeCommand = std::string("alien.py stat ") + mConfig.eventPoolPath + std::string(" 2>/dev/null | grep Type ");
auto typeString = executeCommand(alienStatTypeCommand);
if (typeString.size() == 0) {
return result;
} else if (typeString.size() == 1 && typeString.front() == std::string("Type: f")) {
// this is a file ... simply use it
result.push_back(mConfig.eventPoolPath);
return result;
} else if (typeString.size() == 1 && typeString.front() == std::string("Type: d")) {
// this is a directory
// construct command to find actual event files
std::string alienSearchCommand = std::string("alien.py find ") +
mConfig.eventPoolPath + "/ " + std::string(eventpool_filename);

auto universe_vector = executeCommand(alienSearchCommand);
// check vector
if (!checkFileUniverse(universe_vector)) {
return result;
}
for (auto& f : universe_vector) {
f = std::string(alien_protocol_prefix) + f;
}

return universe_vector;
} else {
LOG(error) << "Unsupported file type";
return result;
}
} else {
// local file case
// check if the path is a regular file
auto is_actual_file = std::filesystem::is_regular_file(path);
if (is_actual_file) {
// The files must match a criteria of being canonical paths ending with eventpool_Kine.root
if (checkFileName(path)) {
TFile rootfile(path.c_str(), "OPEN");
if (!rootfile.IsZombie()) {
result.push_back(path);
return result;
}
} else {
// otherwise assume it is a text file containing a list of files themselves
auto files = readLines(path);
if (checkFileUniverse(files)) {
result = files;
return result;
}
}
} else {
// check if the path is just a path
// In this case we need to search something and check
auto is_dir = std::filesystem::is_directory(path);
if (!is_dir) {
return result;
}
auto files = getLocalFileList(path);
if (checkFileUniverse(files)) {
result = files;
return result;
}
}
}
return result;
}

} // namespace eventgen
} // end namespace o2

ClassImp(o2::eventgen::GeneratorFromEventPool);
ClassImp(o2::eventgen::GeneratorFromFile);
ClassImp(o2::eventgen::GeneratorFromO2Kine);
ClassImp(o2::eventgen::GeneratorFromO2Kine);
1 change: 1 addition & 0 deletions Generators/src/GeneratorFromO2KineParam.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@

#include "Generators/GeneratorFromO2KineParam.h"
O2ParamImpl(o2::eventgen::GeneratorFromO2KineParam);
O2ParamImpl(o2::eventgen::GeneratorEventPoolParam);
Loading

0 comments on commit 808730c

Please sign in to comment.