Skip to content

Commit 01a36ff

Browse files
sawenzelsawenzel
committed
Simplification + allow either ZN or ZP as fast sim
1 parent 748da9e commit 01a36ff

1 file changed

Lines changed: 54 additions & 38 deletions

File tree

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

Lines changed: 54 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -70,9 +70,11 @@ Detector::Detector(Bool_t active)
7070
#ifdef ZDC_FASTSIM_ONNX
7171
// If FastSim module was disabled, log appropriate message
7272
// otherwise check if all necessary parameters were passed, if so try build objects
73-
if (!o2::zdc::ZDCSimParam::Instance().useZDCFastSim) {
73+
auto& simparam = o2::zdc::ZDCSimParam::Instance();
74+
75+
if (!simparam.useZDCFastSim) {
7476
LOG(info) << "FastSim module disabled";
75-
} else if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierPath.empty() && !o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierScales.empty()) {
77+
} else if (simparam.useZDCFastSim && !simparam.ZDCFastSimClassifierPath.empty() && !simparam.ZDCFastSimClassifierScales.empty()) {
7678
if (!mClassifierScaler) {
7779
mClassifierScaler = new fastsim::processors::StandardScaler;
7880
}
@@ -82,38 +84,38 @@ Detector::Detector(Bool_t active)
8284
if (!mModelScalerProton) {
8385
mModelScalerProton = new fastsim::processors::StandardScaler;
8486
}
85-
auto eonScales = o2::zdc::fastsim::loadScales(o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierScales);
87+
auto eonScales = o2::zdc::fastsim::loadScales(simparam.ZDCFastSimClassifierScales);
8688
if (!eonScales.has_value()) {
8789
LOG(error) << "Error while reading model scales from: "
88-
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierScales << "'";
90+
<< "'" << simparam.ZDCFastSimClassifierScales << "'";
8991
LOG(error) << "FastSim module disabled.";
9092
} else {
9193
mClassifierScaler->setScales(eonScales->first, eonScales->second);
92-
mFastSimClassifier = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimClassifierPath, 1);
94+
mFastSimClassifier = new o2::zdc::fastsim::ConditionalModelSimulation(simparam.ZDCFastSimClassifierPath, 1);
9395

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+
if (simparam.useZDCFastSim && !simparam.ZDCFastSimModelPathNeutron.empty() && !simparam.ZDCFastSimModelScalesNeutron.empty()) {
97+
auto modelScalesNeutron = o2::zdc::fastsim::loadScales(simparam.ZDCFastSimModelScalesNeutron);
9698

9799
if (!modelScalesNeutron.has_value()) {
98100
LOG(error) << "Error while reading model scales from: "
99-
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesNeutron << "'";
101+
<< "'" << simparam.ZDCFastSimModelScalesNeutron << "'";
100102
LOG(error) << "FastSim module disabled";
101103
} else {
102104
mModelScalerNeutron->setScales(modelScalesNeutron->first, modelScalesNeutron->second);
103-
mFastSimModelNeutron = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathNeutron, 1);
105+
mFastSimModelNeutron = new o2::zdc::fastsim::ConditionalModelSimulation(simparam.ZDCFastSimModelPathNeutron, 1);
104106
LOG(info) << "FastSim neutron module enabled";
105107
}
106108
}
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);
109+
if (simparam.useZDCFastSim && !simparam.ZDCFastSimModelPathProton.empty() && !simparam.ZDCFastSimModelScalesProton.empty()) {
110+
auto modelScalesProton = o2::zdc::fastsim::loadScales(simparam.ZDCFastSimModelScalesProton);
109111

110112
if (!modelScalesProton.has_value()) {
111113
LOG(error) << "Error while reading model scales from: "
112-
<< "'" << o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelScalesProton << "'";
114+
<< "'" << simparam.ZDCFastSimModelScalesProton << "'";
113115
LOG(error) << "FastSim module disabled";
114116
} else {
115117
mModelScalerProton->setScales(modelScalesProton->first, modelScalesProton->second);
116-
mFastSimModelProton = new o2::zdc::fastsim::ConditionalModelSimulation(o2::zdc::ZDCSimParam::Instance().ZDCFastSimModelPathProton, 1);
118+
mFastSimModelProton = new o2::zdc::fastsim::ConditionalModelSimulation(simparam.ZDCFastSimModelPathProton, 1);
117119
LOG(info) << "FastSim proton module enabled";
118120
}
119121
}
@@ -2488,7 +2490,8 @@ void Detector::FinishPrimary()
24882490

