Skip to content

Commit f96e7c2

Browse files
committed
Implementation of TPC loopers in O2
1 parent 1b04a9a commit f96e7c2

File tree

14 files changed

+1095
-4
lines changed

14 files changed

+1095
-4
lines changed

Generators/CMakeLists.txt

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ o2_add_library(Generators
4141
src/GeneratorTParticleParam.cxx
4242
src/GeneratorService.cxx
4343
src/FlowMapper.cxx
44+
src/TPCLoopers.cxx
45+
src/TPCLoopersParam.cxx
4446
$<$<BOOL:${pythia_FOUND}>:src/GeneratorPythia8.cxx>
4547
$<$<BOOL:${pythia_FOUND}>:src/DecayerPythia8.cxx>
4648
$<$<BOOL:${pythia_FOUND}>:src/GeneratorPythia8Param.cxx>
@@ -53,6 +55,7 @@ o2_add_library(Generators
5355
PUBLIC_LINK_LIBRARIES FairRoot::Base O2::SimConfig O2::CommonUtils O2::DetectorsBase O2::ZDCBase
5456
O2::SimulationDataFormat ${pythiaTarget} ${hepmcTarget}
5557
FairRoot::Gen
58+
onnxruntime::onnxruntime
5659
TARGETVARNAME targetName)
5760

5861
if(pythia_FOUND)
@@ -63,6 +66,8 @@ if(HepMC3_FOUND)
6366
target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_HEPMC3)
6467
endif()
6568

