Skip to content

Commit e839e82

Browse files
majerzemiliasawenzel
majerzemilia
authored and
sawenzel
committed
Generalize ZDC fast sim to include proton ML models
1 parent 81172a8 commit e839e82

3 files changed

Lines changed: 52 additions & 25 deletions

File tree

‎Detectors/ZDC/simulation/include/ZDCSimulation/Detector.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -241,11 +241,13 @@ class Detector : public o2::base::DetImpl<Detector>
241241
// fastsim model wrapper
242242
#ifdef ZDC_FASTSIM_ONNX
243243
fastsim::NeuralFastSimulation* mFastSimClassifier = nullptr; //! no ROOT serialization
244-
fastsim::NeuralFastSimulation* mFastSimModel = nullptr; //!
244+
fastsim::NeuralFastSimulation* mFastSimModelNeutron = nullptr; //!
245+
fastsim::NeuralFastSimulation* mFastSimModelProton = nullptr; //!
245246

246247
// Scalers for models inputs
247248
fastsim::processors::StandardScaler* mClassifierScaler = nullptr; //!
248-
fastsim::processors::StandardScaler* mModelScaler = nullptr; //!
249+
fastsim::processors::StandardScaler* mModelScalerNeutron = nullptr; //!
250+
fastsim::processors::StandardScaler* mModelScalerProton = nullptr; //!
249251

250252
// container for fastsim model responses
251253
using FastSimResults = std::vector<std::array<long, 5>>; //!

‎Detectors/ZDC/simulation/include/ZDCSimulation/ZDCSimParam.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,10 @@ struct ZDCSimParam : public o2::conf::ConfigurableParamHelper<ZDCSimParam> {
3131
bool debugZDCFastSim = false; ///< whether to dump fastsim output as 5ch to seperate file in csv format
3232
std::string ZDCFastSimClassifierPath = ""; ///< path to model file that classify if data are viable for model
3333
std::string ZDCFastSimClassifierScales = ""; ///< path to scales file for classifier
34-
std::string ZDCFastSimModelPath = ""; ///< path to model file
35-
std::string ZDCFastSimModelScales = ""; ///< path to scales file for model
34+
std::string ZDCFastSimModelPathNeutron = ""; ///< path to neutron model file
35+
std::string ZDCFastSimModelScalesNeutron = ""; ///< path to scales file for neutron model
36+
std::string ZDCFastSimModelPathProton = ""; ///< path to proton model file
37+
std::string ZDCFastSimModelScalesProton = ""; ///< path to scales file for proton model
3638

3739
O2ParamDef(ZDCSimParam, "ZDCSimParam");
3840
};

‎Detectors/ZDC/simulation/src/Detector.cxx‎

Lines changed: 44 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,11 @@ Detector::Detector(Bool_t active)
7676
if (!mClassifierScaler) {
7777
mClassifierScaler = new fastsim::processors::StandardScaler;
7878
}
79-
if (!mModelScaler) {
80-
mModelScaler = new fastsim::processors::StandardScaler;
79+
if (!mModelScalerNeutron) {
80+
mModelScalerNeutron = new fastsim::processors::StandardScaler;
81+
}
82+
if (!mModelScalerProton) {
83+
mModelScalerProton = new fastsim::processors::StandardScaler;
8184
}
8285
auto eonScales = o2::zdc::fastsim::loadScales(o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierScales);
8386
if (!eonScales.has_value()) {
@@ -88,17 +91,30 @@ Detector::Detector(Bool_t active)
8891
mClassifierScaler->setScales(eonScales->first, eonScales->second);
8992
mFastSimClassifier = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierPath, 1);
9093

91-
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPath.empty() && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScales.empty()) {
92-
auto modelScales = o2::zdc::fastsim::loadScales(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScales);
94+
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathNeutron.empty() && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesNeutron.empty()) {
95+
auto modelScalesNeutron = o2::zdc::fastsim::loadScales(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesNeutron);
96+
97+
if (!modelScalesNeutron.has_value()) {
98+
LOG(error) << "Error while reading model scales from: "
99+
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesNeutron << "'";
100+
LOG(error) << "FastSim module disabled";
101+
} else {
102+
mModelScalerNeutron->setScales(modelScalesNeutron->first, modelScalesNeutron->second);
103+
mFastSimModelNeutron = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathNeutron, 1);
104+
LOG(info) << "FastSim neutron module enabled";
105+
}
106+
}
107+
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathProton.empty() && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesProton.empty()) {
108+
auto modelScalesProton = o2::zdc::fastsim::loadScales(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesProton);
93109

94-
if (!modelScales.has_value()) {
110+
if (!modelScalesProton.has_value()) {
95111
LOG(error) << "Error while reading model scales from: "
96-
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScales << "'";
112+
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesProton << "'";
97113
LOG(error) << "FastSim module disabled";
98114
} else {
99-
mModelScaler->setScales(modelScales->first, modelScales->second);
100-
mFastSimModel = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPath, 1);
101-
LOG(info) << "FastSim module enabled";
115+
mModelScalerProton->setScales(modelScalesProton->first, modelScalesProton->second);
116+
mFastSimModelProton = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathProton, 1);
117+
LOG(info) << "FastSim proton module enabled";
102118
}
103119
}
104120
}
@@ -118,9 +134,11 @@ Detector::Detector(const Detector& rhs)
118134
Detector::~Detector()
119135
{
120136
delete (mFastSimClassifier);
121-
delete (mFastSimModel);
137+
delete (mFastSimModelNeutron);
138+
delete (mFastSimModelProton);
122139
delete (mClassifierScaler);
123-
delete (mModelScaler);
140+
delete (mModelScalerNeutron);
141+
delete (mModelScalerProton);
124142
}
125143
#endif
126144