24892491
#ifdef ZDC_FASTSIM_ONNX
24902492
// dump to file only if debugZDCFastSim is set to true
2491-
if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim && o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModelNeutron != nullptr && mFastSimModelProton != nullptr && mFastSimClassifier != nullptr) {
2493+
auto& simparam = o2::zdc::ZDCSimParam::Instance();
2494+
if (simparam.debugZDCFastSim && simparam.useZDCFastSim && mFastSimModelNeutron != nullptr && mFastSimModelProton != nullptr && mFastSimClassifier != nullptr) {
24922495
std::fstream output("o2sim-FastSimResult", std::fstream::out | std::fstream::app);
24932496
if (!output.is_open()) {
24942497
LOG(error) << "Could not open file.";
@@ -2516,7 +2519,9 @@ void Detector::BeginPrimary()
25162519
mCurrentPrincipalParticle = *stack->GetCurrentTrack();
25172520

25182521
#ifdef ZDC_FASTSIM_ONNX
2519-
if (o2::zdc::ZDCSimParam::Instance().useZDCFastSim && mFastSimModelNeutron != nullptr && mFastSimModelProton != nullptr && mFastSimClassifier != nullptr) {
2522+
auto& simparam = o2::zdc::ZDCSimParam::Instance();
2523+
using std::vector;
2524+
if (simparam.useZDCFastSim && (mFastSimModelNeutron != nullptr || mFastSimModelProton != nullptr) && mFastSimClassifier != nullptr) {
25202525
const std::vector<float> rawInput = {static_cast<float>(mCurrentPrincipalParticle.Energy()),
25212526
static_cast<float>(mCurrentPrincipalParticle.Vx()),
25222527
static_cast<float>(mCurrentPrincipalParticle.Vy()),
@@ -2531,34 +2536,45 @@ void Detector::BeginPrimary()
25312536
if (!scaledClassParticle.has_value()) {
25322537
LOG(error) << "FastSimModule: error occurred on scaling";
25332538
} else {
2534-
std::vector<std::vector<float>> classifierInput = {std::move(*scaledClassParticle)};
2539+
vector<vector<float>> classifierInput = {std::move(*scaledClassParticle)};
25352540
mFastSimClassifier->setInput(classifierInput);
25362541
mFastSimClassifier->run();
2542+
2543+
// this classifies if particle will leave a trace at all in one of the calos ---> TODO: better do it separately for ZN + ZP?
25372544
if (fastsim::processors::readClassifier(mFastSimClassifier->getResult()[0], 1)[0]) {
2538-
auto scaledModelParticleNeutron = mModelScalerNeutron->scale(rawInput);
2539-
auto scaledModelParticleProton = mModelScalerProton->scale(rawInput);
2540-
if (!scaledModelParticleNeutron.has_value() || !scaledModelParticleProton.has_value()) {
2541-
LOG(error) << "FastSimModule: error occurred on scaling";
2542-
} else {
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();
2549-
2550-
if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim) {
2551-
mFastSimResults.push_back(fastsim::processors::calculateChannels(mFastSimModelNeutron->getResult()[0], 1)[0]);
2552-
// TODO: same for protons
2545+
// let's do the neutron (ZN) part
2546+
if (mModelScalerNeutron && mFastSimModelNeutron) {
2547+
LOG(info) << "Generating fast hits for ZN";
2548+
auto scaledModelParticleNeutron = mModelScalerNeutron->scale(rawInput);
2549+
if (!scaledModelParticleNeutron.has_value()) {
2550+
LOG(error) << "FastSimModule: error occurred on scaling";
2551+
} else {
2552+
vector<vector<float>> modelInputNeutron = {fastsim::normal_distribution(0.0, 1.0, 10), std::move(*scaledModelParticleNeutron)};
2553+
mFastSimModelNeutron->setInput(modelInputNeutron);
2554+
mFastSimModelNeutron->run();
2555+
if (simparam.debugZDCFastSim) {
2556+
mFastSimResults.push_back(fastsim::processors::calculateChannels(mFastSimModelNeutron->getResult()[0], 1)[0]);
2557+
}
2558+
// produce hits from fast sim result
2559+
bool forward = mCurrentPrincipalParticle.Pz() > 0.;
2560+
FastSimToHits(mFastSimModelNeutron->getResult()[0], mCurrentPrincipalParticle, forward ? ZNA : ZNC);
25532561
}
2554-
2555-
// produce hits from fast sim result
2556-
bool forward = mCurrentPrincipalParticle.Pz() > 0.;
2557-
FastSimToHits(mFastSimModelNeutron->getResult()[0], mCurrentPrincipalParticle, forward ? ZNA : ZNC);
2558-
FastSimToHits(mFastSimModelProton->getResult()[0], mCurrentPrincipalParticle, forward ? ZPA : ZPC);
25592562
}
2560-
} else if (o2::zdc::ZDCSimParam::Instance().debugZDCFastSim) {
2561-
mFastSimResults.push_back({0, 0, 0, 0, 0});
2563+
// let's do the proton (ZP) part
2564+
if (mModelScalerProton && mFastSimModelProton) {
2565+
LOG(info) << "Generating fast hits for ZP";
2566+
auto scaledModelParticleProton = mModelScalerProton->scale(rawInput);
2567+
if (!scaledModelParticleProton.has_value()) {
2568+
LOG(error) << "FastSimModule: error occurred on scaling";
2569+
} else {
2570+
vector<vector<float>> modelInputProton = {fastsim::normal_distribution(0.0, 1.0, 10), std::move(*scaledModelParticleProton)};
2571+
mFastSimModelProton->setInput(modelInputProton);
2572+
mFastSimModelProton->run();
2573+
// produce hits from fast sim result
2574+
bool forward = mCurrentPrincipalParticle.Pz() > 0.;
2575+
FastSimToHits(mFastSimModelProton->getResult()[0], mCurrentPrincipalParticle, forward ? ZPA : ZPC);
2576+
}
2577+
} // end proton treatment
25622578
}
25632579
}
25642580
}

0 commit comments

Comments
 (0)