69+
target_compile_definitions(${targetName} PUBLIC GENERATORS_WITH_TPCLOOPERS)
70+
6671
set(headers
6772
include/Generators/Generator.h
6873
include/Generators/Trigger.h
@@ -88,6 +93,10 @@ set(headers
8893
include/Generators/FlowMapper.h
8994
)
9095

96+
list(APPEND headers
97+
include/Generators/TPCLoopers.h
98+
include/Generators/TPCLoopersParam.h)
99+
91100
if(pythia_FOUND)
92101
list(APPEND headers
93102
include/Generators/GeneratorPythia8.h
@@ -158,4 +167,5 @@ endif()
158167

159168
o2_data_file(COPY share/external DESTINATION Generators)
160169
o2_data_file(COPY share/egconfig DESTINATION Generators)
170+
o2_data_file(COPY share/TPCLoopers DESTINATION Generators)
161171
o2_data_file(COPY share/pythia8 DESTINATION Generators)

Generators/include/Generators/Generator.h

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717
#include "FairGenerator.h"
1818
#include "TParticle.h"
1919
#include "Generators/Trigger.h"
20+
#ifdef GENERATORS_WITH_TPCLOOPERS
21+
#include "Generators/TPCLoopers.h"
22+
#include "Generators/TPCLoopersParam.h"
23+
#endif
2024
#include <functional>
2125
#include <vector>
2226
#include <unordered_map>
@@ -73,6 +77,7 @@ class Generator : public FairGenerator
7377
/** methods to override **/
7478
virtual Bool_t generateEvent() = 0; // generates event (in structure internal to generator)
7579
virtual Bool_t importParticles() = 0; // fills the mParticles vector (transfer from generator state)
80+
Bool_t finalizeEvent(); // final part of event generation that can be customised using external macros
7681
virtual void updateHeader(o2::dataformats::MCEventHeader* eventHeader) {};
7782
Bool_t triggerEvent();
7883

@@ -154,6 +159,8 @@ class Generator : public FairGenerator
154159
private:
155160
void updateSubGeneratorInformation(o2::dataformats::MCEventHeader* header) const;
156161

162+
// loopers flag
163+
Bool_t mAddTPCLoopers = kFALSE; // Flag is automatically set to true if TPC is in readout detectors, loopers are not vetoed and transport is enabled
157164
// collect an ID and a short description of sub-generator entities
158165
std::unordered_map<int, std::string> mSubGeneratorsIdToDesc;
159166
// the current ID of the sub-generator used in the current event (if applicable)
@@ -162,6 +169,12 @@ class Generator : public FairGenerator
162169
// global static information about (upper limit of) number of events to be generated
163170
static unsigned int gTotalNEvents;
164171

172+
#ifdef GENERATORS_WITH_TPCLOOPERS
173+
// Loopers generator instance
174+
std::unique_ptr<o2::eventgen::GenTPCLoopers> mTPCLoopersGen = nullptr;
175+
bool initTPCLoopersGen();
176+
#endif
177+
165178
ClassDefOverride(Generator, 2);
166179

167180
}; /** class Generator **/
Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
// Copyright 2024-2025 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
/// \author M+Giacalone - September 2025
13+
14+
#ifndef ALICEO2_EVENTGEN_TPCLOOPERS_H_
15+
#define ALICEO2_EVENTGEN_TPCLOOPERS_H_
16+
17+
#ifdef GENERATORS_WITH_TPCLOOPERS
18+
#include <onnxruntime_cxx_api.h>
19+
#endif
20+
#include <vector>
21+
#include <rapidjson/document.h>
22+
#include "TRandom3.h"
23+
#include <SimulationDataFormat/DigitizationContext.h>
24+
#include "TParticle.h"
25+
26+
#ifdef GENERATORS_WITH_TPCLOOPERS
27+
// Static Ort::Env instance for multiple onnx model loading
28+
extern Ort::Env global_env;
29+
30+
// This class is responsible for loading the scaler parameters from a JSON file
31+
// and applying the inverse transformation to the generated data.
32+
// Inferenced output is scaled (min-max normalization or robust scaling for outlier features) during training,
33+
// so we need to revert this transformation to get physical values.
34+
struct Scaler {
35+
std::vector<double> normal_min;
36+
std::vector<double> normal_max;
37+
std::vector<double> outlier_center;
38+
std::vector<double> outlier_scale;
39+
40+
void load(const std::string& filename);
41+
42+
std::vector<double> inverse_transform(const std::vector<double>& input);
43+
44+
private:
45+
std::vector<double> jsonArrayToVector(const rapidjson::Value& jsonArray);
46+
};
47+
48+
// This class loads the ONNX model and generates samples using it.
49+
class ONNXGenerator
50+
{
51+
public:
52+
ONNXGenerator(Ort::Env& shared_env, const std::string& model_path);
53+
54+
std::vector<double> generate_sample();
55+
56+
private:
57+
Ort::Env& env;
58+
Ort::Session session;
59+
TRandom3 rand_gen;
60+
};
61+
#endif // GENERATORS_WITH_TPCLOOPERS
62+
63+
namespace o2
64+
{
65+
namespace eventgen
66+
{
67+
68+
#ifdef GENERATORS_WITH_TPCLOOPERS
69+
/**
70+
* Generator for TPC Loopers based on pre-trained ONNX models.
71+
* Currently it generates loopers as electron-positron pairs and Compton electrons
72+
* according to specified distributions and parameters.
73+
* This can be extended to other types of background processes in the future (e.g. slow neutron spallation products, saturation tail).
74+
* Multiple configuration options are available:
75+
* - 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.
76+
* ==> 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.
77+
* ==> Number of loopers per orbit can be adjusted via a specific parameter.
78+
* - Poisson + Gaussian sampling: number of loopers are sampled from Poissonian (for pairs) and Gaussian (for Compton electrons) distributions based on provided parameters.
79+
* ==> flat gas must be disabled to use this option.
80+
* - Fixed number of loopers per event
81+
* ==> flat gas must be disabled to use this option and Poissonian/Gaussian parameters file should be set to None
82+
*/
83+
class GenTPCLoopers
84+
{
85+
public:
86+
GenTPCLoopers(std::string model_pairs = "tpcloopmodel.onnx", std::string model_compton = "tpcloopmodelcompton.onnx",
87+
std::string poisson = "poisson.csv", std::string gauss = "gauss.csv", std::string scaler_pair = "scaler_pair.json",
88+
std::string scaler_compton = "scaler_compton.json");
89+
90+
Bool_t generateEvent();
91+
92+
Bool_t generateEvent(double time_limit);
93+
94+
std::vector<TParticle> importParticles();
95+
96+
unsigned int PoissonPairs();
97+
98+
unsigned int GaussianElectrons();
99+
100+
void SetNLoopers(unsigned int nsig_pair, unsigned int nsig_compton);
101+
102+
void SetMultiplier(const std::array<float, 2>& mult);
103+
104+
void setFlatGas(Bool_t flat, Int_t number = -1, Int_t nloopers_orbit = -1);
105+
106+
void setFractionPairs(float fractionPairs);
107+
108+
void SetRate(const std::string& rateFile, bool isPbPb, int intRate = 50000);
109+
110+
void SetAdjust(float adjust = 0.f);
111+
112+
unsigned int getNLoopers() const { return (mNLoopersPairs + mNLoopersCompton); }
113+
114+
private:
115+
std::unique_ptr<ONNXGenerator> mONNX_pair = nullptr;
116+
std::unique_ptr<ONNXGenerator> mONNX_compton = nullptr;
117+
std::unique_ptr<Scaler> mScaler_pair = nullptr;
118+
std::unique_ptr<Scaler> mScaler_compton = nullptr;
119+
double mPoisson[3] = {0.0, 0.0, 0.0}; // Mu, Min and Max of Poissonian
120+
double mGauss[4] = {0.0, 0.0, 0.0, 0.0}; // Mean, Std, Min, Max
121+
std::vector<std::vector<double>> mGenPairs;
122+
std::vector<std::vector<double>> mGenElectrons;
123+
unsigned int mNLoopersPairs = -1;
124+
unsigned int mNLoopersCompton = -1;
125+
std::array<float, 2> mMultiplier = {1., 1.};
126+
bool mPoissonSet = false;
127+
bool mGaussSet = false;
128+
// Random number generator
129+
TRandom3 mRandGen;
130+
int mCurrentEvent = 0; // Current event number, used for adaptive loopers
131+
TFile* mContextFile = nullptr; // Input collision context file
132+
o2::steer::DigitizationContext* mCollisionContext = nullptr; // Pointer to the digitization context
133+
std::vector<o2::InteractionTimeRecord> mInteractionTimeRecords; // Interaction time records from collision context
134+
Bool_t mFlatGas = false; // Flag to indicate if flat gas loopers are used
135+
Bool_t mFlatGasOrbit = false; // Flag to indicate if flat gas loopers are per orbit
136+
Int_t mFlatGasNumber = -1; // Number of flat gas loopers per event
137+
double mIntTimeRecMean = 1.0; // Average interaction time record used for the reference
138+
double mTimeLimit = 0.0; // Time limit for the current event
139+
double mTimeEnd = 0.0; // Time limit for the last event
140+
float mLoopsFractionPairs = 0.08; // Fraction of loopers from Pairs
141+
int mInteractionRate = 50000; // Interaction rate in Hz
142+
};
143+
#endif // GENERATORS_WITH_TPCLOOPERS
144+
145+
} // namespace eventgen
146+
} // namespace o2
147+
148+
#endif // ALICEO2_EVENTGEN_TPCLOOPERS_H_
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Copyright 2024-2025 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
/// \author M+Giacalone - September 2025
13+
14+
#ifndef ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_
15+
#define ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_
16+
17+
#include "CommonUtils/ConfigurableParam.h"
18+
#include "CommonUtils/ConfigurableParamHelper.h"
19+
20+
namespace o2
21+
{
22+
namespace eventgen
23+
{
24+
25+
/**
26+
** a parameter class/struct to keep the settings of
27+
** the TPC loopers event-generator and
28+
** allow the user to modify them
29+
**/
30+
struct GenTPCLoopersParam : public o2::conf::ConfigurableParamHelper<GenTPCLoopersParam> {
31+
bool loopersVeto = false; // if true, no loopers are generated
32+
// Current files are set to custom user CCDB paths, TO BE CHANGED
33+
std::string model_pairs = "ccdb://Users/m/mgiacalo/WGAN_ExtGenPair"; // ONNX model for e+e- pair production
34+
std::string model_compton = "ccdb://Users/m/mgiacalo/WGAN_ExtGenCompton"; // ONNX model for Compton scattering
35+
std::string poisson = "${O2_ROOT}/share/Generators/TPCLoopers/poisson_params.csv"; // file with Poissonian parameters
36+
std::string gauss = "${O2_ROOT}/share/Generators/TPCLoopers/gaussian_params.csv"; // file with Gaussian parameters
37+
std::string scaler_pair = "${O2_ROOT}/share/Generators/TPCLoopers/ScalerPairParams.json"; // file with scaler parameters for e+e- pair production
38+
std::string scaler_compton = "${O2_ROOT}/share/Generators/TPCLoopers/ScalerComptonParams.json"; // file with scaler parameters for Compton scattering
39+
std::string nclxrate = "ccdb://Users/m/mgiacalo/ClustersTrackRatio"; // file with clusters/rate information per orbit
40+
std::string colsys = "PbPb"; // collision system (PbPb or pp)
41+
int intrate = -1; // Automatic IR from collision context if -1, else user-defined interaction rate in Hz
42+
bool flat_gas = true; // if true, the gas density is considered flat in the TPC volume
43+
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]
44+
float fraction_pairs = 0.08; // fraction of loopers [currently unused, kept for possible future debug developments]
45+
float multiplier[2] = {1., 1.}; // multiplier for pairs and compton loopers for Poissonian and Gaussian sampling
46+
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
47+
float adjust_flatgas = 0.f; // adjustment for the number of flat gas loopers per orbit (in percentage, e.g. -0.1 = -10%) [-1, inf)]
48+
O2ParamDef(GenTPCLoopersParam, "GenTPCLoopers");
49+
};
50+
51+
} // end namespace eventgen
52+
} // end namespace o2
53+
54+
#endif // ALICEO2_EVENTGEN_TPCLOOPERSPARAM_H_
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# TPC Loopers Generator - Parameter Files
2+
3+
This directory contains parameter files used by the TPC Loopers event generator in ALICE O2.
4+
5+
## Overview
6+
7+
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:
8+
- Example statistical distribution parameters for sampling the number of loopers per event
9+
- **Mandatory** scaling parameters for transforming the ONNX model outputs to physical values
10+
11+
## Files Description
12+
13+
### Statistical Sampling Parameters
14+
15+
The files provided in the folder are examples based on the training dataset.
16+
17+
#### `gaussian_params.csv`
18+
Parameters for Gaussian distribution used to sample the number of Compton electrons per event.
19+
20+
**Format:** Four values (one per line)
21+
1. Mean (μ)
22+
2. Standard deviation (σ)
23+
3. Minimum value
24+
4. Maximum value
25+
26+
#### `poisson_params.csv`
27+
Parameters for Poisson distribution used to sample the number of electron-positron pairs per event.
28+
29+
**Format:** Three values (one per line)
30+
1. Lambda (λ) parameter
31+
2. Minimum value
32+
3. Maximum value
33+
34+
### Scaler Parameters
35+
36+
These JSON files contain the parameters for inverse transformation of the ONNX models output. They should be kept as they are
37+
unless a new version of the models is released.
38+
39+
#### `ScalerComptonParams.json`
40+
Scaler parameters for Compton electron generation model.
41+
42+
**Structure:**
43+
```json
44+
{
45+
"normal": {
46+
"min": [array of 5 min values for min-max normalization],
47+
"max": [array of 5 max values for min-max normalization]
48+
},
49+
"outlier": {
50+
"center": [array of 2 center values for robust scaling],
51+
"scale": [array of 2 scale values for robust scaling]
52+
}
53+
}
54+
```
55+
56+
- **normal**: Min-max normalization parameters for standard features (`Px`, `Py`, `Pz`, `VertexCoordinatesX`, `VertexCoordinatesY`)
57+
- **outlier**: Robust scaler parameters (center and scale) for outlier features (`VertexCoordinatesZ`,`time`)
58+
59+
#### `ScalerPairParams.json`
60+
Scaler parameters for electron-positron pair generation model.
61+
62+
**Structure:**
63+
```json
64+
{
65+
"normal": {
66+
"min": [array of 8 min values for min-max normalization],
67+
"max": [array of 8 max values for min-max normalization]
68+
},
69+
"outlier": {
70+
"center": [array of 2 center values for robust scaling],
71+
"scale": [array of 2 scale values for robust scaling]
72+
}
73+
}
74+
```
75+
76+
- **normal**: Min-max normalization parameters for standard features (`Px_e`, `Py_e`, `Pz_e`,`Px_p`, `Py_p`, `Pz_p`, `VertexCoordinatesX`, `VertexCoordinatesY`)
77+
- **outlier**: Robust scaler parameters (center and scale) for outlier features (`VertexCoordinatesZ`,`time`)
78+
---
79+
*Author: M. Giacalone - September 2025*
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
{
2+
"normal": {
3+
"min": [
4+
-0.0108811147511005,
5+
-0.0098758740350604,
6+
-0.0103233363479375,
7+
-260.0542297363281,
8+
-259.80059814453125
9+
],
10+
"max": [
11+
0.0108060473576188,
12+
0.0103057539090514,
13+
0.0106524610891938,
14+
260.0343933105469,
15+
259.62890625
16+
]
17+
},
18+
"outlier": {
19+
"center": [
20+
-71.39387130737305,
21+
96791.23828125
22+
],
23+
"scale": [
24+
265.9389114379883,
25+
230762.30981445312
26+
]
27+
}
28+
}

0 commit comments

Comments
 (0)