@@ -2470,7 +2488,7 @@ void Detector::FinishPrimary()
24702488

24712489
#ifdef ZDC_FASTSIM_ONNX
24722490
// dump to file only if debugZDCFastSim is set to true
2473-
if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim && o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModel != nullptr && mFastSimClassifier != nullptr) {
2491+
if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim && o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModelNeutron != nullptr && mFastSimModelProton != nullptr && mFastSimClassifier != nullptr) {
24742492
std::fstream output("o2sim-FastSimResult", std::fstream::out | std::fstream::app);
24752493
if (!output.is_open()) {
24762494
LOG(error) << "Could not open file.";
@@ -2498,7 +2516,7 @@ void Detector::BeginPrimary()
24982516
mCurrentPrincipalParticle = *stack->GetCurrentTrack();
24992517

25002518
#ifdef ZDC_FASTSIM_ONNX
2501-
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModel != nullptr && mFastSimClassifier != nullptr) {
2519+
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModelNeutron != nullptr && mFastSimModelProton != nullptr && mFastSimClassifier != nullptr) {
25022520
const std::vector<float> rawInput = {static_cast<float>(mCurrentPrincipalParticle.Energy()),
25032521
static_cast<float>(mCurrentPrincipalParticle.Vx()),
25042522
static_cast<float>(mCurrentPrincipalParticle.Vy()),
@@ -2517,22 +2535,27 @@ void Detector::BeginPrimary()
25172535
mFastSimClassifier->setInput(classifierInput);
25182536
mFastSimClassifier->run();
25192537
if (fastsim::processors::readClassifier(mFastSimClassifier->getResult()[0], 1)[0]) {
2520-
auto scaledModelParticle = mModelScaler->scale(rawInput);
2521-
if (!scaledModelParticle.has_value()) {
2538+
auto scaledModelParticleNeutron = mModelScalerNeutron->scale(rawInput);
2539+
auto scaledModelParticleProton = mModelScalerProton->scale(rawInput);
2540+
if (!scaledModelParticleNeutron.has_value() || !scaledModelParticleProton.has_value()) {
25222541
LOG(error) << "FastSimModule: error occurred on scaling";
25232542
} else {
2524-
std::vector<std::vector<float>> modelInput = {fastsim::normal_distribution(0.0, 1.0, 10), std::move(*scaledModelParticle)};
2525-
mFastSimModel->setInput(modelInput);
2526-
mFastSimModel->run();
2543+
std::vector<std::vector<float>> modelInputNeutron = {fastsim::normal_distribution(0.0, 1.0, 10), std::move(*scaledModelParticleNeutron)};
2544+
mFastSimModelNeutron->setInput(modelInputNeutron);
2545+
mFastSimModelNeutron->run();
2546+
std::vector<std::vector<float>> modelInputProton = {fastsim::normal_distribution(0.0, 1.0, 10), std::move(*scaledModelParticleProton)};
2547+
mFastSimModelProton->setInput(modelInputProton);
2548+
mFastSimModelProton->run();
25272549

25282550
if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim) {
2529-
mFastSimResults.push_back(fastsim::processors::calculateChannels(mFastSimModel->getResult()[0], 1)[0]);
2551+
mFastSimResults.push_back(fastsim::processors::calculateChannels(mFastSimModelNeutron->getResult()[0], 1)[0]);
2552+
// TODO: same for protons
25302553
}
25312554

25322555
// produce hits from fast sim result
25332556
bool forward = mCurrentPrincipalParticle.Pz() > 0.;
2534-
FastSimToHits(mFastSimModel->getResult()[0], mCurrentPrincipalParticle, forward ? ZNA : ZNC);
2535-
// TODO: call models for all detectors ZNA + ZPA
2557+
FastSimToHits(mFastSimModelNeutron->getResult()[0], mCurrentPrincipalParticle, forward ? ZNA : ZNC);
2558+
FastSimToHits(mFastSimModelProton->getResult()[0], mCurrentPrincipalParticle, forward ? ZPA : ZPC);
25362559
}
25372560
} else if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim) {
25382561
mFastSimResults.push_back({0, 0, 0, 0, 0});

0 commit comments

Comments
 (0)