From f60eb587524f197dfb490ae3fdb79c75fdf0be4b Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Thu, 25 Sep 2025 15:57:50 +0200 Subject: [PATCH 01/10] First implementation of loopers inclusion in base Generator class --- .../SimConfig/include/SimConfig/SimConfig.h | 3 + Common/SimConfig/src/SimConfig.cxx | 2 + Generators/CMakeLists.txt | 13 + Generators/include/Generators/Generator.h | 13 + .../include/Generators/TPCLoopersParam.h | 48 ++ Generators/include/TPCLoopers.h | 127 +++++ .../share/egconfig/ScalerComptonParams.json | 28 ++ .../share/egconfig/ScalerPairParams.json | 34 ++ Generators/share/egconfig/gaussian_params.csv | 4 + Generators/share/egconfig/poisson_params.csv | 3 + Generators/src/Generator.cxx | 442 ++++++++++++------ Generators/src/GeneratorsLinkDef.h | 4 + Generators/src/TPCLoopers.cxx | 417 +++++++++++++++++ Generators/src/TPCLoopersParam.cxx | 15 + 14 files changed, 1007 insertions(+), 146 deletions(-) create mode 100644 Generators/include/Generators/TPCLoopersParam.h create mode 100644 Generators/include/TPCLoopers.h create mode 100644 Generators/share/egconfig/ScalerComptonParams.json create mode 100644 Generators/share/egconfig/ScalerPairParams.json create mode 100644 Generators/share/egconfig/gaussian_params.csv create mode 100644 Generators/share/egconfig/poisson_params.csv create mode 100644 Generators/src/TPCLoopers.cxx create mode 100644 Generators/src/TPCLoopersParam.cxx diff --git a/Common/SimConfig/include/SimConfig/SimConfig.h b/Common/SimConfig/include/SimConfig/SimConfig.h index be88d9fbd8c33..8642a0e5bc225 100644 --- a/Common/SimConfig/include/SimConfig/SimConfig.h +++ b/Common/SimConfig/include/SimConfig/SimConfig.h @@ -52,6 +52,7 @@ struct SimConfigData { std::vector mActiveModules; // list of active modules std::vector mReadoutDetectors; // list of readout detectors std::string mMCEngine; // chosen VMC engine + bool mNoLoopers = false; // Disable automatic TPC loopers std::string mGenerator; // chosen VMC generator std::string mTrigger; // chosen VMC generator trigger unsigned int mNEvents; // number of events to be simulated @@ -138,6 +139,8 @@ class SimConfig // get selected active detectors std::vector const& getActiveModules() const { return mConfigData.mActiveModules; } std::vector const& getReadoutDetectors() const { return mConfigData.mReadoutDetectors; } + // get loopers veto + bool getLoopersVeto() const { return mConfigData.mNoLoopers; } // static helper functions to determine list of active / readout modules // can also be used from outside diff --git a/Common/SimConfig/src/SimConfig.cxx b/Common/SimConfig/src/SimConfig.cxx index 9407a3c556179..e6e0b6adf2093 100644 --- a/Common/SimConfig/src/SimConfig.cxx +++ b/Common/SimConfig/src/SimConfig.cxx @@ -74,6 +74,7 @@ void SimConfig::initOptions(boost::program_options::options_description& options "run", bpo::value()->default_value(-1), "ALICE run number")( "asservice", bpo::value()->default_value(false), "run in service/server mode")( "noGeant", bpo::bool_switch(), "prohibits any Geant transport/physics (by using tight cuts)")( + "noLoopers", bpo::bool_switch(), "disable automatic TPC loopers")( "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()->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\". The format is COLLISIONCONTEXTFILE.root[:SIGNALNAME] where SIGNALNAME is the event part in the context which is relevant."); @@ -290,6 +291,7 @@ bool SimConfig::resetFromParsedMap(boost::program_options::variables_map const& using o2::detectors::DetID; mConfigData.mMCEngine = vm["mcEngine"].as(); mConfigData.mNoGeant = vm["noGeant"].as(); + mConfigData.mNoLoopers = vm["noLoopers"].as(); // Reset modules and detectors as they are anyway re-parsed mConfigData.mReadoutDetectors.clear(); diff --git a/Generators/CMakeLists.txt b/Generators/CMakeLists.txt index 02caa63df0d43..56fe8b8fc2284 100644 --- a/Generators/CMakeLists.txt +++ b/Generators/CMakeLists.txt @@ -41,6 +41,8 @@ o2_add_library(Generators src/GeneratorTParticleParam.cxx src/GeneratorService.cxx src/FlowMapper.cxx + $<$:src/TPCLoopers.cxx> + $<$:src/TPCLoopersParam.cxx> $<$:src/GeneratorPythia8.cxx> $<$:src/DecayerPythia8.cxx> $<$:src/GeneratorPythia8Param.cxx> @@ -53,6 +55,7 @@ o2_add_library(Generators PUBLIC_LINK_LIBRARIES FairRoot::Base O2::SimConfig O2::CommonUtils O2::DetectorsBase O2::ZDCBase O2::SimulationDataFormat ${pythiaTarget} ${hepmcTarget} FairRoot::Gen + $<$:onnxruntime::onnxruntime> TARGETVARNAME targetName) if(pythia_FOUND) @@ -63,6 +66,10 @@ if(HepMC3_FOUND) target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_HEPMC3) endif() +if(onnxruntime_FOUND) + target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_ONNXRUNTIME) +endif() + set(headers include/Generators/Generator.h include/Generators/Trigger.h @@ -88,6 +95,12 @@ set(headers include/Generators/FlowMapper.h ) +if(onnxruntime_FOUND) + list(APPEND headers + include/Generators/TPCLoopers.h + include/Generators/TPCLoopersParam.h) +endif() + if(pythia_FOUND) list(APPEND headers include/Generators/GeneratorPythia8.h diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h index bd35a00793e2d..374d53f324399 100644 --- a/Generators/include/Generators/Generator.h +++ b/Generators/include/Generators/Generator.h @@ -17,6 +17,10 @@ #include "FairGenerator.h" #include "TParticle.h" #include "Generators/Trigger.h" +#ifdef GENERATORS_WITH_ONNXRUNTIME +#include "Generators/TPCLoopers.h" +#include "Generators/TPCLoopersParam.h" +#endif #include #include #include @@ -73,6 +77,7 @@ class Generator : public FairGenerator /** methods to override **/ virtual Bool_t generateEvent() = 0; // generates event (in structure internal to generator) virtual Bool_t importParticles() = 0; // fills the mParticles vector (transfer from generator state) + Bool_t loopers(); // adds loopers to the event in case TPC is used virtual void updateHeader(o2::dataformats::MCEventHeader* eventHeader) {}; Bool_t triggerEvent(); @@ -154,6 +159,8 @@ class Generator : public FairGenerator private: void updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const; + // loopers flag + Bool_t mAddLoopers = kFALSE; // collect an ID and a short description of sub-generator entities std::unordered_map mSubGeneratorsIdToDesc; // the current ID of the sub-generator used in the current event (if applicable) @@ -162,6 +169,12 @@ class Generator : public FairGenerator // global static information about (upper limit of) number of events to be generated static unsigned int gTotalNEvents; +#ifdef GENERATORS_WITH_ONNXRUNTIME + // Loopers generator instance + std::unique_ptr mLoopersGen = nullptr; +#endif + void initLoopersGen(); + ClassDefOverride(Generator, 2); }; /** class Generator **/ diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h new file mode 100644 index 0000000000000..ceeea201538b2 --- /dev/null +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -0,0 +1,48 @@ +// Copyright 2024-2025 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \author M+Giacalone - September 2025 + +#ifndef ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_ +#define ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_ + +#include "CommonUtils/ConfigurableParam.h" +#include "CommonUtils/ConfigurableParamHelper.h" + +namespace o2 +{ +namespace eventgen +{ + +/** + ** a parameter class/struct to keep the settings of + ** the tpc loopers event-generator and + ** allow the user to modify them + **/ +struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper { + std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production + std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering + std::string poisson = "${O2_ROOT}/share/Generators/egconfig/poisson_params.csv"; // file with Poissonian parameters + std::string gauss = "${O2_ROOT}/share/Generators/egconfig/gaussian_params.csv"; // file with Gaussian parameters + std::string scaler_pair = "${O2_ROOT}/share/Generators/egconfig/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production + std::string scaler_compton = "${O2_ROOT}/share/Generators/egconfig/ScalerComptonParams.json"; // file with scaler parameters for Compton scattering + bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume + int nFlatGasLoopers = 500; // number of loopers to be generated per event in case of flat gas + float fraction_pairs = 0.08; // fraction of loopers + std::array multiplier = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling + std::array fixedNLoopers = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty + O2ParamDef(GenTPCLoopersParam, "GenTPCLoopers"); +}; + +} // end namespace eventgen +} // end namespace o2 + +#endif // ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_ diff --git a/Generators/include/TPCLoopers.h b/Generators/include/TPCLoopers.h new file mode 100644 index 0000000000000..70146a82baf60 --- /dev/null +++ b/Generators/include/TPCLoopers.h @@ -0,0 +1,127 @@ +#ifndef ALICEO2_EVENTGEN_TPCLOOPERS_H_ +#define ALICEO2_EVENTGEN_TPCLOOPERS_H_ + +#ifdef GENERATORS_WITH_ONNXRUNTIME +#include +#endif +#include +#include +#include +#include +#include "CCDB/CCDBTimeStampUtils.h" +#include "CCDB/CcdbApi.h" +#include "DetectorsRaw/HBFUtils.h" +#include "TRandom3.h" +#include "TDatabasePDG.h" +#include +#include +#include "SimulationDataFormat/MCGenProperties.h" +#include "TParticle.h" +#include + +#ifdef GENERATORS_WITH_ONNXRUNTIME +// Static Ort::Env instance for multiple onnx model loading +extern Ort::Env global_env; +#endif + +#ifdef GENERATORS_WITH_ONNXRUNTIME +// This class is responsible for loading the scaler parameters from a JSON file +// and applying the inverse transformation to the generated data. +struct Scaler +{ + std::vector normal_min; + std::vector normal_max; + std::vector outlier_center; + std::vector outlier_scale; + + void load(const std::string &filename); + + std::vector inverse_transform(const std::vector &input); + +private: + std::vector jsonArrayToVector(const rapidjson::Value &jsonArray); +}; + +// This class loads the ONNX model and generates samples using it. +class ONNXGenerator +{ +public: + ONNXGenerator(Ort::Env &shared_env, const std::string &model_path); + + std::vector generate_sample(); + +private: + Ort::Env &env; + Ort::Session session; + TRandom3 rand_gen; +}; +#endif // GENERATORS_WITH_ONNXRUNTIME + +namespace o2 +{ +namespace eventgen +{ + +#ifdef GENERATORS_WITH_ONNXRUNTIME +class GenTPCLoopers +{ + public: + GenTPCLoopers(std::string model_pairs = "tpcloopmodel.onnx", std::string model_compton = "tpcloopmodelcompton.onnx", + std::string poisson = "poisson.csv", std::string gauss = "gauss.csv", std::string scaler_pair = "scaler_pair.json", + std::string scaler_compton = "scaler_compton.json"); + + Bool_t generateEvent(); + + Bool_t generateEvent(double &time_limit); + + std::vector importParticles(); + + unsigned int PoissonPairs(); + + unsigned int GaussianElectrons(); + + void SetNLoopers(unsigned int &nsig_pair, unsigned int &nsig_compton); + + void SetMultiplier(std::array &mult); + + void setFlatGas(Bool_t &flat, const Int_t &number = -1); + + void setFractionPairs(float &fractionPairs); + + private: + std::unique_ptr mONNX_pair = nullptr; + std::unique_ptr mONNX_compton = nullptr; + std::unique_ptr mScaler_pair = nullptr; + std::unique_ptr mScaler_compton = nullptr; + double mPoisson[3] = {0.0, 0.0, 0.0}; // Mu, Min and Max of Poissonian + double mGauss[4] = {0.0, 0.0, 0.0, 0.0}; // Mean, Std, Min, Max + std::vector> mGenPairs; + std::vector> mGenElectrons; + unsigned int mNLoopersPairs = -1; + unsigned int mNLoopersCompton = -1; + std::array mMultiplier = {1., 1.}; + bool mPoissonSet = false; + bool mGaussSet = false; + // Random number generator + TRandom3 mRandGen; + // Masses of the electrons and positrons + TDatabasePDG *mPDG = TDatabasePDG::Instance(); + double mMass_e = mPDG->GetParticle(11)->Mass(); + double mMass_p = mPDG->GetParticle(-11)->Mass(); + int mCurrentEvent = 0; // Current event number, used for adaptive loopers + TFile *mContextFile = nullptr; // Input collision context file + o2::steer::DigitizationContext *mCollisionContext = nullptr; // Pointer to the digitization context + std::vector mInteractionTimeRecords; // Interaction time records from collision context + Bool_t mFlatGas = false; // Flag to indicate if flat gas loopers are used + Int_t mFlatGasNumber = -1; // Number of flat gas loopers per event + double mIntTimeRecMean = 1.0; // Average interaction time record used for the reference + double mTimeLimit = 0.0; // Time limit for the current event + double mTimeEnd = 0.0; // Time limit for the last event + float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs +}; +#endif // GENERATORS_WITH_ONNXRUNTIME + +} // namespace eventgen +} // namespace o2 + +#endif // ALICEO2_EVENTGEN_TPCLOOPERS_H_ \ No newline at end of file diff --git a/Generators/share/egconfig/ScalerComptonParams.json b/Generators/share/egconfig/ScalerComptonParams.json new file mode 100644 index 0000000000000..d8e654847f46e --- /dev/null +++ b/Generators/share/egconfig/ScalerComptonParams.json @@ -0,0 +1,28 @@ +{ + "normal": { + "min": [ + -0.0108811147511005, + -0.0098758740350604, + -0.0103233363479375, + -260.0542297363281, + -259.80059814453125 + ], + "max": [ + 0.0108060473576188, + 0.0103057539090514, + 0.0106524610891938, + 260.0343933105469, + 259.62890625 + ] + }, + "outlier": { + "center": [ + -71.39387130737305, + 96791.23828125 + ], + "scale": [ + 265.9389114379883, + 230762.30981445312 + ] + } +} \ No newline at end of file diff --git a/Generators/share/egconfig/ScalerPairParams.json b/Generators/share/egconfig/ScalerPairParams.json new file mode 100644 index 0000000000000..61434bfa2462e --- /dev/null +++ b/Generators/share/egconfig/ScalerPairParams.json @@ -0,0 +1,34 @@ +{ + "normal": { + "min": [ + -0.0073022879660129, + -0.0077305701561272, + -0.0076750442385673, + -0.0082916170358657, + -0.0079681202769279, + -0.0077468422241508, + -255.6164093017578, + -252.9441680908203 + ], + "max": [ + 0.007688719779253, + 0.0077241472899913, + 0.0075828479602932, + 0.00813714787364, + 0.0083825681358575, + 0.0073839174583554, + 256.2904968261719, + 253.4925842285156 + ] + }, + "outlier": { + "center": [ + -79.66580963134766, + 141535.640625 + ], + "scale": [ + 250.8921127319336, + 222363.16015625 + ] + } +} \ No newline at end of file diff --git a/Generators/share/egconfig/gaussian_params.csv b/Generators/share/egconfig/gaussian_params.csv new file mode 100644 index 0000000000000..8e07c22dd30bf --- /dev/null +++ b/Generators/share/egconfig/gaussian_params.csv @@ -0,0 +1,4 @@ +9.611554230339172022e+01 +1.963570744941765867e+01 +4.300000000000000000e+01 +1.690000000000000000e+02 diff --git a/Generators/share/egconfig/poisson_params.csv b/Generators/share/egconfig/poisson_params.csv new file mode 100644 index 0000000000000..ef26bd973d34c --- /dev/null +++ b/Generators/share/egconfig/poisson_params.csv @@ -0,0 +1,3 @@ +3.165383056343737511e+00 +1.000000000000000000e+00 +1.200000000000000000e+01 diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 9204ede98215e..153ef5cd5e35e 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -17,11 +17,14 @@ #include "SimulationDataFormat/MCEventHeader.h" #include "SimulationDataFormat/ParticleStatus.h" #include "SimulationDataFormat/MCGenProperties.h" +#include #include "FairPrimaryGenerator.h" #include #include #include "TClonesArray.h" #include "TParticle.h" +#include "TSystem.h" +#include "TGrid.h" namespace o2 { @@ -39,6 +42,18 @@ Generator::Generator() : FairGenerator("ALICEo2", "ALICEo2 Generator"), /** default constructor **/ mThisInstanceID = Generator::InstanceCounter; Generator::InstanceCounter++; + auto simConfig = o2::conf::SimConfig::Instance(); + auto noLoops = simConfig.getLoopersVeto(); + if (!noLoops) { + bool transport = (simConfig.getMCEngine() != "O2TrivialMCEngine"); + if (transport) { + bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); + if (tpcActive) { + mAddLoopers = kTRUE; + initLoopersGen(); + } + } + } } /*****************************************************************/ @@ -49,6 +64,102 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na /** constructor **/ mThisInstanceID = Generator::InstanceCounter; Generator::InstanceCounter++; + auto simConfig = o2::conf::SimConfig::Instance(); + auto noLoops = simConfig.getLoopersVeto(); + if (!noLoops) { + bool transport = (simConfig.getMCEngine() != "O2TrivialMCEngine"); + if (transport) { + bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); + if (tpcActive) { + mAddLoopers = kTRUE; + initLoopersGen(); + } + } + } +} + +/*****************************************************************/ + +void Generator::initLoopersGen() +{ +#ifdef GENERATORS_WITH_ONNXRUNTIME + // Expand all environment paths + const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); + std::string model_pairs = gSystem->ExpandPathName(loopersParam.model_pairs.c_str()); + std::string model_compton = gSystem->ExpandPathName(loopersParam.model_compton.c_str()); + const auto& scaler_pair = gSystem->ExpandPathName(loopersParam.scaler_pair.c_str()); + const auto& scaler_compton = gSystem->ExpandPathName(loopersParam.scaler_compton.c_str()); + const auto& poisson = gSystem->ExpandPathName(loopersParam.poisson.c_str()); + const auto& gauss = gSystem->ExpandPathName(loopersParam.gauss.c_str()); + auto flat_gas = loopersParam.flat_gas; + const auto& nFlatGasLoopers = loopersParam.nFlatGasLoopers; + auto fraction_pairs = loopersParam.fraction_pairs; + auto multiplier = loopersParam.multiplier; + auto fixedNLoopers = loopersParam.fixedNLoopers; + const std::array models = {model_pairs, model_compton}; + const std::array local_names = {"WGANpair.onnx", "WGANcompton.onnx"}; + const std::array isAlien = {models[0].starts_with("alien://"), models[1].starts_with("alien://")}; + const std::array isCCDB = {models[0].starts_with("ccdb://"), models[1].starts_with("ccdb://")}; + if (std::any_of(isAlien.begin(), isAlien.end(), [](bool v) { return v; })) { + if (!gGrid) { + TGrid::Connect("alien://"); + if (!gGrid) { + LOG(fatal) << "AliEn connection failed, check token."; + exit(1); + } + } + for (size_t i = 0; i < models.size(); ++i) { + if (isAlien[i] && !TFile::Cp(models[i].c_str(), local_names[i].c_str())) { + LOG(fatal) << "Error: Model file " << models[i] << " does not exist!"; + exit(1); + } + } + } + if (std::any_of(isCCDB.begin(), isCCDB.end(), [](bool v) { return v; })) { + o2::ccdb::CcdbApi ccdb_api; + ccdb_api.init("http://alice-ccdb.cern.ch"); + for (size_t i = 0; i < models.size(); ++i) { + if (isCCDB[i]) { + auto model_path = models[i].substr(7); // Remove "ccdb://" + // Treat filename if provided in the CCDB path + auto extension = model_path.find(".onnx"); + if (extension != std::string::npos) { + auto last_slash = model_path.find_last_of('/'); + model_path = model_path.substr(0, last_slash); + } + std::map filter; + if (!ccdb_api.retrieveBlob(model_path, "./", filter, o2::ccdb::getCurrentTimestamp(), false, local_names[i].c_str())) { + LOG(fatal) << "Error: issues in retrieving " << model_path << " from CCDB!"; + exit(1); + } + } + } + } + model_pairs = isAlien[0] || isCCDB[0] ? local_names[0] : model_pairs; + model_compton = isAlien[1] || isCCDB[1] ? local_names[1] : model_compton; + try { + // Create the TPC loopers generator with the provided parameters + mLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); + + // Configure the generator with flat gas loopers if enabled (default) + if (flat_gas) { + mLoopersGen->setFlatGas(flat_gas, nFlatGasLoopers); + mLoopersGen->setFractionPairs(fraction_pairs); + } else { + // Otherwise, Poisson+Gauss sampling or fixed number of loopers will be used + // Multiplier is applied only with distribution sampling + // This configuration can be used for testing purposes, in all other cases flat gas is recommended + mLoopersGen->SetNLoopers(fixedNLoopers[0], fixedNLoopers[1]); + mLoopersGen->SetMultiplier(multiplier); + } + LOG(info) << "TPC Loopers generator initialized successfully"; + } catch (const std::exception& e) { + LOG(error) << "Failed to initialize TPC Loopers generator: " << e.what(); + mLoopersGen.reset(); + } +#else + LOG(warn) << "ONNX Runtime support not available, cannot initialize TPC loopers generator"; +#endif } /*****************************************************************/ @@ -65,191 +176,230 @@ Bool_t /*****************************************************************/ Bool_t - Generator::ReadEvent(FairPrimaryGenerator* primGen) + Generator::loopers() { - /** read event **/ - - /** endless generate-and-trigger loop **/ - while (true) { - mReadEventCounter++; - - /** clear particle vector **/ - mParticles.clear(); - - /** reset the sub-generator ID **/ - mSubGeneratorId = -1; - - /** generate event **/ - if (!generateEvent()) { - LOG(error) << "ReadEvent failed in generateEvent"; - return kFALSE; - } - - /** import particles **/ - if (!importParticles()) { - LOG(error) << "ReadEvent failed in importParticles"; - return kFALSE; - } - - if (mSubGeneratorsIdToDesc.empty() && mSubGeneratorId > -1) { - LOG(fatal) << "ReadEvent failed because no SubGenerator description given"; - } - - if (!mSubGeneratorsIdToDesc.empty() && mSubGeneratorId < 0) { - LOG(fatal) << "ReadEvent failed because SubGenerator description given but sub-generator not set"; - } - - /** trigger event **/ - if (triggerEvent()) { - mTriggerOkHook(mParticles, mReadEventCounter); - break; - } else { - mTriggerFalseHook(mParticles, mReadEventCounter); - } +#ifdef GENERATORS_WITH_ONNXRUNTIME + if (!mLoopersGen) { + LOG(error) << "Loopers generator not initialized"; + return kFALSE; } - /** add tracks **/ - if (!addTracks(primGen)) { - LOG(error) << "ReadEvent failed in addTracks"; + // Generate loopers using the initialized TPC loopers generator + if (!mLoopersGen->generateEvent()) { + LOG(error) << "Failed to generate loopers event"; return kFALSE; } - - /** update header **/ - auto header = primGen->GetEvent(); - auto o2header = dynamic_cast(header); - if (!header) { - LOG(fatal) << "MC event header is not a 'o2::dataformats::MCEventHeader' object"; + const auto& looperParticles = mLoopersGen->importParticles(); + if (looperParticles.empty()) { + LOG(error) << "Failed to import loopers particles"; return kFALSE; } - updateHeader(o2header); - updateSubGeneratorInformation(o2header); + // Append the generated looper particles to the main particle list + mParticles.insert(mParticles.end(), looperParticles.begin(), looperParticles.end()); - /** success **/ + LOG(debug) << "Added " << looperParticles.size() << " looper particles"; + return kTRUE; +#else + LOG(warn) << "ONNX Runtime support not available, skipping TPC loopers generation"; return kTRUE; +#endif } + /*****************************************************************/ + + Bool_t + Generator::ReadEvent(FairPrimaryGenerator * primGen) + { + /** read event **/ + + /** endless generate-and-trigger loop **/ + while (true) { + mReadEventCounter++; + + /** clear particle vector **/ + mParticles.clear(); + + /** reset the sub-generator ID **/ + mSubGeneratorId = -1; + + /** generate event **/ + if (!generateEvent()) { + LOG(error) << "ReadEvent failed in generateEvent"; + return kFALSE; + } + + /** import particles **/ + if (!importParticles()) { + LOG(error) << "ReadEvent failed in importParticles"; + return kFALSE; + } + + /** Add loopers **/ + if(mAddLoopers){ + if (!loopers()) { + LOG(error) << "ReadEvent failed in loopers"; + return kFALSE; + } + } + + if (mSubGeneratorsIdToDesc.empty() && mSubGeneratorId > -1) { + LOG(fatal) << "ReadEvent failed because no SubGenerator description given"; + } + + if (!mSubGeneratorsIdToDesc.empty() && mSubGeneratorId < 0) { + LOG(fatal) << "ReadEvent failed because SubGenerator description given but sub-generator not set"; + } + + /** trigger event **/ + if (triggerEvent()) { + mTriggerOkHook(mParticles, mReadEventCounter); + break; + } else { + mTriggerFalseHook(mParticles, mReadEventCounter); + } + } -/*****************************************************************/ + /** add tracks **/ + if (!addTracks(primGen)) { + LOG(error) << "ReadEvent failed in addTracks"; + return kFALSE; + } -Bool_t - Generator::addTracks(FairPrimaryGenerator* primGen) -{ - /** add tracks **/ + /** update header **/ + auto header = primGen->GetEvent(); + auto o2header = dynamic_cast(header); + if (!header) { + LOG(fatal) << "MC event header is not a 'o2::dataformats::MCEventHeader' object"; + return kFALSE; + } + updateHeader(o2header); + updateSubGeneratorInformation(o2header); - auto o2primGen = dynamic_cast(primGen); - if (!o2primGen) { - LOG(fatal) << "PrimaryGenerator is not a o2::eventgen::PrimaryGenerator"; - return kFALSE; + /** success **/ + return kTRUE; } - /** loop over particles **/ - for (const auto& particle : mParticles) { - o2primGen->AddTrack(particle.GetPdgCode(), - particle.Px() * mMomentumUnit, - particle.Py() * mMomentumUnit, - particle.Pz() * mMomentumUnit, - particle.Vx() * mPositionUnit, - particle.Vy() * mPositionUnit, - particle.Vz() * mPositionUnit, - particle.GetMother(0), - particle.GetMother(1), - particle.GetDaughter(0), - particle.GetDaughter(1), - particle.TestBit(ParticleStatus::kToBeDone), - particle.Energy() * mEnergyUnit, - particle.T() * mTimeUnit, - particle.GetWeight(), - (TMCProcess)particle.GetUniqueID(), - particle.GetStatusCode()); // generator status information passed as status code field - } + /*****************************************************************/ - /** success **/ - return kTRUE; -} + Bool_t + Generator::addTracks(FairPrimaryGenerator * primGen) + { + /** add tracks **/ -/*****************************************************************/ + auto o2primGen = dynamic_cast(primGen); + if (!o2primGen) { + LOG(fatal) << "PrimaryGenerator is not a o2::eventgen::PrimaryGenerator"; + return kFALSE; + } -Bool_t - Generator::boostEvent() -{ - /** boost event **/ + /** loop over particles **/ + for (const auto& particle : mParticles) { + o2primGen->AddTrack(particle.GetPdgCode(), + particle.Px() * mMomentumUnit, + particle.Py() * mMomentumUnit, + particle.Pz() * mMomentumUnit, + particle.Vx() * mPositionUnit, + particle.Vy() * mPositionUnit, + particle.Vz() * mPositionUnit, + particle.GetMother(0), + particle.GetMother(1), + particle.GetDaughter(0), + particle.GetDaughter(1), + particle.TestBit(ParticleStatus::kToBeDone), + particle.Energy() * mEnergyUnit, + particle.T() * mTimeUnit, + particle.GetWeight(), + (TMCProcess)particle.GetUniqueID(), + particle.GetStatusCode()); // generator status information passed as status code field + } - /** success **/ - return kTRUE; -} + /** success **/ + return kTRUE; + } -/*****************************************************************/ + /*****************************************************************/ -Bool_t - Generator::triggerEvent() -{ - /** trigger event **/ + Bool_t + Generator::boostEvent() + { + /** boost event **/ - /** check trigger presence **/ - if (mTriggers.size() == 0 && mDeepTriggers.size() == 0) { + /** success **/ return kTRUE; } - /** check trigger mode **/ - Bool_t triggered; - if (mTriggerMode == kTriggerOFF) { - return kTRUE; - } else if (mTriggerMode == kTriggerOR) { - triggered = kFALSE; - } else if (mTriggerMode == kTriggerAND) { - triggered = kTRUE; - } else { - return kTRUE; - } + /*****************************************************************/ - /** loop over triggers **/ - for (const auto& trigger : mTriggers) { - auto retval = trigger(mParticles); - if (mTriggerMode == kTriggerOR) { - triggered |= retval; + Bool_t + Generator::triggerEvent() + { + /** trigger event **/ + + /** check trigger presence **/ + if (mTriggers.size() == 0 && mDeepTriggers.size() == 0) { + return kTRUE; } - if (mTriggerMode == kTriggerAND) { - triggered &= retval; + + /** check trigger mode **/ + Bool_t triggered; + if (mTriggerMode == kTriggerOFF) { + return kTRUE; + } else if (mTriggerMode == kTriggerOR) { + triggered = kFALSE; + } else if (mTriggerMode == kTriggerAND) { + triggered = kTRUE; + } else { + return kTRUE; } - } - /** loop over deep triggers **/ - for (const auto& trigger : mDeepTriggers) { - auto retval = trigger(mInterface, mInterfaceName); - if (mTriggerMode == kTriggerOR) { - triggered |= retval; + /** loop over triggers **/ + for (const auto& trigger : mTriggers) { + auto retval = trigger(mParticles); + if (mTriggerMode == kTriggerOR) { + triggered |= retval; + } + if (mTriggerMode == kTriggerAND) { + triggered &= retval; + } } - if (mTriggerMode == kTriggerAND) { - triggered &= retval; + + /** loop over deep triggers **/ + for (const auto& trigger : mDeepTriggers) { + auto retval = trigger(mInterface, mInterfaceName); + if (mTriggerMode == kTriggerOR) { + triggered |= retval; + } + if (mTriggerMode == kTriggerAND) { + triggered &= retval; + } } - } - /** return **/ - return triggered; -} + /** return **/ + return triggered; + } -/*****************************************************************/ + /*****************************************************************/ -void Generator::addSubGenerator(int subGeneratorId, std::string const& subGeneratorDescription) -{ - if (subGeneratorId < 0) { - LOG(fatal) << "Sub-generator IDs must be >= 0, instead, passed value is " << subGeneratorId; + void Generator::addSubGenerator(int subGeneratorId, std::string const& subGeneratorDescription) + { + if (subGeneratorId < 0) { + LOG(fatal) << "Sub-generator IDs must be >= 0, instead, passed value is " << subGeneratorId; + } + mSubGeneratorsIdToDesc.insert({subGeneratorId, subGeneratorDescription}); } - mSubGeneratorsIdToDesc.insert({subGeneratorId, subGeneratorDescription}); -} -/*****************************************************************/ + /*****************************************************************/ -void Generator::updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const -{ - if (mSubGeneratorId < 0) { - return; + void Generator::updateSubGeneratorInformation(o2::dataformats::MCEventHeader * header) const + { + if (mSubGeneratorId < 0) { + return; + } + header->putInfo(o2::mcgenid::GeneratorProperty::SUBGENERATORID, mSubGeneratorId); + header->putInfo>(o2::mcgenid::GeneratorProperty::SUBGENERATORDESCRIPTIONMAP, mSubGeneratorsIdToDesc); } - header->putInfo(o2::mcgenid::GeneratorProperty::SUBGENERATORID, mSubGeneratorId); - header->putInfo>(o2::mcgenid::GeneratorProperty::SUBGENERATORDESCRIPTIONMAP, mSubGeneratorsIdToDesc); -} -/*****************************************************************/ -/*****************************************************************/ + /*****************************************************************/ + /*****************************************************************/ } /* namespace eventgen */ } /* namespace o2 */ diff --git a/Generators/src/GeneratorsLinkDef.h b/Generators/src/GeneratorsLinkDef.h index 2b8d42f86bf9b..97896d8225042 100644 --- a/Generators/src/GeneratorsLinkDef.h +++ b/Generators/src/GeneratorsLinkDef.h @@ -35,6 +35,10 @@ #pragma link C++ class o2::eventgen::GeneratorFromEventPool + ; #pragma link C++ class o2::eventgen::GeneratorEventPoolParam + ; #pragma link C++ class o2::eventgen::EventPoolGenConfig + ; +#ifdef GENERATORS_WITH_ONNXRUNTIME +#pragma link C++ class o2::eventgen::GenTPCLoopers + ; +#pragma link C++ class o2::eventgen::GenTPCLoopersParam + ; +#endif #pragma link C++ class o2::conf::ConfigurableParamPromoter < o2::eventgen::GeneratorEventPoolParam, o2::eventgen::EventPoolGenConfig> + ; #ifdef GENERATORS_WITH_HEPMC3 #pragma link C++ class o2::eventgen::GeneratorHepMC + ; diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx new file mode 100644 index 0000000000000..4eacb7674599c --- /dev/null +++ b/Generators/src/TPCLoopers.cxx @@ -0,0 +1,417 @@ +#include "Generators/TPCLoopers.h" + +// Static Ort::Env instance for multiple onnx model loading +Ort::Env global_env(ORT_LOGGING_LEVEL_WARNING, "GlobalEnv"); + +// This class is responsible for loading the scaler parameters from a JSON file +// and applying the inverse transformation to the generated data. + +void Scaler::load(const std::string &filename) +{ + std::ifstream file(filename); + if (!file.is_open()) { + throw std::runtime_error("Error: Could not open scaler file!"); + } + + std::string json_str((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + file.close(); + + rapidjson::Document doc; + doc.Parse(json_str.c_str()); + + if (doc.HasParseError()) { + throw std::runtime_error("Error: JSON parsing failed!"); + } + + normal_min = jsonArrayToVector(doc["normal"]["min"]); + normal_max = jsonArrayToVector(doc["normal"]["max"]); + outlier_center = jsonArrayToVector(doc["outlier"]["center"]); + outlier_scale = jsonArrayToVector(doc["outlier"]["scale"]); + std::vector normal_min; + std::vector normal_max; + std::vector outlier_center; + std::vector outlier_scale; +} + +std::vector Scaler::inverse_transform(const std::vector &input) +{ + std::vector output; + for (int i = 0; i < input.size(); ++i) + { + if (i < input.size() - 2) + output.push_back(input[i] * (normal_max[i] - normal_min[i]) + normal_min[i]); + else + output.push_back(input[i] * outlier_scale[i - (input.size() - 2)] + outlier_center[i - (input.size() - 2)]); + } + + return output; +} + +std::vector Scaler::jsonArrayToVector(const rapidjson::Value &jsonArray) +{ + std::vector vec; + for (int i = 0; i < jsonArray.Size(); ++i) + { + vec.push_back(jsonArray[i].GetDouble()); + } + return vec; +} + +// This class loads the ONNX model and generates samples using it. + +ONNXGenerator::ONNXGenerator(Ort::Env& shared_env, const std::string& model_path) +: env(shared_env), session(env, model_path.c_str(), Ort::SessionOptions{}) +{ + // Create session options + Ort::SessionOptions session_options; + session = Ort::Session(env, model_path.c_str(), session_options); +} + +std::vector ONNXGenerator::generate_sample() +{ + Ort::AllocatorWithDefaultOptions allocator; + + // Generate a latent vector (z) + std::vector z(100); + for (auto &v : z) + v = rand_gen.Gaus(0.0, 1.0); + + // Prepare input tensor + std::vector input_shape = {1, 100}; + // Get memory information + Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + + // Create input tensor correctly + Ort::Value input_tensor = Ort::Value::CreateTensor( + memory_info, z.data(), z.size(), input_shape.data(), input_shape.size()); + // Run inference + const char *input_names[] = {"z"}; + const char *output_names[] = {"output"}; + auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, 1); + + // Extract output + float *output_data = output_tensors.front().GetTensorMutableData(); + // Get the size of the output tensor + auto output_tensor_info = output_tensors.front().GetTensorTypeAndShapeInfo(); + size_t output_data_size = output_tensor_info.GetElementCount(); // Total number of elements in the tensor + std::vector output; + for (int i = 0; i < output_data_size; ++i) + { + output.push_back(output_data[i]); + } + + return output; +} + +namespace o2 +{ +namespace eventgen +{ + +GenTPCLoopers::GenTPCLoopers(std::string model_pairs, std::string model_compton, + std::string poisson, std::string gauss, std::string scaler_pair, + std::string scaler_compton) +{ + // Checking if the model files exist and are not empty + std::ifstream model_file[2]; + model_file[0].open(model_pairs); + model_file[1].open(model_compton); + if (!model_file[0].is_open() || model_file[0].peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Pairs model file is empty or does not exist!"; + exit(1); + } + if (!model_file[1].is_open() || model_file[1].peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Compton model file is empty or does not exist!"; + exit(1); + } + model_file[0].close(); + model_file[1].close(); + // Checking if the scaler files exist and are not empty + std::ifstream scaler_file[2]; + scaler_file[0].open(scaler_pair); + scaler_file[1].open(scaler_compton); + if (!scaler_file[0].is_open() || scaler_file[0].peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Pairs scaler file is empty or does not exist!"; + exit(1); + } + if (!scaler_file[1].is_open() || scaler_file[1].peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Compton scaler file is empty or does not exist!"; + exit(1); + } + scaler_file[0].close(); + scaler_file[1].close(); + // Checking if the poisson file exists and it's not empty + if (poisson != "") + { + std::ifstream poisson_file(poisson); + if (!poisson_file.is_open() || poisson_file.peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Poisson file is empty or does not exist!"; + exit(1); + } + else + { + poisson_file >> mPoisson[0] >> mPoisson[1] >> mPoisson[2]; + poisson_file.close(); + mPoissonSet = true; + } + } + // Checking if the gauss file exists and it's not empty + if (gauss != "") + { + std::ifstream gauss_file(gauss); + if (!gauss_file.is_open() || gauss_file.peek() == std::ifstream::traits_type::eof()) + { + LOG(fatal) << "Error: Gauss file is empty or does not exist!"; + exit(1); + } + else + { + gauss_file >> mGauss[0] >> mGauss[1] >> mGauss[2] >> mGauss[3]; + gauss_file.close(); + mGaussSet = true; + } + } + mONNX_pair = std::make_unique(global_env, model_pairs); + mScaler_pair = std::make_unique(); + mScaler_pair->load(scaler_pair); + mONNX_compton = std::make_unique(global_env, model_compton); + mScaler_compton = std::make_unique(); + mScaler_compton->load(scaler_compton); +} + +Bool_t GenTPCLoopers::generateEvent() +{ + // Clear the vector of pairs + mGenPairs.clear(); + // Clear the vector of compton electrons + mGenElectrons.clear(); + if (mFlatGas) { + unsigned int nLoopers, nLoopersPairs, nLoopersCompton; + LOG(debug) << "mCurrentEvent is " << mCurrentEvent; + LOG(debug) << "Current event time: " << ((mCurrentEvent < mInteractionTimeRecords.size() - 1) ? std::to_string(mInteractionTimeRecords[mCurrentEvent + 1].bc2ns() - mInteractionTimeRecords[mCurrentEvent].bc2ns()) : std::to_string(mTimeEnd - mInteractionTimeRecords[mCurrentEvent].bc2ns())) << " ns"; + LOG(debug) << "Current time offset wrt BC: " << mInteractionTimeRecords[mCurrentEvent].getTimeOffsetWrtBC() << " ns"; + mTimeLimit = (mCurrentEvent < mInteractionTimeRecords.size() - 1) ? mInteractionTimeRecords[mCurrentEvent + 1].bc2ns() - mInteractionTimeRecords[mCurrentEvent].bc2ns() : mTimeEnd - mInteractionTimeRecords[mCurrentEvent].bc2ns(); + // With flat gas the number of loopers are adapted based on time interval widths + nLoopers = mFlatGasNumber * (mTimeLimit / mIntTimeRecMean); + nLoopersPairs = static_cast(std::round(nLoopers * mLoopsFractionPairs)); + nLoopersCompton = nLoopers - nLoopersPairs; + SetNLoopers(nLoopersPairs, nLoopersCompton); + LOG(info) << "Flat gas loopers: " << nLoopers << " (pairs: " << nLoopersPairs << ", compton: " << nLoopersCompton << ")"; + generateEvent(mTimeLimit); + mCurrentEvent++; + } else { + // Set number of loopers if poissonian params are available + if (mPoissonSet) { + mNLoopersPairs = static_cast(std::round(mMultiplier[0] * PoissonPairs())); + } + if (mGaussSet) { + mNLoopersCompton = static_cast(std::round(mMultiplier[1] * GaussianElectrons())); + } + // Generate pairs + for (int i = 0; i < mNLoopersPairs; ++i) { + std::vector pair = mONNX_pair->generate_sample(); + // Apply the inverse transformation using the scaler + std::vector transformed_pair = mScaler_pair->inverse_transform(pair); + mGenPairs.push_back(transformed_pair); + } + // Generate compton electrons + for (int i = 0; i < mNLoopersCompton; ++i) { + std::vector electron = mONNX_compton->generate_sample(); + // Apply the inverse transformation using the scaler + std::vector transformed_electron = mScaler_compton->inverse_transform(electron); + mGenElectrons.push_back(transformed_electron); + } + } + return true; +} + +Bool_t GenTPCLoopers::generateEvent(double& time_limit) +{ + LOG(info) << "Time constraint for loopers: " << time_limit << " ns"; + // Generate pairs + for (int i = 0; i < mNLoopersPairs; ++i) { + std::vector pair = mONNX_pair->generate_sample(); + // Apply the inverse transformation using the scaler + std::vector transformed_pair = mScaler_pair->inverse_transform(pair); + transformed_pair[9] = gRandom->Uniform(0., time_limit); // Regenerate time, scaling is not needed because time_limit is already in nanoseconds + mGenPairs.push_back(transformed_pair); + } + // Generate compton electrons + for (int i = 0; i < mNLoopersCompton; ++i) { + std::vector electron = mONNX_compton->generate_sample(); + // Apply the inverse transformation using the scaler + std::vector transformed_electron = mScaler_compton->inverse_transform(electron); + transformed_electron[6] = gRandom->Uniform(0., time_limit); // Regenerate time, scaling is not needed because time_limit is already in nanoseconds + mGenElectrons.push_back(transformed_electron); + } + LOG(info) << "Generated Particles with time limit"; + return true; +} + +std::vector GenTPCLoopers::importParticles() +{ + std::vector particles; + // Get looper pairs from the event + for (auto& pair : mGenPairs) { + double px_e, py_e, pz_e, px_p, py_p, pz_p; + double vx, vy, vz, time; + double e_etot, p_etot; + px_e = pair[0]; + py_e = pair[1]; + pz_e = pair[2]; + px_p = pair[3]; + py_p = pair[4]; + pz_p = pair[5]; + vx = pair[6]; + vy = pair[7]; + vz = pair[8]; + time = pair[9]; + e_etot = TMath::Sqrt(px_e * px_e + py_e * py_e + pz_e * pz_e + mMass_e * mMass_e); + p_etot = TMath::Sqrt(px_p * px_p + py_p * py_p + pz_p * pz_p + mMass_p * mMass_p); + // Push the electron + TParticle electron(11, 1, -1, -1, -1, -1, px_e, py_e, pz_e, e_etot, vx, vy, vz, time / 1e9); + electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(electron.GetStatusCode(), 0).fullEncoding); + electron.SetBit(ParticleStatus::kToBeDone, // + o2::mcgenstatus::getHepMCStatusCode(electron.GetStatusCode()) == 1); + particles.push_back(electron); + // Push the positron + TParticle positron(-11, 1, -1, -1, -1, -1, px_p, py_p, pz_p, p_etot, vx, vy, vz, time / 1e9); + positron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(positron.GetStatusCode(), 0).fullEncoding); + positron.SetBit(ParticleStatus::kToBeDone, // + o2::mcgenstatus::getHepMCStatusCode(positron.GetStatusCode()) == 1); + particles.push_back(positron); + } + // Get compton electrons from the event + for (auto& compton : mGenElectrons) { + double px, py, pz; + double vx, vy, vz, time; + double etot; + px = compton[0]; + py = compton[1]; + pz = compton[2]; + vx = compton[3]; + vy = compton[4]; + vz = compton[5]; + time = compton[6]; + etot = TMath::Sqrt(px * px + py * py + pz * pz + mMass_e * mMass_e); + // Push the electron + TParticle electron(11, 1, -1, -1, -1, -1, px, py, pz, etot, vx, vy, vz, time / 1e9); + electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(electron.GetStatusCode(), 0).fullEncoding); + electron.SetBit(ParticleStatus::kToBeDone, // + o2::mcgenstatus::getHepMCStatusCode(electron.GetStatusCode()) == 1); + particles.push_back(electron); + } + + return particles; +} + +unsigned int GenTPCLoopers::PoissonPairs() +{ + unsigned int poissonValue; + do { + // Generate a Poisson-distributed random number with mean mPoisson[0] + poissonValue = mRandGen.Poisson(mPoisson[0]); + } while (poissonValue < mPoisson[1] || poissonValue > mPoisson[2]); // Regenerate if out of range + + return poissonValue; +} + +unsigned int GenTPCLoopers::GaussianElectrons() +{ + unsigned int gaussValue; + do { + // Generate a Normal-distributed random number with mean mGass[0] and stddev mGauss[1] + gaussValue = mRandGen.Gaus(mGauss[0], mGauss[1]); + } while (gaussValue < mGauss[2] || gaussValue > mGauss[3]); // Regenerate if out of range + + return gaussValue; +} + +void GenTPCLoopers::SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_compton) +{ + if (mFlatGas) { + mNLoopersPairs = nsig_pair; + mNLoopersCompton = nsig_compton; + } else { + if (mPoissonSet) { + LOG(info) << "Poissonian parameters correctly loaded."; + } else { + mNLoopersPairs = nsig_pair; + } + if (mGaussSet) { + LOG(info) << "Gaussian parameters correctly loaded."; + } else { + mNLoopersCompton = nsig_compton; + } + } +} + +void GenTPCLoopers::SetMultiplier(std::array& mult) +{ + // Multipliers will work only if the poissonian and gaussian parameters are set + // otherwise they will be ignored + if (mult[0] < 0 || mult[1] < 0) + { + LOG(fatal) << "Error: Multiplier values must be non-negative!"; + exit(1); + } else { + LOG(info) << "Multiplier values set to: Pair = " << mult[0] << ", Compton = " << mult[1]; + mMultiplier[0] = mult[0]; + mMultiplier[1] = mult[1]; + } +} + +void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number) +{ + mFlatGas = flat; + if (mFlatGas) { + if (number < 0) { + LOG(warn) << "Warning: Number of loopers per event must be non-negative! Switching option off."; + mFlatGas = false; + mFlatGasNumber = -1; + } else { + mFlatGasNumber = number; + mContextFile = std::filesystem::exists("collisioncontext.root") ? TFile::Open("collisioncontext.root") : nullptr; + mCollisionContext = mContextFile ? (o2::steer::DigitizationContext*)mContextFile->Get("DigitizationContext") : nullptr; + mInteractionTimeRecords = mCollisionContext ? mCollisionContext->getEventRecords() : std::vector{}; + if (mInteractionTimeRecords.empty()) { + LOG(error) << "Error: No interaction time records found in the collision context!"; + exit(1); + } else { + LOG(info) << "Interaction Time records has " << mInteractionTimeRecords.size() << " entries."; + mCollisionContext->printCollisionSummary(); + } + for (int c = 0; c < mInteractionTimeRecords.size() - 1; c++) { + mIntTimeRecMean += mInteractionTimeRecords[c + 1].bc2ns() - mInteractionTimeRecords[c].bc2ns(); + } + mIntTimeRecMean /= (mInteractionTimeRecords.size() - 1); // Average interaction time record used as reference + const auto& hbfUtils = o2::raw::HBFUtils::Instance(); + // Get the start time of the second orbit after the last interaction record + const auto& lastIR = mInteractionTimeRecords.back(); + o2::InteractionRecord finalOrbitIR(0, lastIR.orbit + 2); // Final orbit, BC = 0 + mTimeEnd = finalOrbitIR.bc2ns(); + LOG(debug) << "Final orbit start time: " << mTimeEnd << " ns while last interaction record time is " << mInteractionTimeRecords.back().bc2ns() << " ns"; + } + } else { + mFlatGasNumber = -1; + } + LOG(info) << "Flat gas loopers: " << (mFlatGas ? "ON" : "OFF") << ", Reference loopers number per event: " << mFlatGasNumber; +} + +void GenTPCLoopers::setFractionPairs(float& fractionPairs) +{ + if (fractionPairs < 0 || fractionPairs > 1) { + LOG(fatal) << "Error: Loops fraction for pairs must be in the range [0, 1]."; + exit(1); + } + mLoopsFractionPairs = fractionPairs; + LOG(info) << "Pairs fraction set to: " << mLoopsFractionPairs; +} + +} // namespace eventgen +} // namespace o2 \ No newline at end of file diff --git a/Generators/src/TPCLoopersParam.cxx b/Generators/src/TPCLoopersParam.cxx new file mode 100644 index 0000000000000..0202a8ced0535 --- /dev/null +++ b/Generators/src/TPCLoopersParam.cxx @@ -0,0 +1,15 @@ +// Copyright 2024-2025 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \author M+Giacalone - September 2025 + +#include "Generators/TPCLoopersParam.h" +O2ParamImpl(o2::eventgen::GenTPCLoopersParam); From 045213b1f559f9288429b356c503ff48570433bd Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Fri, 26 Sep 2025 15:57:46 +0200 Subject: [PATCH 02/10] Various improvements Fixed arrays in TPCLoopersParams and implemented comments --- .../SimConfig/include/SimConfig/SimConfig.h | 3 - Common/SimConfig/src/SimConfig.cxx | 2 - Generators/CMakeLists.txt | 4 +- Generators/include/Generators/Generator.h | 11 +- .../include/Generators/TPCLoopersParam.h | 5 +- Generators/include/TPCLoopers.h | 12 +- Generators/src/Generator.cxx | 392 +++++++++--------- Generators/src/GeneratorsLinkDef.h | 2 +- Generators/src/TPCLoopers.cxx | 3 - 9 files changed, 214 insertions(+), 220 deletions(-) diff --git a/Common/SimConfig/include/SimConfig/SimConfig.h b/Common/SimConfig/include/SimConfig/SimConfig.h index 8642a0e5bc225..be88d9fbd8c33 100644 --- a/Common/SimConfig/include/SimConfig/SimConfig.h +++ b/Common/SimConfig/include/SimConfig/SimConfig.h @@ -52,7 +52,6 @@ struct SimConfigData { std::vector mActiveModules; // list of active modules std::vector mReadoutDetectors; // list of readout detectors std::string mMCEngine; // chosen VMC engine - bool mNoLoopers = false; // Disable automatic TPC loopers std::string mGenerator; // chosen VMC generator std::string mTrigger; // chosen VMC generator trigger unsigned int mNEvents; // number of events to be simulated @@ -139,8 +138,6 @@ class SimConfig // get selected active detectors std::vector const& getActiveModules() const { return mConfigData.mActiveModules; } std::vector const& getReadoutDetectors() const { return mConfigData.mReadoutDetectors; } - // get loopers veto - bool getLoopersVeto() const { return mConfigData.mNoLoopers; } // static helper functions to determine list of active / readout modules // can also be used from outside diff --git a/Common/SimConfig/src/SimConfig.cxx b/Common/SimConfig/src/SimConfig.cxx index e6e0b6adf2093..9407a3c556179 100644 --- a/Common/SimConfig/src/SimConfig.cxx +++ b/Common/SimConfig/src/SimConfig.cxx @@ -74,7 +74,6 @@ void SimConfig::initOptions(boost::program_options::options_description& options "run", bpo::value()->default_value(-1), "ALICE run number")( "asservice", bpo::value()->default_value(false), "run in service/server mode")( "noGeant", bpo::bool_switch(), "prohibits any Geant transport/physics (by using tight cuts)")( - "noLoopers", bpo::bool_switch(), "disable automatic TPC loopers")( "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()->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\". The format is COLLISIONCONTEXTFILE.root[:SIGNALNAME] where SIGNALNAME is the event part in the context which is relevant."); @@ -291,7 +290,6 @@ bool SimConfig::resetFromParsedMap(boost::program_options::variables_map const& using o2::detectors::DetID; mConfigData.mMCEngine = vm["mcEngine"].as(); mConfigData.mNoGeant = vm["noGeant"].as(); - mConfigData.mNoLoopers = vm["noLoopers"].as(); // Reset modules and detectors as they are anyway re-parsed mConfigData.mReadoutDetectors.clear(); diff --git a/Generators/CMakeLists.txt b/Generators/CMakeLists.txt index 56fe8b8fc2284..f1921b8d8d72a 100644 --- a/Generators/CMakeLists.txt +++ b/Generators/CMakeLists.txt @@ -67,7 +67,7 @@ if(HepMC3_FOUND) endif() if(onnxruntime_FOUND) - target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_ONNXRUNTIME) + target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_TPCLOOPERS) endif() set(headers @@ -96,7 +96,7 @@ set(headers ) if(onnxruntime_FOUND) - list(APPEND headers + list(APPEND headers include/Generators/TPCLoopers.h include/Generators/TPCLoopersParam.h) endif() diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h index 374d53f324399..4b68112517893 100644 --- a/Generators/include/Generators/Generator.h +++ b/Generators/include/Generators/Generator.h @@ -17,7 +17,8 @@ #include "FairGenerator.h" #include "TParticle.h" #include "Generators/Trigger.h" -#ifdef GENERATORS_WITH_ONNXRUNTIME +#include "CCDB/BasicCCDBManager.h" +#ifdef GENERATORS_WITH_TPCLOOPERS #include "Generators/TPCLoopers.h" #include "Generators/TPCLoopersParam.h" #endif @@ -77,7 +78,7 @@ class Generator : public FairGenerator /** methods to override **/ virtual Bool_t generateEvent() = 0; // generates event (in structure internal to generator) virtual Bool_t importParticles() = 0; // fills the mParticles vector (transfer from generator state) - Bool_t loopers(); // adds loopers to the event in case TPC is used + Bool_t finalizeEvent(); // final part of event generation that can be customised using external macros virtual void updateHeader(o2::dataformats::MCEventHeader* eventHeader) {}; Bool_t triggerEvent(); @@ -160,7 +161,7 @@ class Generator : public FairGenerator void updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const; // loopers flag - Bool_t mAddLoopers = kFALSE; + Bool_t mAddTPCLoopers = kFALSE; // Flag is automatically set to true if TPC is in readout detectors, loopers are not vetoed and transport is enabled // collect an ID and a short description of sub-generator entities std::unordered_map mSubGeneratorsIdToDesc; // the current ID of the sub-generator used in the current event (if applicable) @@ -169,11 +170,11 @@ class Generator : public FairGenerator // global static information about (upper limit of) number of events to be generated static unsigned int gTotalNEvents; -#ifdef GENERATORS_WITH_ONNXRUNTIME +#ifdef GENERATORS_WITH_TPCLOOPERS // Loopers generator instance std::unique_ptr mLoopersGen = nullptr; -#endif void initLoopersGen(); +#endif ClassDefOverride(Generator, 2); diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index ceeea201538b2..9430f4e05ac6e 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -28,6 +28,7 @@ namespace eventgen ** allow the user to modify them **/ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper { + bool loopersVeto = false; // if true, no loopers are generated std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering std::string poisson = "${O2_ROOT}/share/Generators/egconfig/poisson_params.csv"; // file with Poissonian parameters @@ -37,8 +38,8 @@ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper multiplier = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling - std::array fixedNLoopers = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty + float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling + unsigned int fixedNLoopers[2] = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty O2ParamDef(GenTPCLoopersParam, "GenTPCLoopers"); }; diff --git a/Generators/include/TPCLoopers.h b/Generators/include/TPCLoopers.h index 70146a82baf60..1c1f3585eb3ab 100644 --- a/Generators/include/TPCLoopers.h +++ b/Generators/include/TPCLoopers.h @@ -1,7 +1,7 @@ #ifndef ALICEO2_EVENTGEN_TPCLOOPERS_H_ #define ALICEO2_EVENTGEN_TPCLOOPERS_H_ -#ifdef GENERATORS_WITH_ONNXRUNTIME +#ifdef GENERATORS_WITH_TPCLOOPERS #include #endif #include @@ -19,12 +19,10 @@ #include "TParticle.h" #include -#ifdef GENERATORS_WITH_ONNXRUNTIME +#ifdef GENERATORS_WITH_TPCLOOPERS // Static Ort::Env instance for multiple onnx model loading extern Ort::Env global_env; -#endif -#ifdef GENERATORS_WITH_ONNXRUNTIME // This class is responsible for loading the scaler parameters from a JSON file // and applying the inverse transformation to the generated data. struct Scaler @@ -55,14 +53,14 @@ class ONNXGenerator Ort::Session session; TRandom3 rand_gen; }; -#endif // GENERATORS_WITH_ONNXRUNTIME +#endif // GENERATORS_WITH_TPCLOOPERS namespace o2 { namespace eventgen { -#ifdef GENERATORS_WITH_ONNXRUNTIME +#ifdef GENERATORS_WITH_TPCLOOPERS class GenTPCLoopers { public: @@ -119,7 +117,7 @@ class GenTPCLoopers double mTimeEnd = 0.0; // Time limit for the last event float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs }; -#endif // GENERATORS_WITH_ONNXRUNTIME +#endif // GENERATORS_WITH_TPCLOOPERS } // namespace eventgen } // namespace o2 diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 153ef5cd5e35e..6fc9f378148d3 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -42,18 +42,20 @@ Generator::Generator() : FairGenerator("ALICEo2", "ALICEo2 Generator"), /** default constructor **/ mThisInstanceID = Generator::InstanceCounter; Generator::InstanceCounter++; - auto simConfig = o2::conf::SimConfig::Instance(); - auto noLoops = simConfig.getLoopersVeto(); - if (!noLoops) { +#ifdef GENERATORS_WITH_TPCLOOPERS + const auto& simConfig = o2::conf::SimConfig::Instance(); + const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); + if (!loopersParam.loopersVeto) { bool transport = (simConfig.getMCEngine() != "O2TrivialMCEngine"); if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - mAddLoopers = kTRUE; + mAddTPCLoopers = kTRUE; initLoopersGen(); } } } +#endif } /*****************************************************************/ @@ -64,25 +66,26 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na /** constructor **/ mThisInstanceID = Generator::InstanceCounter; Generator::InstanceCounter++; - auto simConfig = o2::conf::SimConfig::Instance(); - auto noLoops = simConfig.getLoopersVeto(); - if (!noLoops) { +#ifdef GENERATORS_WITH_TPCLOOPERS + const auto& simConfig = o2::conf::SimConfig::Instance(); + const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); + if (!loopersParam.loopersVeto) { bool transport = (simConfig.getMCEngine() != "O2TrivialMCEngine"); if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - mAddLoopers = kTRUE; + mAddTPCLoopers = kTRUE; initLoopersGen(); } } } +#endif } /*****************************************************************/ - +#ifdef GENERATORS_WITH_TPCLOOPERS void Generator::initLoopersGen() { -#ifdef GENERATORS_WITH_ONNXRUNTIME // Expand all environment paths const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); std::string model_pairs = gSystem->ExpandPathName(loopersParam.model_pairs.c_str()); @@ -94,8 +97,9 @@ void Generator::initLoopersGen() auto flat_gas = loopersParam.flat_gas; const auto& nFlatGasLoopers = loopersParam.nFlatGasLoopers; auto fraction_pairs = loopersParam.fraction_pairs; - auto multiplier = loopersParam.multiplier; - auto fixedNLoopers = loopersParam.fixedNLoopers; + std::array multiplier = {loopersParam.multiplier[0], loopersParam.multiplier[1]}; + unsigned int nLoopersPairs = loopersParam.fixedNLoopers[0]; + unsigned int nLoopersCompton = loopersParam.fixedNLoopers[1]; const std::array models = {model_pairs, model_compton}; const std::array local_names = {"WGANpair.onnx", "WGANcompton.onnx"}; const std::array isAlien = {models[0].starts_with("alien://"), models[1].starts_with("alien://")}; @@ -116,8 +120,10 @@ void Generator::initLoopersGen() } } if (std::any_of(isCCDB.begin(), isCCDB.end(), [](bool v) { return v; })) { - o2::ccdb::CcdbApi ccdb_api; - ccdb_api.init("http://alice-ccdb.cern.ch"); + auto& ccdb = o2::ccdb::BasicCCDBManager::instance(); + ccdb.setURL("http://alice-ccdb.cern.ch"); + // Get underlying CCDB API from BasicCCDBManager + auto& ccdb_api = ccdb.getCCDBAccessor(); for (size_t i = 0; i < models.size(); ++i) { if (isCCDB[i]) { auto model_path = models[i].substr(7); // Remove "ccdb://" @@ -149,7 +155,7 @@ void Generator::initLoopersGen() // Otherwise, Poisson+Gauss sampling or fixed number of loopers will be used // Multiplier is applied only with distribution sampling // This configuration can be used for testing purposes, in all other cases flat gas is recommended - mLoopersGen->SetNLoopers(fixedNLoopers[0], fixedNLoopers[1]); + mLoopersGen->SetNLoopers(nLoopersPairs, nLoopersCompton); mLoopersGen->SetMultiplier(multiplier); } LOG(info) << "TPC Loopers generator initialized successfully"; @@ -157,10 +163,8 @@ void Generator::initLoopersGen() LOG(error) << "Failed to initialize TPC Loopers generator: " << e.what(); mLoopersGen.reset(); } -#else - LOG(warn) << "ONNX Runtime support not available, cannot initialize TPC loopers generator"; -#endif } +#endif /*****************************************************************/ @@ -176,230 +180,228 @@ Bool_t /*****************************************************************/ Bool_t - Generator::loopers() + Generator::finalizeEvent() { -#ifdef GENERATORS_WITH_ONNXRUNTIME - if (!mLoopersGen) { - LOG(error) << "Loopers generator not initialized"; - return kFALSE; - } +#ifdef GENERATORS_WITH_TPCLOOPERS + if(mAddTPCLoopers) { + if (!mLoopersGen) { + LOG(error) << "Loopers generator not initialized"; + return kFALSE; + } - // Generate loopers using the initialized TPC loopers generator - if (!mLoopersGen->generateEvent()) { - LOG(error) << "Failed to generate loopers event"; - return kFALSE; - } - const auto& looperParticles = mLoopersGen->importParticles(); - if (looperParticles.empty()) { - LOG(error) << "Failed to import loopers particles"; - return kFALSE; - } - // Append the generated looper particles to the main particle list - mParticles.insert(mParticles.end(), looperParticles.begin(), looperParticles.end()); + // Generate loopers using the initialized TPC loopers generator + if (!mLoopersGen->generateEvent()) { + LOG(error) << "Failed to generate loopers event"; + return kFALSE; + } + const auto& looperParticles = mLoopersGen->importParticles(); + if (looperParticles.empty()) { + LOG(error) << "Failed to import loopers particles"; + return kFALSE; + } + // Append the generated looper particles to the main particle list + mParticles.insert(mParticles.end(), looperParticles.begin(), looperParticles.end()); - LOG(debug) << "Added " << looperParticles.size() << " looper particles"; - return kTRUE; -#else - LOG(warn) << "ONNX Runtime support not available, skipping TPC loopers generation"; - return kTRUE; + LOG(debug) << "Added " << looperParticles.size() << " looper particles"; + } #endif + return kTRUE; } - /*****************************************************************/ - Bool_t - Generator::ReadEvent(FairPrimaryGenerator * primGen) - { - /** read event **/ - - /** endless generate-and-trigger loop **/ - while (true) { - mReadEventCounter++; +/*****************************************************************/ - /** clear particle vector **/ - mParticles.clear(); +Bool_t + Generator::ReadEvent(FairPrimaryGenerator* primGen) +{ + /** read event **/ - /** reset the sub-generator ID **/ - mSubGeneratorId = -1; + /** endless generate-and-trigger loop **/ + while (true) { + mReadEventCounter++; - /** generate event **/ - if (!generateEvent()) { - LOG(error) << "ReadEvent failed in generateEvent"; - return kFALSE; - } + /** clear particle vector **/ + mParticles.clear(); - /** import particles **/ - if (!importParticles()) { - LOG(error) << "ReadEvent failed in importParticles"; - return kFALSE; - } + /** reset the sub-generator ID **/ + mSubGeneratorId = -1; - /** Add loopers **/ - if(mAddLoopers){ - if (!loopers()) { - LOG(error) << "ReadEvent failed in loopers"; - return kFALSE; - } - } + /** generate event **/ + if (!generateEvent()) { + LOG(error) << "ReadEvent failed in generateEvent"; + return kFALSE; + } - if (mSubGeneratorsIdToDesc.empty() && mSubGeneratorId > -1) { - LOG(fatal) << "ReadEvent failed because no SubGenerator description given"; - } + /** import particles **/ + if (!importParticles()) { + LOG(error) << "ReadEvent failed in importParticles"; + return kFALSE; + } - if (!mSubGeneratorsIdToDesc.empty() && mSubGeneratorId < 0) { - LOG(fatal) << "ReadEvent failed because SubGenerator description given but sub-generator not set"; - } + /** Event finalization**/ + if(!finalizeEvent()) { + LOG(error) << "ReadEvent failed in finalizeEvent"; + return kFALSE; + } - /** trigger event **/ - if (triggerEvent()) { - mTriggerOkHook(mParticles, mReadEventCounter); - break; - } else { - mTriggerFalseHook(mParticles, mReadEventCounter); - } + if (mSubGeneratorsIdToDesc.empty() && mSubGeneratorId > -1) { + LOG(fatal) << "ReadEvent failed because no SubGenerator description given"; } - /** add tracks **/ - if (!addTracks(primGen)) { - LOG(error) << "ReadEvent failed in addTracks"; - return kFALSE; + if (!mSubGeneratorsIdToDesc.empty() && mSubGeneratorId < 0) { + LOG(fatal) << "ReadEvent failed because SubGenerator description given but sub-generator not set"; } - /** update header **/ - auto header = primGen->GetEvent(); - auto o2header = dynamic_cast(header); - if (!header) { - LOG(fatal) << "MC event header is not a 'o2::dataformats::MCEventHeader' object"; - return kFALSE; + /** trigger event **/ + if (triggerEvent()) { + mTriggerOkHook(mParticles, mReadEventCounter); + break; + } else { + mTriggerFalseHook(mParticles, mReadEventCounter); } - updateHeader(o2header); - updateSubGeneratorInformation(o2header); + } - /** success **/ - return kTRUE; + /** add tracks **/ + if (!addTracks(primGen)) { + LOG(error) << "ReadEvent failed in addTracks"; + return kFALSE; + } + + /** update header **/ + auto header = primGen->GetEvent(); + auto o2header = dynamic_cast(header); + if (!header) { + LOG(fatal) << "MC event header is not a 'o2::dataformats::MCEventHeader' object"; + return kFALSE; } + updateHeader(o2header); + updateSubGeneratorInformation(o2header); - /*****************************************************************/ + /** success **/ + return kTRUE; +} - Bool_t - Generator::addTracks(FairPrimaryGenerator * primGen) - { - /** add tracks **/ +/*****************************************************************/ - auto o2primGen = dynamic_cast(primGen); - if (!o2primGen) { - LOG(fatal) << "PrimaryGenerator is not a o2::eventgen::PrimaryGenerator"; - return kFALSE; - } +Bool_t + Generator::addTracks(FairPrimaryGenerator* primGen) +{ + /** add tracks **/ - /** loop over particles **/ - for (const auto& particle : mParticles) { - o2primGen->AddTrack(particle.GetPdgCode(), - particle.Px() * mMomentumUnit, - particle.Py() * mMomentumUnit, - particle.Pz() * mMomentumUnit, - particle.Vx() * mPositionUnit, - particle.Vy() * mPositionUnit, - particle.Vz() * mPositionUnit, - particle.GetMother(0), - particle.GetMother(1), - particle.GetDaughter(0), - particle.GetDaughter(1), - particle.TestBit(ParticleStatus::kToBeDone), - particle.Energy() * mEnergyUnit, - particle.T() * mTimeUnit, - particle.GetWeight(), - (TMCProcess)particle.GetUniqueID(), - particle.GetStatusCode()); // generator status information passed as status code field - } + auto o2primGen = dynamic_cast(primGen); + if (!o2primGen) { + LOG(fatal) << "PrimaryGenerator is not a o2::eventgen::PrimaryGenerator"; + return kFALSE; + } - /** success **/ - return kTRUE; + /** loop over particles **/ + for (const auto& particle : mParticles) { + o2primGen->AddTrack(particle.GetPdgCode(), + particle.Px() * mMomentumUnit, + particle.Py() * mMomentumUnit, + particle.Pz() * mMomentumUnit, + particle.Vx() * mPositionUnit, + particle.Vy() * mPositionUnit, + particle.Vz() * mPositionUnit, + particle.GetMother(0), + particle.GetMother(1), + particle.GetDaughter(0), + particle.GetDaughter(1), + particle.TestBit(ParticleStatus::kToBeDone), + particle.Energy() * mEnergyUnit, + particle.T() * mTimeUnit, + particle.GetWeight(), + (TMCProcess)particle.GetUniqueID(), + particle.GetStatusCode()); // generator status information passed as status code field } - /*****************************************************************/ + /** success **/ + return kTRUE; +} + +/*****************************************************************/ - Bool_t - Generator::boostEvent() - { - /** boost event **/ +Bool_t + Generator::boostEvent() +{ + /** boost event **/ + + /** success **/ + return kTRUE; +} + +/*****************************************************************/ + +Bool_t + Generator::triggerEvent() +{ + /** trigger event **/ - /** success **/ + /** check trigger presence **/ + if (mTriggers.size() == 0 && mDeepTriggers.size() == 0) { return kTRUE; } - /*****************************************************************/ - - Bool_t - Generator::triggerEvent() - { - /** trigger event **/ + /** check trigger mode **/ + Bool_t triggered; + if (mTriggerMode == kTriggerOFF) { + return kTRUE; + } else if (mTriggerMode == kTriggerOR) { + triggered = kFALSE; + } else if (mTriggerMode == kTriggerAND) { + triggered = kTRUE; + } else { + return kTRUE; + } - /** check trigger presence **/ - if (mTriggers.size() == 0 && mDeepTriggers.size() == 0) { - return kTRUE; + /** loop over triggers **/ + for (const auto& trigger : mTriggers) { + auto retval = trigger(mParticles); + if (mTriggerMode == kTriggerOR) { + triggered |= retval; } - - /** check trigger mode **/ - Bool_t triggered; - if (mTriggerMode == kTriggerOFF) { - return kTRUE; - } else if (mTriggerMode == kTriggerOR) { - triggered = kFALSE; - } else if (mTriggerMode == kTriggerAND) { - triggered = kTRUE; - } else { - return kTRUE; + if (mTriggerMode == kTriggerAND) { + triggered &= retval; } + } - /** loop over triggers **/ - for (const auto& trigger : mTriggers) { - auto retval = trigger(mParticles); - if (mTriggerMode == kTriggerOR) { - triggered |= retval; - } - if (mTriggerMode == kTriggerAND) { - triggered &= retval; - } + /** loop over deep triggers **/ + for (const auto& trigger : mDeepTriggers) { + auto retval = trigger(mInterface, mInterfaceName); + if (mTriggerMode == kTriggerOR) { + triggered |= retval; } - - /** loop over deep triggers **/ - for (const auto& trigger : mDeepTriggers) { - auto retval = trigger(mInterface, mInterfaceName); - if (mTriggerMode == kTriggerOR) { - triggered |= retval; - } - if (mTriggerMode == kTriggerAND) { - triggered &= retval; - } + if (mTriggerMode == kTriggerAND) { + triggered &= retval; } - - /** return **/ - return triggered; } - /*****************************************************************/ + /** return **/ + return triggered; +} - void Generator::addSubGenerator(int subGeneratorId, std::string const& subGeneratorDescription) - { - if (subGeneratorId < 0) { - LOG(fatal) << "Sub-generator IDs must be >= 0, instead, passed value is " << subGeneratorId; - } - mSubGeneratorsIdToDesc.insert({subGeneratorId, subGeneratorDescription}); +/*****************************************************************/ + +void Generator::addSubGenerator(int subGeneratorId, std::string const& subGeneratorDescription) +{ + if (subGeneratorId < 0) { + LOG(fatal) << "Sub-generator IDs must be >= 0, instead, passed value is " << subGeneratorId; } + mSubGeneratorsIdToDesc.insert({subGeneratorId, subGeneratorDescription}); +} - /*****************************************************************/ +/*****************************************************************/ - void Generator::updateSubGeneratorInformation(o2::dataformats::MCEventHeader * header) const - { - if (mSubGeneratorId < 0) { - return; - } - header->putInfo(o2::mcgenid::GeneratorProperty::SUBGENERATORID, mSubGeneratorId); - header->putInfo>(o2::mcgenid::GeneratorProperty::SUBGENERATORDESCRIPTIONMAP, mSubGeneratorsIdToDesc); +void Generator::updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const +{ + if (mSubGeneratorId < 0) { + return; } + header->putInfo(o2::mcgenid::GeneratorProperty::SUBGENERATORID, mSubGeneratorId); + header->putInfo>(o2::mcgenid::GeneratorProperty::SUBGENERATORDESCRIPTIONMAP, mSubGeneratorsIdToDesc); +} - /*****************************************************************/ - /*****************************************************************/ +/*****************************************************************/ +/*****************************************************************/ } /* namespace eventgen */ } /* namespace o2 */ diff --git a/Generators/src/GeneratorsLinkDef.h b/Generators/src/GeneratorsLinkDef.h index 97896d8225042..24b3f2e452498 100644 --- a/Generators/src/GeneratorsLinkDef.h +++ b/Generators/src/GeneratorsLinkDef.h @@ -35,7 +35,7 @@ #pragma link C++ class o2::eventgen::GeneratorFromEventPool + ; #pragma link C++ class o2::eventgen::GeneratorEventPoolParam + ; #pragma link C++ class o2::eventgen::EventPoolGenConfig + ; -#ifdef GENERATORS_WITH_ONNXRUNTIME +#ifdef GENERATORS_WITH_TPCLOOPERS #pragma link C++ class o2::eventgen::GenTPCLoopers + ; #pragma link C++ class o2::eventgen::GenTPCLoopersParam + ; #endif diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 4eacb7674599c..109461ab71dfa 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -382,9 +382,6 @@ void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number) if (mInteractionTimeRecords.empty()) { LOG(error) << "Error: No interaction time records found in the collision context!"; exit(1); - } else { - LOG(info) << "Interaction Time records has " << mInteractionTimeRecords.size() << " entries."; - mCollisionContext->printCollisionSummary(); } for (int c = 0; c < mInteractionTimeRecords.size() - 1; c++) { mIntTimeRecMean += mInteractionTimeRecords[c + 1].bc2ns() - mInteractionTimeRecords[c].bc2ns(); From bc55d0adb227fe9a886971173867ea08281ea589 Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Sun, 12 Oct 2025 17:24:31 +0200 Subject: [PATCH 03/10] Vetoing loopers for FlatGas and \!collisioncontext --- Generators/include/Generators/Generator.h | 2 +- Generators/src/Generator.cxx | 21 ++++++++++++++++----- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h index 4b68112517893..67277e20736ce 100644 --- a/Generators/include/Generators/Generator.h +++ b/Generators/include/Generators/Generator.h @@ -173,7 +173,7 @@ class Generator : public FairGenerator #ifdef GENERATORS_WITH_TPCLOOPERS // Loopers generator instance std::unique_ptr mLoopersGen = nullptr; - void initLoopersGen(); + bool initLoopersGen(); #endif ClassDefOverride(Generator, 2); diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 6fc9f378148d3..50b11c0c7bb53 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -50,8 +50,9 @@ Generator::Generator() : FairGenerator("ALICEo2", "ALICEo2 Generator"), if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - mAddTPCLoopers = kTRUE; - initLoopersGen(); + if(initLoopersGen()){ + mAddTPCLoopers = kTRUE; + } } } } @@ -74,8 +75,9 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - mAddTPCLoopers = kTRUE; - initLoopersGen(); + if (initLoopersGen()) { + mAddTPCLoopers = kTRUE; + } } } } @@ -84,7 +86,7 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na /*****************************************************************/ #ifdef GENERATORS_WITH_TPCLOOPERS -void Generator::initLoopersGen() +bool Generator::initLoopersGen() { // Expand all environment paths const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); @@ -95,6 +97,14 @@ void Generator::initLoopersGen() const auto& poisson = gSystem->ExpandPathName(loopersParam.poisson.c_str()); const auto& gauss = gSystem->ExpandPathName(loopersParam.gauss.c_str()); auto flat_gas = loopersParam.flat_gas; + if (flat_gas) { + bool isContext = std::filesystem::exists("collisioncontext.root"); + if (!isContext) { + LOG(warning) << "Warning: No collisioncontext.root file found!"; + LOG(warning) << "Loopers will be kept OFF."; + return kFALSE; + } + } const auto& nFlatGasLoopers = loopersParam.nFlatGasLoopers; auto fraction_pairs = loopersParam.fraction_pairs; std::array multiplier = {loopersParam.multiplier[0], loopersParam.multiplier[1]}; @@ -163,6 +173,7 @@ void Generator::initLoopersGen() LOG(error) << "Failed to initialize TPC Loopers generator: " << e.what(); mLoopersGen.reset(); } + return kTRUE; } #endif From 883bc8eac19a21f961a871236ba4c0a478f61898 Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Fri, 24 Oct 2025 16:42:27 +0200 Subject: [PATCH 04/10] Implemented rate and collision system dependence (default) --- .../include/Generators/TPCLoopersParam.h | 4 ++ Generators/include/TPCLoopers.h | 8 ++- Generators/src/Generator.cxx | 28 +++++--- Generators/src/TPCLoopers.cxx | 64 ++++++++++++++++--- 4 files changed, 87 insertions(+), 17 deletions(-) diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index 9430f4e05ac6e..8571013cdec48 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -35,11 +35,15 @@ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper #include "SimulationDataFormat/MCGenProperties.h" #include "TParticle.h" +#include "TF1.h" #include #ifdef GENERATORS_WITH_TPCLOOPERS @@ -82,10 +83,14 @@ class GenTPCLoopers void SetMultiplier(std::array &mult); - void setFlatGas(Bool_t &flat, const Int_t &number = -1); + void setFlatGas(Bool_t& flat, const Int_t& number, const Int_t& nloopers_orbit); void setFractionPairs(float &fractionPairs); + void SetRate(const std::string &rateFile, const bool &isPbPb, const int &intRate); + + void SetAdjust(const float &adjust); + private: std::unique_ptr mONNX_pair = nullptr; std::unique_ptr mONNX_compton = nullptr; @@ -111,6 +116,7 @@ class GenTPCLoopers o2::steer::DigitizationContext *mCollisionContext = nullptr; // Pointer to the digitization context std::vector mInteractionTimeRecords; // Interaction time records from collision context Bool_t mFlatGas = false; // Flag to indicate if flat gas loopers are used + Bool_t mFlatGasOrbit = false; // Flag to indicate if flat gas loopers are per orbit Int_t mFlatGasNumber = -1; // Number of flat gas loopers per event double mIntTimeRecMean = 1.0; // Average interaction time record used for the reference double mTimeLimit = 0.0; // Time limit for the current event diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 50b11c0c7bb53..fea1a38f1a146 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -92,6 +92,7 @@ bool Generator::initLoopersGen() const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); std::string model_pairs = gSystem->ExpandPathName(loopersParam.model_pairs.c_str()); std::string model_compton = gSystem->ExpandPathName(loopersParam.model_compton.c_str()); + std::string nclxrate = gSystem->ExpandPathName(loopersParam.nclxrate.c_str()); const auto& scaler_pair = gSystem->ExpandPathName(loopersParam.scaler_pair.c_str()); const auto& scaler_compton = gSystem->ExpandPathName(loopersParam.scaler_compton.c_str()); const auto& poisson = gSystem->ExpandPathName(loopersParam.poisson.c_str()); @@ -110,10 +111,10 @@ bool Generator::initLoopersGen() std::array multiplier = {loopersParam.multiplier[0], loopersParam.multiplier[1]}; unsigned int nLoopersPairs = loopersParam.fixedNLoopers[0]; unsigned int nLoopersCompton = loopersParam.fixedNLoopers[1]; - const std::array models = {model_pairs, model_compton}; - const std::array local_names = {"WGANpair.onnx", "WGANcompton.onnx"}; - const std::array isAlien = {models[0].starts_with("alien://"), models[1].starts_with("alien://")}; - const std::array isCCDB = {models[0].starts_with("ccdb://"), models[1].starts_with("ccdb://")}; + const std::array models = {model_pairs, model_compton, nclxrate}; + const std::array local_names = {"WGANpair.onnx", "WGANcompton.onnx", "nclxrate.root"}; + const std::array isAlien = {models[0].starts_with("alien://"), models[1].starts_with("alien://"), models[2].starts_with("alien://")}; + const std::array isCCDB = {models[0].starts_with("ccdb://"), models[1].starts_with("ccdb://"), models[2].starts_with("ccdb://")}; if (std::any_of(isAlien.begin(), isAlien.end(), [](bool v) { return v; })) { if (!gGrid) { TGrid::Connect("alien://"); @@ -153,14 +154,25 @@ bool Generator::initLoopersGen() } model_pairs = isAlien[0] || isCCDB[0] ? local_names[0] : model_pairs; model_compton = isAlien[1] || isCCDB[1] ? local_names[1] : model_compton; + nclxrate = isAlien[2] || isCCDB[2] ? local_names[2] : nclxrate; try { // Create the TPC loopers generator with the provided parameters mLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); - - // Configure the generator with flat gas loopers if enabled (default) + auto& colsys = loopersParam.colsys; + auto &intrate = loopersParam.intrate; + // Configure the generator with flat gas loopers defined per orbit with clusters/track info if (flat_gas) { - mLoopersGen->setFlatGas(flat_gas, nFlatGasLoopers); - mLoopersGen->setFractionPairs(fraction_pairs); + if (colsys != "PbPb" && colsys != "pp") { + LOG(fatal) << "Error: collision system must be either 'PbPb' or 'pp'"; + exit(1); + } else { + if (intrate <= 0) { + LOG(fatal) << "Error: interaction rate must be positive!"; + exit(1); + } + mLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); + mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); + } } else { // Otherwise, Poisson+Gauss sampling or fixed number of loopers will be used // Multiplier is applied only with distribution sampling diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 109461ab71dfa..b771b53ed33d2 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -197,7 +197,8 @@ Bool_t GenTPCLoopers::generateEvent() LOG(debug) << "Current time offset wrt BC: " << mInteractionTimeRecords[mCurrentEvent].getTimeOffsetWrtBC() << " ns"; mTimeLimit = (mCurrentEvent < mInteractionTimeRecords.size() - 1) ? mInteractionTimeRecords[mCurrentEvent + 1].bc2ns() - mInteractionTimeRecords[mCurrentEvent].bc2ns() : mTimeEnd - mInteractionTimeRecords[mCurrentEvent].bc2ns(); // With flat gas the number of loopers are adapted based on time interval widths - nLoopers = mFlatGasNumber * (mTimeLimit / mIntTimeRecMean); + // The denominator is either the LHC orbit (if mFlatGasOrbit is true) or the mean interaction time record interval + nLoopers = mFlatGasOrbit ? (mFlatGasNumber * (mTimeLimit / o2::constants::lhc::LHCOrbitNS)) : (mFlatGasNumber * (mTimeLimit / mIntTimeRecMean)); nLoopersPairs = static_cast(std::round(nLoopers * mLoopsFractionPairs)); nLoopersCompton = nLoopers - nLoopersPairs; SetNLoopers(nLoopersPairs, nLoopersCompton); @@ -366,22 +367,34 @@ void GenTPCLoopers::SetMultiplier(std::array& mult) } } -void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number) +void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number = -1, const Int_t& nloopers_orbit = -1) { mFlatGas = flat; if (mFlatGas) { - if (number < 0) { - LOG(warn) << "Warning: Number of loopers per event must be non-negative! Switching option off."; - mFlatGas = false; - mFlatGasNumber = -1; + if (nloopers_orbit > 0) { + mFlatGasOrbit = true; + mFlatGasNumber = nloopers_orbit; + LOG(info) << "Flat gas loopers will be generated using orbit reference."; } else { - mFlatGasNumber = number; + mFlatGasOrbit = false; + if (number < 0) { + LOG(warn) << "Warning: Number of loopers per event must be non-negative! Switching option off."; + mFlatGas = false; + mFlatGasNumber = -1; + } else { + mFlatGasNumber = number; + } + } + if (mFlatGas) { mContextFile = std::filesystem::exists("collisioncontext.root") ? TFile::Open("collisioncontext.root") : nullptr; mCollisionContext = mContextFile ? (o2::steer::DigitizationContext*)mContextFile->Get("DigitizationContext") : nullptr; mInteractionTimeRecords = mCollisionContext ? mCollisionContext->getEventRecords() : std::vector{}; if (mInteractionTimeRecords.empty()) { LOG(error) << "Error: No interaction time records found in the collision context!"; exit(1); + } else { + LOG(info) << "Interaction Time records has " << mInteractionTimeRecords.size() << " entries."; + mCollisionContext->printCollisionSummary(); } for (int c = 0; c < mInteractionTimeRecords.size() - 1; c++) { mIntTimeRecMean += mInteractionTimeRecords[c + 1].bc2ns() - mInteractionTimeRecords[c].bc2ns(); @@ -397,7 +410,7 @@ void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number) } else { mFlatGasNumber = -1; } - LOG(info) << "Flat gas loopers: " << (mFlatGas ? "ON" : "OFF") << ", Reference loopers number per event: " << mFlatGasNumber; + LOG(info) << "Flat gas loopers: " << (mFlatGas ? "ON" : "OFF") << ", Reference loopers number per " << (mFlatGasOrbit ? "orbit " : "event ") << mFlatGasNumber; } void GenTPCLoopers::setFractionPairs(float& fractionPairs) @@ -410,5 +423,40 @@ void GenTPCLoopers::setFractionPairs(float& fractionPairs) LOG(info) << "Pairs fraction set to: " << mLoopsFractionPairs; } +void GenTPCLoopers::SetRate(const std::string &rateFile, const bool &isPbPb = true, const int &intRate = 50000) +{ + // Checking if the rate file exists and is not empty + TFile rate_file(rateFile.c_str(), "READ"); + if (!rate_file.IsOpen() || rate_file.IsZombie()) { + LOG(fatal) << "Error: Rate file is empty or does not exist!"; + exit(1); + } + const char* fitName = isPbPb ? "fitPbPb" : "fitpp"; + auto fit = (TF1*)rate_file.Get(fitName); + if (!fit) { + LOG(fatal) << "Error: Could not find fit function '" << fitName << "' in rate file!"; + exit(1); + } + auto ref = static_cast(std::floor(fit->Eval(intRate / 1000.))); // fit expects rate in kHz + rate_file.Close(); + if (ref <= 0) { + LOG(fatal) << "Computed flat gas number reference per orbit is <=0"; + exit(1); + } else { + LOG(info) << "Set flat gas number to " << ref << " loopers per orbit using " << fitName << " from " << intRate << " Hz interaction rate."; + auto flat = true; + setFlatGas(flat, -1, ref); + } +} + +void GenTPCLoopers::SetAdjust(const float& adjust = 0.f) +{ + if (mFlatGas && mFlatGasOrbit && adjust >= -1.f && adjust != 0.f) { + LOG(info) << "Adjusting flat gas number per orbit by " << adjust * 100.f << "%"; + mFlatGasNumber = static_cast(std::round(mFlatGasNumber * (1.f + adjust))); + LOG(info) << "New flat gas number per orbit: " << mFlatGasNumber; + } +} + } // namespace eventgen } // namespace o2 \ No newline at end of file From c156a31d5f393de5d646aa9415fe10fbacbf504e Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Thu, 20 Nov 2025 22:40:33 +0100 Subject: [PATCH 05/10] Set automatic interaction rate from collision context --- .../include/Generators/TPCLoopersParam.h | 4 ++-- Generators/include/TPCLoopers.h | 3 +++ Generators/src/Generator.cxx | 10 +++++----- Generators/src/TPCLoopers.cxx | 19 +++++++++++++++++-- 4 files changed, 27 insertions(+), 9 deletions(-) diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index 8571013cdec48..74c3cf4cff0ad 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -37,9 +37,9 @@ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper mONNX_pair = nullptr; std::unique_ptr mONNX_compton = nullptr; @@ -122,6 +124,7 @@ class GenTPCLoopers double mTimeLimit = 0.0; // Time limit for the current event double mTimeEnd = 0.0; // Time limit for the last event float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs + int mInteractionRate = 50000; // Interaction rate in Hz }; #endif // GENERATORS_WITH_TPCLOOPERS diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index fea1a38f1a146..18e28e4cc2668 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -107,7 +107,7 @@ bool Generator::initLoopersGen() } } const auto& nFlatGasLoopers = loopersParam.nFlatGasLoopers; - auto fraction_pairs = loopersParam.fraction_pairs; + const auto& fraction_pairs = loopersParam.fraction_pairs; std::array multiplier = {loopersParam.multiplier[0], loopersParam.multiplier[1]}; unsigned int nLoopersPairs = loopersParam.fixedNLoopers[0]; unsigned int nLoopersCompton = loopersParam.fixedNLoopers[1]; @@ -166,10 +166,6 @@ bool Generator::initLoopersGen() LOG(fatal) << "Error: collision system must be either 'PbPb' or 'pp'"; exit(1); } else { - if (intrate <= 0) { - LOG(fatal) << "Error: interaction rate must be positive!"; - exit(1); - } mLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); } @@ -217,6 +213,10 @@ Bool_t LOG(error) << "Failed to generate loopers event"; return kFALSE; } + if (mLoopersGen->getNLoopers() == 0) { + LOG(warning) << "No loopers generated for this event"; + return kTRUE; + } const auto& looperParticles = mLoopersGen->importParticles(); if (looperParticles.empty()) { LOG(error) << "Failed to import loopers particles"; diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index b771b53ed33d2..0fb76fcd8c3a9 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -437,13 +437,28 @@ void GenTPCLoopers::SetRate(const std::string &rateFile, const bool &isPbPb = tr LOG(fatal) << "Error: Could not find fit function '" << fitName << "' in rate file!"; exit(1); } - auto ref = static_cast(std::floor(fit->Eval(intRate / 1000.))); // fit expects rate in kHz + mInteractionRate = intRate; + if (mInteractionRate < 0) { + mContextFile = std::filesystem::exists("collisioncontext.root") ? TFile::Open("collisioncontext.root") : nullptr; + if (!mContextFile || mContextFile->IsZombie()) { + LOG(fatal) << "Error: Interaction rate not provided and collision context file not found!"; + exit(1); + } + mCollisionContext = (o2::steer::DigitizationContext*)mContextFile->Get("DigitizationContext"); + mInteractionRate = std::floor(mCollisionContext->getDigitizerInteractionRate()); + LOG(info) << "Interaction rate retrieved from collision context: " << mInteractionRate << " Hz"; + if (mInteractionRate < 0) { + LOG(fatal) << "Error: Invalid interaction rate retrieved from collision context!"; + exit(1); + } + } + auto ref = static_cast(std::floor(fit->Eval(mInteractionRate / 1000.))); // fit expects rate in kHz rate_file.Close(); if (ref <= 0) { LOG(fatal) << "Computed flat gas number reference per orbit is <=0"; exit(1); } else { - LOG(info) << "Set flat gas number to " << ref << " loopers per orbit using " << fitName << " from " << intRate << " Hz interaction rate."; + LOG(info) << "Set flat gas number to " << ref << " loopers per orbit using " << fitName << " from " << mInteractionRate << " Hz interaction rate."; auto flat = true; setFlatGas(flat, -1, ref); } From 19e77a34afacafea149ce4d49dcd28fe08b232c6 Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Mon, 24 Nov 2025 10:11:42 +0100 Subject: [PATCH 06/10] Fixed bug + cleaned code --- Generators/include/Generators/TPCLoopersParam.h | 4 ++-- Generators/src/Generator.cxx | 4 +--- Generators/src/TPCLoopers.cxx | 4 ---- 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index 74c3cf4cff0ad..24d905c59c967 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -39,8 +39,8 @@ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper multiplier = {loopersParam.multiplier[0], loopersParam.multiplier[1]}; unsigned int nLoopersPairs = loopersParam.fixedNLoopers[0]; unsigned int nLoopersCompton = loopersParam.fixedNLoopers[1]; @@ -170,7 +168,7 @@ bool Generator::initLoopersGen() mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); } } else { - // Otherwise, Poisson+Gauss sampling or fixed number of loopers will be used + // Otherwise, Poisson+Gauss sampling or fixed number of loopers per event will be used // Multiplier is applied only with distribution sampling // This configuration can be used for testing purposes, in all other cases flat gas is recommended mLoopersGen->SetNLoopers(nLoopersPairs, nLoopersCompton); diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 0fb76fcd8c3a9..07af5b25f99f9 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -27,10 +27,6 @@ void Scaler::load(const std::string &filename) normal_max = jsonArrayToVector(doc["normal"]["max"]); outlier_center = jsonArrayToVector(doc["outlier"]["center"]); outlier_scale = jsonArrayToVector(doc["outlier"]["scale"]); - std::vector normal_min; - std::vector normal_max; - std::vector outlier_center; - std::vector outlier_scale; } std::vector Scaler::inverse_transform(const std::vector &input) From b773e933010f4200d281662137f0d6c19ddee66f Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Mon, 24 Nov 2025 20:22:39 +0100 Subject: [PATCH 07/10] Improved logging + colsys check --- Generators/src/Generator.cxx | 31 ++++++++++++++++++++----------- Generators/src/TPCLoopers.cxx | 6 ++++-- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 9e083913c3bc7..9c16c0dfb7e92 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -50,11 +50,15 @@ Generator::Generator() : FairGenerator("ALICEo2", "ALICEo2 Generator"), if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - if(initLoopersGen()){ + if (initLoopersGen()) { mAddTPCLoopers = kTRUE; } + } else { + LOG(info) << "TPC not active in readout detectors: loopers fast generator disabled."; } } + } else { + LOG(info) << "Loopers fast generator turned OFF with veto flag."; } #endif } @@ -78,8 +82,12 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na if (initLoopersGen()) { mAddTPCLoopers = kTRUE; } + } else { + LOG(info) << "TPC not active in readout detectors: loopers fast generator disabled."; } } + } else { + LOG(info) << "Loopers fast generator turned OFF with veto flag."; } #endif } @@ -97,8 +105,14 @@ bool Generator::initLoopersGen() const auto& scaler_compton = gSystem->ExpandPathName(loopersParam.scaler_compton.c_str()); const auto& poisson = gSystem->ExpandPathName(loopersParam.poisson.c_str()); const auto& gauss = gSystem->ExpandPathName(loopersParam.gauss.c_str()); - auto flat_gas = loopersParam.flat_gas; + const auto& flat_gas = loopersParam.flat_gas; + const auto& colsys = loopersParam.colsys; if (flat_gas) { + if (colsys != "PbPb" && colsys != "pp") { + LOG(warning) << "Automatic background loopers configuration supports only 'pp' and 'PbPb' systems."; + LOG(warning) << "Fast loopers generator will remain OFF."; + return kFALSE; + } bool isContext = std::filesystem::exists("collisioncontext.root"); if (!isContext) { LOG(warning) << "Warning: No collisioncontext.root file found!"; @@ -156,17 +170,12 @@ bool Generator::initLoopersGen() try { // Create the TPC loopers generator with the provided parameters mLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); - auto& colsys = loopersParam.colsys; - auto &intrate = loopersParam.intrate; + const auto &intrate = loopersParam.intrate; // Configure the generator with flat gas loopers defined per orbit with clusters/track info + // If intrate is negative (default), automatic IR from collisioncontext.root will be used if (flat_gas) { - if (colsys != "PbPb" && colsys != "pp") { - LOG(fatal) << "Error: collision system must be either 'PbPb' or 'pp'"; - exit(1); - } else { - mLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); - mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); - } + mLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); + mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); } else { // Otherwise, Poisson+Gauss sampling or fixed number of loopers per event will be used // Multiplier is applied only with distribution sampling diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 07af5b25f99f9..ac1123b8d0bbd 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -141,7 +141,7 @@ GenTPCLoopers::GenTPCLoopers(std::string model_pairs, std::string model_compton, scaler_file[0].close(); scaler_file[1].close(); // Checking if the poisson file exists and it's not empty - if (poisson != "") + if (poisson != "" && poisson != "None" && poisson != "none") { std::ifstream poisson_file(poisson); if (!poisson_file.is_open() || poisson_file.peek() == std::ifstream::traits_type::eof()) @@ -157,7 +157,7 @@ GenTPCLoopers::GenTPCLoopers(std::string model_pairs, std::string model_compton, } } // Checking if the gauss file exists and it's not empty - if (gauss != "") + if (gauss != "" && gauss != "None" && gauss != "none") { std::ifstream gauss_file(gauss); if (!gauss_file.is_open() || gauss_file.peek() == std::ifstream::traits_type::eof()) @@ -205,9 +205,11 @@ Bool_t GenTPCLoopers::generateEvent() // Set number of loopers if poissonian params are available if (mPoissonSet) { mNLoopersPairs = static_cast(std::round(mMultiplier[0] * PoissonPairs())); + LOG(debug) << "Generated loopers pairs (Poisson): " << mNLoopersPairs; } if (mGaussSet) { mNLoopersCompton = static_cast(std::round(mMultiplier[1] * GaussianElectrons())); + LOG(debug) << "Generated compton electrons (Gauss): " << mNLoopersCompton; } // Generate pairs for (int i = 0; i < mNLoopersPairs; ++i) { From b81a2ee7e5807d486b8be46b6d0b10f3572f18ce Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Tue, 25 Nov 2025 07:47:04 +0000 Subject: [PATCH 08/10] Please consider the following formatting changes --- Generators/include/Generators/Generator.h | 2 +- .../include/Generators/TPCLoopersParam.h | 24 +- Generators/include/TPCLoopers.h | 155 ++++++----- .../share/egconfig/ScalerComptonParams.json | 52 ++-- .../share/egconfig/ScalerPairParams.json | 64 ++--- Generators/src/Generator.cxx | 6 +- Generators/src/TPCLoopers.cxx | 256 ++++++++---------- 7 files changed, 271 insertions(+), 288 deletions(-) diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h index 67277e20736ce..5a4921e036ca3 100644 --- a/Generators/include/Generators/Generator.h +++ b/Generators/include/Generators/Generator.h @@ -161,7 +161,7 @@ class Generator : public FairGenerator void updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const; // loopers flag - Bool_t mAddTPCLoopers = kFALSE; // Flag is automatically set to true if TPC is in readout detectors, loopers are not vetoed and transport is enabled + Bool_t mAddTPCLoopers = kFALSE; // Flag is automatically set to true if TPC is in readout detectors, loopers are not vetoed and transport is enabled // collect an ID and a short description of sub-generator entities std::unordered_map mSubGeneratorsIdToDesc; // the current ID of the sub-generator used in the current event (if applicable) diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index 24d905c59c967..49c8e5f5927b6 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -28,22 +28,22 @@ namespace eventgen ** allow the user to modify them **/ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper { - bool loopersVeto = false; // if true, no loopers are generated - std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production - std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering - std::string poisson = "${O2_ROOT}/share/Generators/egconfig/poisson_params.csv"; // file with Poissonian parameters - std::string gauss = "${O2_ROOT}/share/Generators/egconfig/gaussian_params.csv"; // file with Gaussian parameters - std::string scaler_pair = "${O2_ROOT}/share/Generators/egconfig/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production + bool loopersVeto = false; // if true, no loopers are generated + std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production + std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering + std::string poisson = "${O2_ROOT}/share/Generators/egconfig/poisson_params.csv"; // file with Poissonian parameters + std::string gauss = "${O2_ROOT}/share/Generators/egconfig/gaussian_params.csv"; // file with Gaussian parameters + std::string scaler_pair = "${O2_ROOT}/share/Generators/egconfig/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production std::string scaler_compton = "${O2_ROOT}/share/Generators/egconfig/ScalerComptonParams.json"; // file with scaler parameters for Compton scattering std::string nclxrate = "ccdb://Users/m/mgiacalo/ClustersTrackRatio"; // file with clusters/rate information per orbit std::string colsys = "PbPb"; // collision system (PbPb or pp) int intrate = -1; // Automatic IR from collision context if -1, else user-defined interaction rate in Hz - bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume - unsigned int nFlatGasLoopers = 500; // number of loopers to be generated per event in case of flat gas [currently unused, kept for possible future debug developments] - float fraction_pairs = 0.08; // fraction of loopers [currently unused, kept for possible future debug developments] - float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling - unsigned int fixedNLoopers[2] = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty - float adjust_flatgas = 0.f; // adjustment for the number of flat gas loopers per orbit (in percentage, e.g. -0.1 = -10%) [-1, inf)] + bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume + unsigned int nFlatGasLoopers = 500; // number of loopers to be generated per event in case of flat gas [currently unused, kept for possible future debug developments] + float fraction_pairs = 0.08; // fraction of loopers [currently unused, kept for possible future debug developments] + float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling + unsigned int fixedNLoopers[2] = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty + float adjust_flatgas = 0.f; // adjustment for the number of flat gas loopers per orbit (in percentage, e.g. -0.1 = -10%) [-1, inf)] O2ParamDef(GenTPCLoopersParam, "GenTPCLoopers"); }; diff --git a/Generators/include/TPCLoopers.h b/Generators/include/TPCLoopers.h index 8a4dc0030aa21..9addcf844e09d 100644 --- a/Generators/include/TPCLoopers.h +++ b/Generators/include/TPCLoopers.h @@ -26,33 +26,32 @@ extern Ort::Env global_env; // This class is responsible for loading the scaler parameters from a JSON file // and applying the inverse transformation to the generated data. -struct Scaler -{ - std::vector normal_min; - std::vector normal_max; - std::vector outlier_center; - std::vector outlier_scale; +struct Scaler { + std::vector normal_min; + std::vector normal_max; + std::vector outlier_center; + std::vector outlier_scale; - void load(const std::string &filename); + void load(const std::string& filename); - std::vector inverse_transform(const std::vector &input); + std::vector inverse_transform(const std::vector& input); -private: - std::vector jsonArrayToVector(const rapidjson::Value &jsonArray); + private: + std::vector jsonArrayToVector(const rapidjson::Value& jsonArray); }; // This class loads the ONNX model and generates samples using it. class ONNXGenerator { -public: - ONNXGenerator(Ort::Env &shared_env, const std::string &model_path); + public: + ONNXGenerator(Ort::Env& shared_env, const std::string& model_path); - std::vector generate_sample(); + std::vector generate_sample(); -private: - Ort::Env &env; - Ort::Session session; - TRandom3 rand_gen; + private: + Ort::Env& env; + Ort::Session session; + TRandom3 rand_gen; }; #endif // GENERATORS_WITH_TPCLOOPERS @@ -64,67 +63,67 @@ namespace eventgen #ifdef GENERATORS_WITH_TPCLOOPERS class GenTPCLoopers { - public: - GenTPCLoopers(std::string model_pairs = "tpcloopmodel.onnx", std::string model_compton = "tpcloopmodelcompton.onnx", - std::string poisson = "poisson.csv", std::string gauss = "gauss.csv", std::string scaler_pair = "scaler_pair.json", - std::string scaler_compton = "scaler_compton.json"); - - Bool_t generateEvent(); - - Bool_t generateEvent(double &time_limit); - - std::vector importParticles(); - - unsigned int PoissonPairs(); - - unsigned int GaussianElectrons(); - - void SetNLoopers(unsigned int &nsig_pair, unsigned int &nsig_compton); - - void SetMultiplier(std::array &mult); - - void setFlatGas(Bool_t& flat, const Int_t& number, const Int_t& nloopers_orbit); - - void setFractionPairs(float &fractionPairs); - - void SetRate(const std::string &rateFile, const bool &isPbPb, const int &intRate); - - void SetAdjust(const float &adjust); - - unsigned int getNLoopers() const { return (mNLoopersPairs + mNLoopersCompton); } - - private: - std::unique_ptr mONNX_pair = nullptr; - std::unique_ptr mONNX_compton = nullptr; - std::unique_ptr mScaler_pair = nullptr; - std::unique_ptr mScaler_compton = nullptr; - double mPoisson[3] = {0.0, 0.0, 0.0}; // Mu, Min and Max of Poissonian - double mGauss[4] = {0.0, 0.0, 0.0, 0.0}; // Mean, Std, Min, Max - std::vector> mGenPairs; - std::vector> mGenElectrons; - unsigned int mNLoopersPairs = -1; - unsigned int mNLoopersCompton = -1; - std::array mMultiplier = {1., 1.}; - bool mPoissonSet = false; - bool mGaussSet = false; - // Random number generator - TRandom3 mRandGen; - // Masses of the electrons and positrons - TDatabasePDG *mPDG = TDatabasePDG::Instance(); - double mMass_e = mPDG->GetParticle(11)->Mass(); - double mMass_p = mPDG->GetParticle(-11)->Mass(); - int mCurrentEvent = 0; // Current event number, used for adaptive loopers - TFile *mContextFile = nullptr; // Input collision context file - o2::steer::DigitizationContext *mCollisionContext = nullptr; // Pointer to the digitization context - std::vector mInteractionTimeRecords; // Interaction time records from collision context - Bool_t mFlatGas = false; // Flag to indicate if flat gas loopers are used - Bool_t mFlatGasOrbit = false; // Flag to indicate if flat gas loopers are per orbit - Int_t mFlatGasNumber = -1; // Number of flat gas loopers per event - double mIntTimeRecMean = 1.0; // Average interaction time record used for the reference - double mTimeLimit = 0.0; // Time limit for the current event - double mTimeEnd = 0.0; // Time limit for the last event - float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs - int mInteractionRate = 50000; // Interaction rate in Hz + public: + GenTPCLoopers(std::string model_pairs = "tpcloopmodel.onnx", std::string model_compton = "tpcloopmodelcompton.onnx", + std::string poisson = "poisson.csv", std::string gauss = "gauss.csv", std::string scaler_pair = "scaler_pair.json", + std::string scaler_compton = "scaler_compton.json"); + + Bool_t generateEvent(); + + Bool_t generateEvent(double& time_limit); + + std::vector importParticles(); + + unsigned int PoissonPairs(); + + unsigned int GaussianElectrons(); + + void SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_compton); + + void SetMultiplier(std::array& mult); + + void setFlatGas(Bool_t& flat, const Int_t& number, const Int_t& nloopers_orbit); + + void setFractionPairs(float& fractionPairs); + + void SetRate(const std::string& rateFile, const bool& isPbPb, const int& intRate); + + void SetAdjust(const float& adjust); + + unsigned int getNLoopers() const { return (mNLoopersPairs + mNLoopersCompton); } + + private: + std::unique_ptr mONNX_pair = nullptr; + std::unique_ptr mONNX_compton = nullptr; + std::unique_ptr mScaler_pair = nullptr; + std::unique_ptr mScaler_compton = nullptr; + double mPoisson[3] = {0.0, 0.0, 0.0}; // Mu, Min and Max of Poissonian + double mGauss[4] = {0.0, 0.0, 0.0, 0.0}; // Mean, Std, Min, Max + std::vector> mGenPairs; + std::vector> mGenElectrons; + unsigned int mNLoopersPairs = -1; + unsigned int mNLoopersCompton = -1; + std::array mMultiplier = {1., 1.}; + bool mPoissonSet = false; + bool mGaussSet = false; + // Random number generator + TRandom3 mRandGen; + // Masses of the electrons and positrons + TDatabasePDG* mPDG = TDatabasePDG::Instance(); + double mMass_e = mPDG->GetParticle(11)->Mass(); + double mMass_p = mPDG->GetParticle(-11)->Mass(); + int mCurrentEvent = 0; // Current event number, used for adaptive loopers + TFile* mContextFile = nullptr; // Input collision context file + o2::steer::DigitizationContext* mCollisionContext = nullptr; // Pointer to the digitization context + std::vector mInteractionTimeRecords; // Interaction time records from collision context + Bool_t mFlatGas = false; // Flag to indicate if flat gas loopers are used + Bool_t mFlatGasOrbit = false; // Flag to indicate if flat gas loopers are per orbit + Int_t mFlatGasNumber = -1; // Number of flat gas loopers per event + double mIntTimeRecMean = 1.0; // Average interaction time record used for the reference + double mTimeLimit = 0.0; // Time limit for the current event + double mTimeEnd = 0.0; // Time limit for the last event + float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs + int mInteractionRate = 50000; // Interaction rate in Hz }; #endif // GENERATORS_WITH_TPCLOOPERS diff --git a/Generators/share/egconfig/ScalerComptonParams.json b/Generators/share/egconfig/ScalerComptonParams.json index d8e654847f46e..157647fee2db7 100644 --- a/Generators/share/egconfig/ScalerComptonParams.json +++ b/Generators/share/egconfig/ScalerComptonParams.json @@ -1,28 +1,28 @@ { - "normal": { - "min": [ - -0.0108811147511005, - -0.0098758740350604, - -0.0103233363479375, - -260.0542297363281, - -259.80059814453125 - ], - "max": [ - 0.0108060473576188, - 0.0103057539090514, - 0.0106524610891938, - 260.0343933105469, - 259.62890625 - ] - }, - "outlier": { - "center": [ - -71.39387130737305, - 96791.23828125 - ], - "scale": [ - 265.9389114379883, - 230762.30981445312 - ] - } + "normal": { + "min": [ + -0.0108811147511005, + -0.0098758740350604, + -0.0103233363479375, + -260.0542297363281, + -259.80059814453125 + ], + "max": [ + 0.0108060473576188, + 0.0103057539090514, + 0.0106524610891938, + 260.0343933105469, + 259.62890625 + ] + }, + "outlier": { + "center": [ + -71.39387130737305, + 96791.23828125 + ], + "scale": [ + 265.9389114379883, + 230762.30981445312 + ] + } } \ No newline at end of file diff --git a/Generators/share/egconfig/ScalerPairParams.json b/Generators/share/egconfig/ScalerPairParams.json index 61434bfa2462e..57cdac421d3f6 100644 --- a/Generators/share/egconfig/ScalerPairParams.json +++ b/Generators/share/egconfig/ScalerPairParams.json @@ -1,34 +1,34 @@ { - "normal": { - "min": [ - -0.0073022879660129, - -0.0077305701561272, - -0.0076750442385673, - -0.0082916170358657, - -0.0079681202769279, - -0.0077468422241508, - -255.6164093017578, - -252.9441680908203 - ], - "max": [ - 0.007688719779253, - 0.0077241472899913, - 0.0075828479602932, - 0.00813714787364, - 0.0083825681358575, - 0.0073839174583554, - 256.2904968261719, - 253.4925842285156 - ] - }, - "outlier": { - "center": [ - -79.66580963134766, - 141535.640625 - ], - "scale": [ - 250.8921127319336, - 222363.16015625 - ] - } + "normal": { + "min": [ + -0.0073022879660129, + -0.0077305701561272, + -0.0076750442385673, + -0.0082916170358657, + -0.0079681202769279, + -0.0077468422241508, + -255.6164093017578, + -252.9441680908203 + ], + "max": [ + 0.007688719779253, + 0.0077241472899913, + 0.0075828479602932, + 0.00813714787364, + 0.0083825681358575, + 0.0073839174583554, + 256.2904968261719, + 253.4925842285156 + ] + }, + "outlier": { + "center": [ + -79.66580963134766, + 141535.640625 + ], + "scale": [ + 250.8921127319336, + 222363.16015625 + ] + } } \ No newline at end of file diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index 9c16c0dfb7e92..ce49254799587 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -170,7 +170,7 @@ bool Generator::initLoopersGen() try { // Create the TPC loopers generator with the provided parameters mLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); - const auto &intrate = loopersParam.intrate; + const auto& intrate = loopersParam.intrate; // Configure the generator with flat gas loopers defined per orbit with clusters/track info // If intrate is negative (default), automatic IR from collisioncontext.root will be used if (flat_gas) { @@ -209,7 +209,7 @@ Bool_t Generator::finalizeEvent() { #ifdef GENERATORS_WITH_TPCLOOPERS - if(mAddTPCLoopers) { + if (mAddTPCLoopers) { if (!mLoopersGen) { LOG(error) << "Loopers generator not initialized"; return kFALSE; @@ -268,7 +268,7 @@ Bool_t } /** Event finalization**/ - if(!finalizeEvent()) { + if (!finalizeEvent()) { LOG(error) << "ReadEvent failed in finalizeEvent"; return kFALSE; } diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index ac1123b8d0bbd..258b6cce07b5b 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -6,7 +6,7 @@ Ort::Env global_env(ORT_LOGGING_LEVEL_WARNING, "GlobalEnv"); // This class is responsible for loading the scaler parameters from a JSON file // and applying the inverse transformation to the generated data. -void Scaler::load(const std::string &filename) +void Scaler::load(const std::string& filename) { std::ifstream file(filename); if (!file.is_open()) { @@ -27,76 +27,73 @@ void Scaler::load(const std::string &filename) normal_max = jsonArrayToVector(doc["normal"]["max"]); outlier_center = jsonArrayToVector(doc["outlier"]["center"]); outlier_scale = jsonArrayToVector(doc["outlier"]["scale"]); -} +} -std::vector Scaler::inverse_transform(const std::vector &input) +std::vector Scaler::inverse_transform(const std::vector& input) { - std::vector output; - for (int i = 0; i < input.size(); ++i) - { - if (i < input.size() - 2) - output.push_back(input[i] * (normal_max[i] - normal_min[i]) + normal_min[i]); - else - output.push_back(input[i] * outlier_scale[i - (input.size() - 2)] + outlier_center[i - (input.size() - 2)]); - } + std::vector output; + for (int i = 0; i < input.size(); ++i) { + if (i < input.size() - 2) + output.push_back(input[i] * (normal_max[i] - normal_min[i]) + normal_min[i]); + else + output.push_back(input[i] * outlier_scale[i - (input.size() - 2)] + outlier_center[i - (input.size() - 2)]); + } - return output; + return output; } -std::vector Scaler::jsonArrayToVector(const rapidjson::Value &jsonArray) +std::vector Scaler::jsonArrayToVector(const rapidjson::Value& jsonArray) { - std::vector vec; - for (int i = 0; i < jsonArray.Size(); ++i) - { - vec.push_back(jsonArray[i].GetDouble()); - } - return vec; + std::vector vec; + for (int i = 0; i < jsonArray.Size(); ++i) { + vec.push_back(jsonArray[i].GetDouble()); + } + return vec; } // This class loads the ONNX model and generates samples using it. ONNXGenerator::ONNXGenerator(Ort::Env& shared_env, const std::string& model_path) -: env(shared_env), session(env, model_path.c_str(), Ort::SessionOptions{}) + : env(shared_env), session(env, model_path.c_str(), Ort::SessionOptions{}) { - // Create session options - Ort::SessionOptions session_options; - session = Ort::Session(env, model_path.c_str(), session_options); + // Create session options + Ort::SessionOptions session_options; + session = Ort::Session(env, model_path.c_str(), session_options); } std::vector ONNXGenerator::generate_sample() { - Ort::AllocatorWithDefaultOptions allocator; - - // Generate a latent vector (z) - std::vector z(100); - for (auto &v : z) - v = rand_gen.Gaus(0.0, 1.0); - - // Prepare input tensor - std::vector input_shape = {1, 100}; - // Get memory information - Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); - - // Create input tensor correctly - Ort::Value input_tensor = Ort::Value::CreateTensor( - memory_info, z.data(), z.size(), input_shape.data(), input_shape.size()); - // Run inference - const char *input_names[] = {"z"}; - const char *output_names[] = {"output"}; - auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, 1); - - // Extract output - float *output_data = output_tensors.front().GetTensorMutableData(); - // Get the size of the output tensor - auto output_tensor_info = output_tensors.front().GetTensorTypeAndShapeInfo(); - size_t output_data_size = output_tensor_info.GetElementCount(); // Total number of elements in the tensor - std::vector output; - for (int i = 0; i < output_data_size; ++i) - { - output.push_back(output_data[i]); - } + Ort::AllocatorWithDefaultOptions allocator; + + // Generate a latent vector (z) + std::vector z(100); + for (auto& v : z) + v = rand_gen.Gaus(0.0, 1.0); + + // Prepare input tensor + std::vector input_shape = {1, 100}; + // Get memory information + Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + + // Create input tensor correctly + Ort::Value input_tensor = Ort::Value::CreateTensor( + memory_info, z.data(), z.size(), input_shape.data(), input_shape.size()); + // Run inference + const char* input_names[] = {"z"}; + const char* output_names[] = {"output"}; + auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, 1); + + // Extract output + float* output_data = output_tensors.front().GetTensorMutableData(); + // Get the size of the output tensor + auto output_tensor_info = output_tensors.front().GetTensorTypeAndShapeInfo(); + size_t output_data_size = output_tensor_info.GetElementCount(); // Total number of elements in the tensor + std::vector output; + for (int i = 0; i < output_data_size; ++i) { + output.push_back(output_data[i]); + } - return output; + return output; } namespace o2 @@ -105,79 +102,67 @@ namespace eventgen { GenTPCLoopers::GenTPCLoopers(std::string model_pairs, std::string model_compton, - std::string poisson, std::string gauss, std::string scaler_pair, - std::string scaler_compton) + std::string poisson, std::string gauss, std::string scaler_pair, + std::string scaler_compton) { - // Checking if the model files exist and are not empty - std::ifstream model_file[2]; - model_file[0].open(model_pairs); - model_file[1].open(model_compton); - if (!model_file[0].is_open() || model_file[0].peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Pairs model file is empty or does not exist!"; - exit(1); - } - if (!model_file[1].is_open() || model_file[1].peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Compton model file is empty or does not exist!"; - exit(1); - } - model_file[0].close(); - model_file[1].close(); - // Checking if the scaler files exist and are not empty - std::ifstream scaler_file[2]; - scaler_file[0].open(scaler_pair); - scaler_file[1].open(scaler_compton); - if (!scaler_file[0].is_open() || scaler_file[0].peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Pairs scaler file is empty or does not exist!"; - exit(1); - } - if (!scaler_file[1].is_open() || scaler_file[1].peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Compton scaler file is empty or does not exist!"; - exit(1); - } - scaler_file[0].close(); - scaler_file[1].close(); - // Checking if the poisson file exists and it's not empty - if (poisson != "" && poisson != "None" && poisson != "none") - { - std::ifstream poisson_file(poisson); - if (!poisson_file.is_open() || poisson_file.peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Poisson file is empty or does not exist!"; - exit(1); - } - else - { - poisson_file >> mPoisson[0] >> mPoisson[1] >> mPoisson[2]; - poisson_file.close(); - mPoissonSet = true; - } + // Checking if the model files exist and are not empty + std::ifstream model_file[2]; + model_file[0].open(model_pairs); + model_file[1].open(model_compton); + if (!model_file[0].is_open() || model_file[0].peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Pairs model file is empty or does not exist!"; + exit(1); + } + if (!model_file[1].is_open() || model_file[1].peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Compton model file is empty or does not exist!"; + exit(1); + } + model_file[0].close(); + model_file[1].close(); + // Checking if the scaler files exist and are not empty + std::ifstream scaler_file[2]; + scaler_file[0].open(scaler_pair); + scaler_file[1].open(scaler_compton); + if (!scaler_file[0].is_open() || scaler_file[0].peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Pairs scaler file is empty or does not exist!"; + exit(1); + } + if (!scaler_file[1].is_open() || scaler_file[1].peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Compton scaler file is empty or does not exist!"; + exit(1); + } + scaler_file[0].close(); + scaler_file[1].close(); + // Checking if the poisson file exists and it's not empty + if (poisson != "" && poisson != "None" && poisson != "none") { + std::ifstream poisson_file(poisson); + if (!poisson_file.is_open() || poisson_file.peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Poisson file is empty or does not exist!"; + exit(1); + } else { + poisson_file >> mPoisson[0] >> mPoisson[1] >> mPoisson[2]; + poisson_file.close(); + mPoissonSet = true; } - // Checking if the gauss file exists and it's not empty - if (gauss != "" && gauss != "None" && gauss != "none") - { - std::ifstream gauss_file(gauss); - if (!gauss_file.is_open() || gauss_file.peek() == std::ifstream::traits_type::eof()) - { - LOG(fatal) << "Error: Gauss file is empty or does not exist!"; - exit(1); - } - else - { - gauss_file >> mGauss[0] >> mGauss[1] >> mGauss[2] >> mGauss[3]; - gauss_file.close(); - mGaussSet = true; - } + } + // Checking if the gauss file exists and it's not empty + if (gauss != "" && gauss != "None" && gauss != "none") { + std::ifstream gauss_file(gauss); + if (!gauss_file.is_open() || gauss_file.peek() == std::ifstream::traits_type::eof()) { + LOG(fatal) << "Error: Gauss file is empty or does not exist!"; + exit(1); + } else { + gauss_file >> mGauss[0] >> mGauss[1] >> mGauss[2] >> mGauss[3]; + gauss_file.close(); + mGaussSet = true; } - mONNX_pair = std::make_unique(global_env, model_pairs); - mScaler_pair = std::make_unique(); - mScaler_pair->load(scaler_pair); - mONNX_compton = std::make_unique(global_env, model_compton); - mScaler_compton = std::make_unique(); - mScaler_compton->load(scaler_compton); + } + mONNX_pair = std::make_unique(global_env, model_pairs); + mScaler_pair = std::make_unique(); + mScaler_pair->load(scaler_pair); + mONNX_compton = std::make_unique(global_env, model_compton); + mScaler_compton = std::make_unique(); + mScaler_compton->load(scaler_compton); } Bool_t GenTPCLoopers::generateEvent() @@ -352,17 +337,16 @@ void GenTPCLoopers::SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_comp void GenTPCLoopers::SetMultiplier(std::array& mult) { - // Multipliers will work only if the poissonian and gaussian parameters are set - // otherwise they will be ignored - if (mult[0] < 0 || mult[1] < 0) - { - LOG(fatal) << "Error: Multiplier values must be non-negative!"; - exit(1); - } else { - LOG(info) << "Multiplier values set to: Pair = " << mult[0] << ", Compton = " << mult[1]; - mMultiplier[0] = mult[0]; - mMultiplier[1] = mult[1]; - } + // Multipliers will work only if the poissonian and gaussian parameters are set + // otherwise they will be ignored + if (mult[0] < 0 || mult[1] < 0) { + LOG(fatal) << "Error: Multiplier values must be non-negative!"; + exit(1); + } else { + LOG(info) << "Multiplier values set to: Pair = " << mult[0] << ", Compton = " << mult[1]; + mMultiplier[0] = mult[0]; + mMultiplier[1] = mult[1]; + } } void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number = -1, const Int_t& nloopers_orbit = -1) @@ -421,7 +405,7 @@ void GenTPCLoopers::setFractionPairs(float& fractionPairs) LOG(info) << "Pairs fraction set to: " << mLoopsFractionPairs; } -void GenTPCLoopers::SetRate(const std::string &rateFile, const bool &isPbPb = true, const int &intRate = 50000) +void GenTPCLoopers::SetRate(const std::string& rateFile, const bool& isPbPb = true, const int& intRate = 50000) { // Checking if the rate file exists and is not empty TFile rate_file(rateFile.c_str(), "READ"); From 797f283e8954fccb52e2854d31eecfe4d480254e Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Tue, 25 Nov 2025 08:54:19 +0100 Subject: [PATCH 09/10] Add copyright headers --- Generators/include/TPCLoopers.h | 13 +++++++++++++ Generators/src/TPCLoopers.cxx | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/Generators/include/TPCLoopers.h b/Generators/include/TPCLoopers.h index 9addcf844e09d..57d178667b497 100644 --- a/Generators/include/TPCLoopers.h +++ b/Generators/include/TPCLoopers.h @@ -1,3 +1,16 @@ +// Copyright 2024-2025 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \author M+Giacalone - September 2025 + #ifndef ALICEO2_EVENTGEN_TPCLOOPERS_H_ #define ALICEO2_EVENTGEN_TPCLOOPERS_H_ diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 258b6cce07b5b..8dff795de40a3 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -1,3 +1,16 @@ +// Copyright 2024-2025 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \author M+Giacalone - September 2025 + #include "Generators/TPCLoopers.h" // Static Ort::Env instance for multiple onnx model loading From ba6a1a7cf32ef49db4bfd14add1c60a298f449c7 Mon Sep 17 00:00:00 2001 From: Marco Giacalone Date: Tue, 25 Nov 2025 12:55:13 +0100 Subject: [PATCH 10/10] Corrected include folder --- Generators/CMakeLists.txt | 19 ++--- Generators/include/Generators/Generator.h | 5 +- .../include/{ => Generators}/TPCLoopers.h | 44 ++++++----- .../include/Generators/TPCLoopersParam.h | 35 ++++---- Generators/share/TPCLoopers/README.md | 79 +++++++++++++++++++ .../ScalerComptonParams.json | 0 .../ScalerPairParams.json | 0 .../gaussian_params.csv | 0 .../poisson_params.csv | 0 Generators/src/Generator.cxx | 28 ++++--- Generators/src/TPCLoopers.cxx | 40 +++++++--- 11 files changed, 172 insertions(+), 78 deletions(-) rename Generators/include/{ => Generators}/TPCLoopers.h (71%) create mode 100644 Generators/share/TPCLoopers/README.md rename Generators/share/{egconfig => TPCLoopers}/ScalerComptonParams.json (100%) rename Generators/share/{egconfig => TPCLoopers}/ScalerPairParams.json (100%) rename Generators/share/{egconfig => TPCLoopers}/gaussian_params.csv (100%) rename Generators/share/{egconfig => TPCLoopers}/poisson_params.csv (100%) diff --git a/Generators/CMakeLists.txt b/Generators/CMakeLists.txt index f1921b8d8d72a..287536ff118f7 100644 --- a/Generators/CMakeLists.txt +++ b/Generators/CMakeLists.txt @@ -41,8 +41,8 @@ o2_add_library(Generators src/GeneratorTParticleParam.cxx src/GeneratorService.cxx src/FlowMapper.cxx - $<$:src/TPCLoopers.cxx> - $<$:src/TPCLoopersParam.cxx> + src/TPCLoopers.cxx + src/TPCLoopersParam.cxx $<$:src/GeneratorPythia8.cxx> $<$:src/DecayerPythia8.cxx> $<$:src/GeneratorPythia8Param.cxx> @@ -55,7 +55,7 @@ o2_add_library(Generators PUBLIC_LINK_LIBRARIES FairRoot::Base O2::SimConfig O2::CommonUtils O2::DetectorsBase O2::ZDCBase O2::SimulationDataFormat ${pythiaTarget} ${hepmcTarget} FairRoot::Gen - $<$:onnxruntime::onnxruntime> + onnxruntime::onnxruntime TARGETVARNAME targetName) if(pythia_FOUND) @@ -66,9 +66,7 @@ if(HepMC3_FOUND) target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_HEPMC3) endif() -if(onnxruntime_FOUND) - target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_TPCLOOPERS) -endif() +target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_TPCLOOPERS) set(headers include/Generators/Generator.h @@ -95,11 +93,9 @@ set(headers include/Generators/FlowMapper.h ) -if(onnxruntime_FOUND) - list(APPEND headers - include/Generators/TPCLoopers.h - include/Generators/TPCLoopersParam.h) -endif() +list(APPEND headers + include/Generators/TPCLoopers.h + include/Generators/TPCLoopersParam.h) if(pythia_FOUND) list(APPEND headers @@ -171,4 +167,5 @@ endif() o2_data_file(COPY share/external DESTINATION Generators) o2_data_file(COPY share/egconfig DESTINATION Generators) +o2_data_file(COPY share/TPCLoopers DESTINATION Generators) o2_data_file(COPY share/pythia8 DESTINATION Generators) diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h index 5a4921e036ca3..3484601aa42bb 100644 --- a/Generators/include/Generators/Generator.h +++ b/Generators/include/Generators/Generator.h @@ -17,7 +17,6 @@ #include "FairGenerator.h" #include "TParticle.h" #include "Generators/Trigger.h" -#include "CCDB/BasicCCDBManager.h" #ifdef GENERATORS_WITH_TPCLOOPERS #include "Generators/TPCLoopers.h" #include "Generators/TPCLoopersParam.h" @@ -172,8 +171,8 @@ class Generator : public FairGenerator #ifdef GENERATORS_WITH_TPCLOOPERS // Loopers generator instance - std::unique_ptr mLoopersGen = nullptr; - bool initLoopersGen(); + std::unique_ptr mTPCLoopersGen = nullptr; + bool initTPCLoopersGen(); #endif ClassDefOverride(Generator, 2); diff --git a/Generators/include/TPCLoopers.h b/Generators/include/Generators/TPCLoopers.h similarity index 71% rename from Generators/include/TPCLoopers.h rename to Generators/include/Generators/TPCLoopers.h index 57d178667b497..6a1d3ef262e22 100644 --- a/Generators/include/TPCLoopers.h +++ b/Generators/include/Generators/TPCLoopers.h @@ -17,21 +17,11 @@ #ifdef GENERATORS_WITH_TPCLOOPERS #include #endif -#include #include -#include #include -#include "CCDB/CCDBTimeStampUtils.h" -#include "CCDB/CcdbApi.h" -#include "DetectorsRaw/HBFUtils.h" #include "TRandom3.h" -#include "TDatabasePDG.h" #include -#include -#include "SimulationDataFormat/MCGenProperties.h" #include "TParticle.h" -#include "TF1.h" -#include #ifdef GENERATORS_WITH_TPCLOOPERS // Static Ort::Env instance for multiple onnx model loading @@ -39,6 +29,8 @@ extern Ort::Env global_env; // This class is responsible for loading the scaler parameters from a JSON file // and applying the inverse transformation to the generated data. +// Inferenced output is scaled (min-max normalization or robust scaling for outlier features) during training, +// so we need to revert this transformation to get physical values. struct Scaler { std::vector normal_min; std::vector normal_max; @@ -74,6 +66,20 @@ namespace eventgen { #ifdef GENERATORS_WITH_TPCLOOPERS +/** + * Generator for TPC Loopers based on pre-trained ONNX models. + * Currently it generates loopers as electron-positron pairs and Compton electrons + * according to specified distributions and parameters. + * This can be extended to other types of background processes in the future (e.g. slow neutron spallation products, saturation tail). + * Multiple configuration options are available: + * - Flat gas: loopers are generated uniformly per event taking a reference value which can be either the LHC orbit time or the average interaction time record interval from the collision context. + * ==> Current automatic setup (default) sets the interaction rate automatically from the collision context and the reference value per orbit is calculated from an external file. + * ==> Number of loopers per orbit can be adjusted via a specific parameter. + * - Poisson + Gaussian sampling: number of loopers are sampled from Poissonian (for pairs) and Gaussian (for Compton electrons) distributions based on provided parameters. + * ==> flat gas must be disabled to use this option. + * - Fixed number of loopers per event + * ==> flat gas must be disabled to use this option and Poissonian/Gaussian parameters file should be set to None + */ class GenTPCLoopers { public: @@ -83,7 +89,7 @@ class GenTPCLoopers Bool_t generateEvent(); - Bool_t generateEvent(double& time_limit); + Bool_t generateEvent(double time_limit); std::vector importParticles(); @@ -91,17 +97,17 @@ class GenTPCLoopers unsigned int GaussianElectrons(); - void SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_compton); + void SetNLoopers(unsigned int nsig_pair, unsigned int nsig_compton); - void SetMultiplier(std::array& mult); + void SetMultiplier(const std::array& mult); - void setFlatGas(Bool_t& flat, const Int_t& number, const Int_t& nloopers_orbit); + void setFlatGas(Bool_t flat, Int_t number = -1, Int_t nloopers_orbit = -1); - void setFractionPairs(float& fractionPairs); + void setFractionPairs(float fractionPairs); - void SetRate(const std::string& rateFile, const bool& isPbPb, const int& intRate); + void SetRate(const std::string& rateFile, bool isPbPb, int intRate = 50000); - void SetAdjust(const float& adjust); + void SetAdjust(float adjust = 0.f); unsigned int getNLoopers() const { return (mNLoopersPairs + mNLoopersCompton); } @@ -121,10 +127,6 @@ class GenTPCLoopers bool mGaussSet = false; // Random number generator TRandom3 mRandGen; - // Masses of the electrons and positrons - TDatabasePDG* mPDG = TDatabasePDG::Instance(); - double mMass_e = mPDG->GetParticle(11)->Mass(); - double mMass_p = mPDG->GetParticle(-11)->Mass(); int mCurrentEvent = 0; // Current event number, used for adaptive loopers TFile* mContextFile = nullptr; // Input collision context file o2::steer::DigitizationContext* mCollisionContext = nullptr; // Pointer to the digitization context diff --git a/Generators/include/Generators/TPCLoopersParam.h b/Generators/include/Generators/TPCLoopersParam.h index 49c8e5f5927b6..87e4510d6e617 100644 --- a/Generators/include/Generators/TPCLoopersParam.h +++ b/Generators/include/Generators/TPCLoopersParam.h @@ -24,26 +24,27 @@ namespace eventgen /** ** a parameter class/struct to keep the settings of - ** the tpc loopers event-generator and + ** the TPC loopers event-generator and ** allow the user to modify them **/ struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper { - bool loopersVeto = false; // if true, no loopers are generated - std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production - std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering - std::string poisson = "${O2_ROOT}/share/Generators/egconfig/poisson_params.csv"; // file with Poissonian parameters - std::string gauss = "${O2_ROOT}/share/Generators/egconfig/gaussian_params.csv"; // file with Gaussian parameters - std::string scaler_pair = "${O2_ROOT}/share/Generators/egconfig/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production - std::string scaler_compton = "${O2_ROOT}/share/Generators/egconfig/ScalerComptonParams.json"; // file with scaler parameters for Compton scattering - std::string nclxrate = "ccdb://Users/m/mgiacalo/ClustersTrackRatio"; // file with clusters/rate information per orbit - std::string colsys = "PbPb"; // collision system (PbPb or pp) - int intrate = -1; // Automatic IR from collision context if -1, else user-defined interaction rate in Hz - bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume - unsigned int nFlatGasLoopers = 500; // number of loopers to be generated per event in case of flat gas [currently unused, kept for possible future debug developments] - float fraction_pairs = 0.08; // fraction of loopers [currently unused, kept for possible future debug developments] - float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling - unsigned int fixedNLoopers[2] = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty - float adjust_flatgas = 0.f; // adjustment for the number of flat gas loopers per orbit (in percentage, e.g. -0.1 = -10%) [-1, inf)] + bool loopersVeto = false; // if true, no loopers are generated + // Current files are set to custom user CCDB paths, TO BE CHANGED + std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production + std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering + std::string poisson = "${O2_ROOT}/share/Generators/TPCLoopers/poisson_params.csv"; // file with Poissonian parameters + std::string gauss = "${O2_ROOT}/share/Generators/TPCLoopers/gaussian_params.csv"; // file with Gaussian parameters + std::string scaler_pair = "${O2_ROOT}/share/Generators/TPCLoopers/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production + std::string scaler_compton = "${O2_ROOT}/share/Generators/TPCLoopers/ScalerComptonParams.json"; // file with scaler parameters for Compton scattering + std::string nclxrate = "ccdb://Users/m/mgiacalo/ClustersTrackRatio"; // file with clusters/rate information per orbit + std::string colsys = "PbPb"; // collision system (PbPb or pp) + int intrate = -1; // Automatic IR from collision context if -1, else user-defined interaction rate in Hz + bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume + unsigned int nFlatGasLoopers = 500; // number of loopers to be generated per event in case of flat gas [currently unused, kept for possible future debug developments] + float fraction_pairs = 0.08; // fraction of loopers [currently unused, kept for possible future debug developments] + float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling + unsigned int fixedNLoopers[2] = {1, 1}; // fixed number of loopers coming from pairs and compton electrons - valid if flat gas is false and both Poisson and Gaussian params files are empty + float adjust_flatgas = 0.f; // adjustment for the number of flat gas loopers per orbit (in percentage, e.g. -0.1 = -10%) [-1, inf)] O2ParamDef(GenTPCLoopersParam, "GenTPCLoopers"); }; diff --git a/Generators/share/TPCLoopers/README.md b/Generators/share/TPCLoopers/README.md new file mode 100644 index 0000000000000..0e0ac858b8809 --- /dev/null +++ b/Generators/share/TPCLoopers/README.md @@ -0,0 +1,79 @@ +# TPC Loopers Generator - Parameter Files + +This directory contains parameter files used by the TPC Loopers event generator in ALICE O2. + +## Overview + +The TPC Loopers generator uses pre-trained ONNX models to generate realistic looper particles based on machine learning models trained on full GEANT4 slow neutron transport simulations. The parameter files in this directory provide: +- Example statistical distribution parameters for sampling the number of loopers per event +- **Mandatory** scaling parameters for transforming the ONNX model outputs to physical values + +## Files Description + +### Statistical Sampling Parameters + +The files provided in the folder are examples based on the training dataset. + +#### `gaussian_params.csv` +Parameters for Gaussian distribution used to sample the number of Compton electrons per event. + +**Format:** Four values (one per line) +1. Mean (μ) +2. Standard deviation (σ) +3. Minimum value +4. Maximum value + +#### `poisson_params.csv` +Parameters for Poisson distribution used to sample the number of electron-positron pairs per event. + +**Format:** Three values (one per line) +1. Lambda (λ) parameter +2. Minimum value +3. Maximum value + +### Scaler Parameters + +These JSON files contain the parameters for inverse transformation of the ONNX models output. They should be kept as they are +unless a new version of the models is released. + +#### `ScalerComptonParams.json` +Scaler parameters for Compton electron generation model. + +**Structure:** +```json +{ + "normal": { + "min": [array of 5 min values for min-max normalization], + "max": [array of 5 max values for min-max normalization] + }, + "outlier": { + "center": [array of 2 center values for robust scaling], + "scale": [array of 2 scale values for robust scaling] + } +} +``` + +- **normal**: Min-max normalization parameters for standard features (`Px`, `Py`, `Pz`, `VertexCoordinatesX`, `VertexCoordinatesY`) +- **outlier**: Robust scaler parameters (center and scale) for outlier features (`VertexCoordinatesZ`,`time`) + +#### `ScalerPairParams.json` +Scaler parameters for electron-positron pair generation model. + +**Structure:** +```json +{ + "normal": { + "min": [array of 8 min values for min-max normalization], + "max": [array of 8 max values for min-max normalization] + }, + "outlier": { + "center": [array of 2 center values for robust scaling], + "scale": [array of 2 scale values for robust scaling] + } +} +``` + +- **normal**: Min-max normalization parameters for standard features (`Px_e`, `Py_e`, `Pz_e`,`Px_p`, `Py_p`, `Pz_p`, `VertexCoordinatesX`, `VertexCoordinatesY`) +- **outlier**: Robust scaler parameters (center and scale) for outlier features (`VertexCoordinatesZ`,`time`) +--- +*Author: M. Giacalone - September 2025* diff --git a/Generators/share/egconfig/ScalerComptonParams.json b/Generators/share/TPCLoopers/ScalerComptonParams.json similarity index 100% rename from Generators/share/egconfig/ScalerComptonParams.json rename to Generators/share/TPCLoopers/ScalerComptonParams.json diff --git a/Generators/share/egconfig/ScalerPairParams.json b/Generators/share/TPCLoopers/ScalerPairParams.json similarity index 100% rename from Generators/share/egconfig/ScalerPairParams.json rename to Generators/share/TPCLoopers/ScalerPairParams.json diff --git a/Generators/share/egconfig/gaussian_params.csv b/Generators/share/TPCLoopers/gaussian_params.csv similarity index 100% rename from Generators/share/egconfig/gaussian_params.csv rename to Generators/share/TPCLoopers/gaussian_params.csv diff --git a/Generators/share/egconfig/poisson_params.csv b/Generators/share/TPCLoopers/poisson_params.csv similarity index 100% rename from Generators/share/egconfig/poisson_params.csv rename to Generators/share/TPCLoopers/poisson_params.csv diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx index ce49254799587..465a8ffb7ee22 100644 --- a/Generators/src/Generator.cxx +++ b/Generators/src/Generator.cxx @@ -25,6 +25,8 @@ #include "TParticle.h" #include "TSystem.h" #include "TGrid.h" +#include "CCDB/BasicCCDBManager.h" +#include namespace o2 { @@ -50,7 +52,7 @@ Generator::Generator() : FairGenerator("ALICEo2", "ALICEo2 Generator"), if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - if (initLoopersGen()) { + if (initTPCLoopersGen()) { mAddTPCLoopers = kTRUE; } } else { @@ -79,7 +81,7 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na if (transport) { bool tpcActive = (std::find(simConfig.getReadoutDetectors().begin(), simConfig.getReadoutDetectors().end(), "TPC") != simConfig.getReadoutDetectors().end()); if (tpcActive) { - if (initLoopersGen()) { + if (initTPCLoopersGen()) { mAddTPCLoopers = kTRUE; } } else { @@ -94,7 +96,7 @@ Generator::Generator(const Char_t* name, const Char_t* title) : FairGenerator(na /*****************************************************************/ #ifdef GENERATORS_WITH_TPCLOOPERS -bool Generator::initLoopersGen() +bool Generator::initTPCLoopersGen() { // Expand all environment paths const auto& loopersParam = o2::eventgen::GenTPCLoopersParam::Instance(); @@ -169,24 +171,24 @@ bool Generator::initLoopersGen() nclxrate = isAlien[2] || isCCDB[2] ? local_names[2] : nclxrate; try { // Create the TPC loopers generator with the provided parameters - mLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); + mTPCLoopersGen = std::make_unique(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton); const auto& intrate = loopersParam.intrate; // Configure the generator with flat gas loopers defined per orbit with clusters/track info // If intrate is negative (default), automatic IR from collisioncontext.root will be used if (flat_gas) { - mLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); - mLoopersGen->SetAdjust(loopersParam.adjust_flatgas); + mTPCLoopersGen->SetRate(nclxrate, (colsys == "PbPb") ? true : false, intrate); + mTPCLoopersGen->SetAdjust(loopersParam.adjust_flatgas); } else { // Otherwise, Poisson+Gauss sampling or fixed number of loopers per event will be used // Multiplier is applied only with distribution sampling // This configuration can be used for testing purposes, in all other cases flat gas is recommended - mLoopersGen->SetNLoopers(nLoopersPairs, nLoopersCompton); - mLoopersGen->SetMultiplier(multiplier); + mTPCLoopersGen->SetNLoopers(nLoopersPairs, nLoopersCompton); + mTPCLoopersGen->SetMultiplier(multiplier); } LOG(info) << "TPC Loopers generator initialized successfully"; } catch (const std::exception& e) { LOG(error) << "Failed to initialize TPC Loopers generator: " << e.what(); - mLoopersGen.reset(); + mTPCLoopersGen.reset(); } return kTRUE; } @@ -210,21 +212,21 @@ Bool_t { #ifdef GENERATORS_WITH_TPCLOOPERS if (mAddTPCLoopers) { - if (!mLoopersGen) { + if (!mTPCLoopersGen) { LOG(error) << "Loopers generator not initialized"; return kFALSE; } // Generate loopers using the initialized TPC loopers generator - if (!mLoopersGen->generateEvent()) { + if (!mTPCLoopersGen->generateEvent()) { LOG(error) << "Failed to generate loopers event"; return kFALSE; } - if (mLoopersGen->getNLoopers() == 0) { + if (mTPCLoopersGen->getNLoopers() == 0) { LOG(warning) << "No loopers generated for this event"; return kTRUE; } - const auto& looperParticles = mLoopersGen->importParticles(); + const auto& looperParticles = mTPCLoopersGen->importParticles(); if (looperParticles.empty()) { LOG(error) << "Failed to import loopers particles"; return kFALSE; diff --git a/Generators/src/TPCLoopers.cxx b/Generators/src/TPCLoopers.cxx index 8dff795de40a3..6e5af7c0c84d8 100644 --- a/Generators/src/TPCLoopers.cxx +++ b/Generators/src/TPCLoopers.cxx @@ -12,6 +12,16 @@ /// \author M+Giacalone - September 2025 #include "Generators/TPCLoopers.h" +#include "CCDB/CCDBTimeStampUtils.h" +#include "CCDB/CcdbApi.h" +#include "DetectorsRaw/HBFUtils.h" +#include "TF1.h" +#include +#include +#include "SimulationDataFormat/MCGenProperties.h" +#include +#include +#include "TDatabasePDG.h" // Static Ort::Env instance for multiple onnx model loading Ort::Env global_env(ORT_LOGGING_LEVEL_WARNING, "GlobalEnv"); @@ -46,10 +56,11 @@ std::vector Scaler::inverse_transform(const std::vector& input) { std::vector output; for (int i = 0; i < input.size(); ++i) { - if (i < input.size() - 2) + if (i < input.size() - 2) { output.push_back(input[i] * (normal_max[i] - normal_min[i]) + normal_min[i]); - else + } else { output.push_back(input[i] * outlier_scale[i - (input.size() - 2)] + outlier_center[i - (input.size() - 2)]); + } } return output; @@ -80,8 +91,9 @@ std::vector ONNXGenerator::generate_sample() // Generate a latent vector (z) std::vector z(100); - for (auto& v : z) + for (auto& v : z) { v = rand_gen.Gaus(0.0, 1.0); + } // Prepare input tensor std::vector input_shape = {1, 100}; @@ -227,7 +239,7 @@ Bool_t GenTPCLoopers::generateEvent() return true; } -Bool_t GenTPCLoopers::generateEvent(double& time_limit) +Bool_t GenTPCLoopers::generateEvent(double time_limit) { LOG(info) << "Time constraint for loopers: " << time_limit << " ns"; // Generate pairs @@ -253,6 +265,8 @@ Bool_t GenTPCLoopers::generateEvent(double& time_limit) std::vector GenTPCLoopers::importParticles() { std::vector particles; + const double mass_e = TDatabasePDG::Instance()->GetParticle(11)->Mass(); + const double mass_p = TDatabasePDG::Instance()->GetParticle(-11)->Mass(); // Get looper pairs from the event for (auto& pair : mGenPairs) { double px_e, py_e, pz_e, px_p, py_p, pz_p; @@ -268,8 +282,8 @@ std::vector GenTPCLoopers::importParticles() vy = pair[7]; vz = pair[8]; time = pair[9]; - e_etot = TMath::Sqrt(px_e * px_e + py_e * py_e + pz_e * pz_e + mMass_e * mMass_e); - p_etot = TMath::Sqrt(px_p * px_p + py_p * py_p + pz_p * pz_p + mMass_p * mMass_p); + e_etot = TMath::Sqrt(px_e * px_e + py_e * py_e + pz_e * pz_e + mass_e * mass_e); + p_etot = TMath::Sqrt(px_p * px_p + py_p * py_p + pz_p * pz_p + mass_p * mass_p); // Push the electron TParticle electron(11, 1, -1, -1, -1, -1, px_e, py_e, pz_e, e_etot, vx, vy, vz, time / 1e9); electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(electron.GetStatusCode(), 0).fullEncoding); @@ -295,7 +309,7 @@ std::vector GenTPCLoopers::importParticles() vy = compton[4]; vz = compton[5]; time = compton[6]; - etot = TMath::Sqrt(px * px + py * py + pz * pz + mMass_e * mMass_e); + etot = TMath::Sqrt(px * px + py * py + pz * pz + mass_e * mass_e); // Push the electron TParticle electron(11, 1, -1, -1, -1, -1, px, py, pz, etot, vx, vy, vz, time / 1e9); electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(electron.GetStatusCode(), 0).fullEncoding); @@ -329,7 +343,7 @@ unsigned int GenTPCLoopers::GaussianElectrons() return gaussValue; } -void GenTPCLoopers::SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_compton) +void GenTPCLoopers::SetNLoopers(unsigned int nsig_pair, unsigned int nsig_compton) { if (mFlatGas) { mNLoopersPairs = nsig_pair; @@ -348,7 +362,7 @@ void GenTPCLoopers::SetNLoopers(unsigned int& nsig_pair, unsigned int& nsig_comp } } -void GenTPCLoopers::SetMultiplier(std::array& mult) +void GenTPCLoopers::SetMultiplier(const std::array& mult) { // Multipliers will work only if the poissonian and gaussian parameters are set // otherwise they will be ignored @@ -362,7 +376,7 @@ void GenTPCLoopers::SetMultiplier(std::array& mult) } } -void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number = -1, const Int_t& nloopers_orbit = -1) +void GenTPCLoopers::setFlatGas(Bool_t flat, Int_t number, Int_t nloopers_orbit) { mFlatGas = flat; if (mFlatGas) { @@ -408,7 +422,7 @@ void GenTPCLoopers::setFlatGas(Bool_t& flat, const Int_t& number = -1, const Int LOG(info) << "Flat gas loopers: " << (mFlatGas ? "ON" : "OFF") << ", Reference loopers number per " << (mFlatGasOrbit ? "orbit " : "event ") << mFlatGasNumber; } -void GenTPCLoopers::setFractionPairs(float& fractionPairs) +void GenTPCLoopers::setFractionPairs(float fractionPairs) { if (fractionPairs < 0 || fractionPairs > 1) { LOG(fatal) << "Error: Loops fraction for pairs must be in the range [0, 1]."; @@ -418,7 +432,7 @@ void GenTPCLoopers::setFractionPairs(float& fractionPairs) LOG(info) << "Pairs fraction set to: " << mLoopsFractionPairs; } -void GenTPCLoopers::SetRate(const std::string& rateFile, const bool& isPbPb = true, const int& intRate = 50000) +void GenTPCLoopers::SetRate(const std::string& rateFile, bool isPbPb = true, int intRate) { // Checking if the rate file exists and is not empty TFile rate_file(rateFile.c_str(), "READ"); @@ -459,7 +473,7 @@ void GenTPCLoopers::SetRate(const std::string& rateFile, const bool& isPbPb = tr } } -void GenTPCLoopers::SetAdjust(const float& adjust = 0.f) +void GenTPCLoopers::SetAdjust(float adjust) { if (mFlatGas && mFlatGasOrbit && adjust >= -1.f && adjust != 0.f) { LOG(info) << "Adjusting flat gas number per orbit by " << adjust * 100.f << "%";