Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 8 additions & 11 deletions Generators/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,8 @@ o2_add_library(Generators
src/GeneratorTParticleParam.cxx
src/GeneratorService.cxx
src/FlowMapper.cxx
$<$<BOOL:${onnxruntime_FOUND}>:src/TPCLoopers.cxx>
$<$<BOOL:${onnxruntime_FOUND}>:src/TPCLoopersParam.cxx>
src/TPCLoopers.cxx
src/TPCLoopersParam.cxx
$<$<BOOL:${pythia_FOUND}>:src/GeneratorPythia8.cxx>
$<$<BOOL:${pythia_FOUND}>:src/DecayerPythia8.cxx>
$<$<BOOL:${pythia_FOUND}>:src/GeneratorPythia8Param.cxx>
Expand All @@ -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
$<$<BOOL:${onnxruntime_FOUND}>:onnxruntime::onnxruntime>
onnxruntime::onnxruntime
TARGETVARNAME targetName)

if(pythia_FOUND)
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
5 changes: 2 additions & 3 deletions Generators/include/Generators/Generator.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -172,8 +171,8 @@ class Generator : public FairGenerator

#ifdef GENERATORS_WITH_TPCLOOPERS
// Loopers generator instance
std::unique_ptr<o2::eventgen::GenTPCLoopers> mLoopersGen = nullptr;
bool initLoopersGen();
std::unique_ptr<o2::eventgen::GenTPCLoopers> mTPCLoopersGen = nullptr;
bool initTPCLoopersGen();
#endif

ClassDefOverride(Generator, 2);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,28 +17,20 @@
#ifdef GENERATORS_WITH_TPCLOOPERS
#include <onnxruntime_cxx_api.h>
#endif
#include <iostream>
#include <vector>
#include <fstream>
#include <rapidjson/document.h>
#include "CCDB/CCDBTimeStampUtils.h"
#include "CCDB/CcdbApi.h"
#include "DetectorsRaw/HBFUtils.h"
#include "TRandom3.h"
#include "TDatabasePDG.h"
#include <SimulationDataFormat/DigitizationContext.h>
#include <SimulationDataFormat/ParticleStatus.h>
#include "SimulationDataFormat/MCGenProperties.h"
#include "TParticle.h"
#include "TF1.h"
#include <filesystem>

#ifdef GENERATORS_WITH_TPCLOOPERS
// Static Ort::Env instance for multiple onnx model loading
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<double> normal_min;
std::vector<double> normal_max;
Expand Down Expand Up @@ -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:
Expand All @@ -83,25 +89,25 @@ class GenTPCLoopers

Bool_t generateEvent();

Bool_t generateEvent(double& time_limit);
Bool_t generateEvent(double time_limit);

std::vector<TParticle> importParticles();

unsigned int PoissonPairs();

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<float, 2>& mult);
void SetMultiplier(const std::array<float, 2>& 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); }

Expand All @@ -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
Expand Down
35 changes: 18 additions & 17 deletions Generators/include/Generators/TPCLoopersParam.h
Original file line number Diff line number Diff line change
Expand Up @@ -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<GenTPCLoopersParam> {
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");
};

Expand Down
79 changes: 79 additions & 0 deletions Generators/share/TPCLoopers/README.md
Original file line number Diff line number Diff line change
@@ -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*
28 changes: 15 additions & 13 deletions Generators/src/Generator.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
#include "TParticle.h"
#include "TSystem.h"
#include "TGrid.h"
#include "CCDB/BasicCCDBManager.h"
#include <filesystem>

namespace o2
{
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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();
Expand Down Expand Up @@ -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<o2::eventgen::GenTPCLoopers>(model_pairs, model_compton, poisson, gauss, scaler_pair, scaler_compton);
mTPCLoopersGen = std::make_unique<o2::eventgen::GenTPCLoopers>(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;
}
Expand All @@ -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;
Expand Down
Loading