Skip to content

Commit 4de6bde

Browse files
GNiendorfGNiendorf
committed
Add T3 DNN for LST
1 parent 357fdc2 commit 4de6bde

File tree

10 files changed

+2058
-114
lines changed

10 files changed

+2058
-114
lines changed

‎RecoTracker/LSTCore/interface/alpaka/Common.h‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -49,18 +49,34 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
4949
{0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.18f, 0.18f, /*10*/ 0.18f, 0.18f, 0.18f, 0.18f, 0.18f},
5050
{0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.18f, /*10*/ 0.18f, 0.18f, 0.18f, 0.18f, 0.18f}};
5151

52-
namespace t5dnn {
53-
HOST_DEVICE_CONSTANT float kZ_max = 267.2349854f;
54-
HOST_DEVICE_CONSTANT float kR_max = 110.1099396f;
52+
namespace dnn {
53+
54+
// Common constants for both DNNs
5555
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
5656
HOST_DEVICE_CONSTANT float kPhi_norm = kPi;
57-
// pt, eta binned
5857
constexpr unsigned int kPtBins = 2;
5958
constexpr unsigned int kEtaBins = 10;
60-
HOST_DEVICE_CONSTANT float kWp[kPtBins][kEtaBins] = {
61-
{0.4493, 0.4939, 0.5715, 0.6488, 0.5709, 0.5938, 0.7164, 0.7565, 0.8103, 0.8593},
62-
{0.4488, 0.4448, 0.5067, 0.5929, 0.4836, 0.4112, 0.4968, 0.4403, 0.5597, 0.5067}};
63-
} // namespace t5dnn
59+
60+
namespace t3dnn {
61+
HOST_DEVICE_CONSTANT float kZ_max = 224.149505f;
62+
HOST_DEVICE_CONSTANT float kR_max = 98.932365f;
63+
HOST_DEVICE_CONSTANT float kWp_prompt[kPtBins][kEtaBins] = {
64+
{0.4957, 0.5052, 0.5201, 0.5340, 0.4275, 0.4708, 0.4890, 0.4932, 0.5400, 0.5449},
65+
{0.0302, 0.0415, 0.0994, 0.1791, 0.1960, 0.2467, 0.3227, 0.3242, 0.2367, 0.2187}};
66+
HOST_DEVICE_CONSTANT float kWp_displaced[kPtBins][kEtaBins] = {
67+
{0.0334, 0.0504, 0.0748, 0.0994, 0.1128, 0.1123, 0.1118, 0.1525, 0.1867, 0.1847},
68+
{0.0091, 0.0075, 0.0350, 0.0213, 0.0435, 0.0676, 0.1957, 0.1649, 0.1080, 0.1046}};
69+
} // namespace t3dnn
70+
71+
namespace t5dnn {
72+
HOST_DEVICE_CONSTANT float kZ_max = 267.2349854f;
73+
HOST_DEVICE_CONSTANT float kR_max = 110.1099396f;
74+
HOST_DEVICE_CONSTANT float kWp[kPtBins][kEtaBins] = {
75+
{0.4493, 0.4939, 0.5715, 0.6488, 0.5709, 0.5938, 0.7164, 0.7565, 0.8103, 0.8593},
76+
{0.4488, 0.4448, 0.5067, 0.5929, 0.4836, 0.4112, 0.4968, 0.4403, 0.5597, 0.5067}};
77+
} // namespace t5dnn
78+
79+
} // namespace dnn
6480

6581
} // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst
6682
#endif

‎RecoTracker/LSTCore/src/alpaka/NeuralNetwork.h‎

Lines changed: 203 additions & 99 deletions
Large diffs are not rendered by default.

‎RecoTracker/LSTCore/src/alpaka/T3NeuralNetworkWeights.h‎

Lines changed: 235 additions & 0 deletions
Large diffs are not rendered by default.

RecoTracker/LSTCore/src/alpaka/NeuralNetworkWeights.h renamed to RecoTracker/LSTCore/src/alpaka/T5NeuralNetworkWeights.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
#ifndef RecoTracker_LSTCore_src_alpaka_NeuralNetworkWeights_h
2-
#define RecoTracker_LSTCore_src_alpaka_NeuralNetworkWeights_h
1+
#ifndef RecoTracker_LSTCore_src_alpaka_T5NeuralNetworkWeights_h
2+
#define RecoTracker_LSTCore_src_alpaka_T5NeuralNetworkWeights_h
33

44
#include <alpaka/alpaka.hpp>
55

66
#include "FWCore/Utilities/interface/HostDeviceConstant.h"
77

8-
namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::t5dnn {
8+
namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::dnn::t5dnn {
99
HOST_DEVICE_CONSTANT float bias_layer1[32] = {
1010
-1.3837075f, -0.0653152f, -0.7900129f, 0.0714758f, -1.1574365f, -1.4634879f, -0.9317133f, -0.1455518f,
1111
-0.0459635f, -0.2055620f, 0.0586231f, -0.8943899f, -0.1009487f, 0.0166031f, -0.5451909f, -0.1384538f,
@@ -254,6 +254,6 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::t5dnn {
254254
{-0.5268564f}, {0.0735416f}, {0.0270067f}, {-0.5614370f},
255255
};
256256

257-
} // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::t5dnn
257+
} // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::dnn::t5dnn
258258

259259
#endif

‎RecoTracker/LSTCore/src/alpaka/Triplet.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
#include "RecoTracker/LSTCore/interface/TripletsSoA.h"
1111
#include "RecoTracker/LSTCore/interface/Circle.h"
1212

13+
#include "NeuralNetwork.h"
14+
1315
namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
1416

1517
ALPAKA_FN_ACC ALPAKA_FN_INLINE void addTripletToMemory(ModulesConst modules,
@@ -698,6 +700,11 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
698700
ptCut))
699701
return false;
700702

703+
bool inference =
704+
lst::t3dnn::runInference(acc, mds, firstMDIndex, secondMDIndex, thirdMDIndex, circleRadius, betaIn);
705+
if (!inference) // T3-building cut
706+
return false;
707+
701708
return true;
702709
}
703710

‎RecoTracker/LSTCore/standalone/analysis/DNN/train_T3_DNN.ipynb‎

Lines changed: 1430 additions & 0 deletions
Large diffs are not rendered by default.

‎RecoTracker/LSTCore/standalone/code/core/AccessHelper.cc‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -207,12 +207,24 @@ std::vector<unsigned int> getModuleIdxsFromT5(LSTEvent* event, unsigned int T5)
207207
}
208208
return module_idxs;
209209
}
210+
210211
//____________________________________________________________________________________________
211-
std::vector<unsigned int> getHitTypesFromT5(LSTEvent* event, unsigned int T5) {
212-
return {4, 4, 4, 4, 4, 4, 4, 4, 4, 4};
213-
;
212+
std::vector<unsigned int> getModuleIdxsFromT3(LSTEvent* event, unsigned int T3) {
213+
std::vector<unsigned int> hits = getHitsFromT3(event, T3);
214+
std::vector<unsigned int> module_idxs;
215+
auto hitsEvt = event->getHits<HitsSoA>();
216+
for (auto& hitIdx : hits) {
217+
module_idxs.push_back(hitsEvt.moduleIndices()[hitIdx]);
218+
}
219+
return module_idxs;
214220
}
215221

222+
//____________________________________________________________________________________________
223+
std::vector<unsigned int> getHitTypesFromT5(LSTEvent* event, unsigned int T5) { return {4, 4, 4, 4, 4, 4, 4, 4, 4, 4}; }
224+
225+
//____________________________________________________________________________________________
226+
std::vector<unsigned int> getHitTypesFromT3(LSTEvent* event, unsigned int T5) { return {4, 4, 4, 4, 4, 4}; }
227+
216228
//____________________________________________________________________________________________
217229
std::tuple<std::vector<unsigned int>, std::vector<unsigned int>> getHitIdxsAndHitTypesFromT5(LSTEvent* event,
218230
unsigned T5) {

‎RecoTracker/LSTCore/standalone/code/core/AccessHelper.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ std::tuple<std::vector<unsigned int>, std::vector<unsigned int>> getHitIdxsAndHi
3333
std::vector<unsigned int> getLSsFromT3(LSTEvent* event, unsigned int T3);
3434
std::vector<unsigned int> getMDsFromT3(LSTEvent* event, unsigned int T3);
3535
std::vector<unsigned int> getHitsFromT3(LSTEvent* event, unsigned int T3);
36+
std::vector<unsigned int> getHitTypesFromT3(LSTEvent* event, unsigned int T3);
37+
std::vector<unsigned int> getModuleIdxsFromT3(LSTEvent* event, unsigned int T3);
3638
std::tuple<std::vector<unsigned int>, std::vector<unsigned int>> getHitIdxsAndHitTypesFromT3(LSTEvent* event,
3739
unsigned T3);
3840

‎RecoTracker/LSTCore/standalone/code/core/write_lst_ntuple.cc‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,7 @@ void createOptionalOutputBranches() {
160160

161161
// T5 DNN branches
162162
createT5DNNBranches();
163+
createT3DNNBranches();
163164

164165
#endif
165166
}
@@ -191,6 +192,40 @@ void createT5DNNBranches() {
191192
}
192193
}
193194

195+
//________________________________________________________________________________________________________________________________
196+
void createT3DNNBranches() {
197+
// Common branches for T3 properties based on TripletsSoA fields
198+
ana.tx->createBranch<std::vector<float>>("t3_betaIn");
199+
ana.tx->createBranch<std::vector<float>>("t3_centerX");
200+
ana.tx->createBranch<std::vector<float>>("t3_centerY");
201+
ana.tx->createBranch<std::vector<float>>("t3_radius");
202+
ana.tx->createBranch<std::vector<bool>>("t3_partOfPT5");
203+
ana.tx->createBranch<std::vector<bool>>("t3_partOfT5");
204+
ana.tx->createBranch<std::vector<bool>>("t3_partOfPT3");
205+
ana.tx->createBranch<std::vector<float>>("t3_pMatched");
206+
ana.tx->createBranch<std::vector<float>>("t3_sim_vxy");
207+
ana.tx->createBranch<std::vector<float>>("t3_sim_vz");
208+
209+
// Hit-specific branches (T3 has 4 hits from two segments)
210+
std::vector<std::string> hitIndices = {"0", "1", "2", "3", "4", "5"};
211+
std::vector<std::string> hitProperties = {"r", "x", "y", "z", "eta", "phi", "detId", "layer", "moduleType"};
212+
213+
for (const auto& idx : hitIndices) {
214+
for (const auto& prop : hitProperties) {
215+
std::string branchName = "t3_hit_" + idx + "_" + prop;
216+
if (prop == "detId" || prop == "layer" || prop == "moduleType") {
217+
ana.tx->createBranch<std::vector<int>>(branchName);
218+
} else {
219+
ana.tx->createBranch<std::vector<float>>(branchName);
220+
}
221+
}
222+
}
223+
224+
// Additional metadata branches
225+
ana.tx->createBranch<std::vector<int>>("t3_layer_binary");
226+
ana.tx->createBranch<std::vector<std::vector<int>>>("t3_matched_simIdx");
227+
}
228+
194229
//________________________________________________________________________________________________________________________________
195230
void createGnnNtupleBranches() {
196231
// Mini Doublets
@@ -339,6 +374,7 @@ void setOptionalOutputBranches(LSTEvent* event) {
339374
setQuintupletOutputBranches(event);
340375
setPixelTripletOutputBranches(event);
341376
setOccupancyBranches(event);
377+
setT3DNNBranches(event);
342378
setT5DNNBranches(event);
343379

344380
#endif
@@ -637,6 +673,42 @@ void setPixelTripletOutputBranches(LSTEvent* event) {
637673
ana.tx->setBranch<std::vector<int>>("pT3_isDuplicate", pT3_isDuplicate);
638674
}
639675

676+
//________________________________________________________________________________________________________________________________
677+
void fillT3DNNBranches(LSTEvent* event, unsigned int iT3) {
678+
auto hits = event->getHits<HitsSoA>();
679+
auto modules = event->getModules<ModulesSoA>();
680+
681+
std::vector<unsigned int> hitIdx = getHitsFromT3(event, iT3);
682+
std::vector<lst_math::Hit> hitObjects;
683+
684+
for (int i = 0; i < hitIdx.size(); ++i) {
685+
unsigned int hit = hitIdx[i];
686+
float x = hits.xs()[hit];
687+
float y = hits.ys()[hit];
688+
float z = hits.zs()[hit];
689+
lst_math::Hit hitObj(x, y, z);
690+
hitObjects.push_back(hitObj);
691+
692+
std::string idx = std::to_string(i);
693+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_r", sqrt(x * x + y * y));
694+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_x", x);
695+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_y", y);
696+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_z", z);
697+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_eta", hitObj.eta());
698+
ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_phi", hitObj.phi());
699+
700+
int subdet = trk.ph2_subdet()[hits.idxs()[hit]];
701+
int is_endcap = subdet == 4;
702+
int layer = trk.ph2_layer()[hits.idxs()[hit]] + 6 * is_endcap;
703+
int detId = trk.ph2_detId()[hits.idxs()[hit]];
704+
unsigned int module = hits.moduleIndices()[hit];
705+
706+
ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_detId", detId);
707+
ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_layer", layer);
708+
ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_moduleType", modules.moduleType()[module]);
709+
}
710+
}
711+
640712
//________________________________________________________________________________________________________________________________
641713
void fillT5DNNBranches(LSTEvent* event, unsigned int iT3) {
642714
auto hits = event->getHits<HitsSoA>();
@@ -687,6 +759,69 @@ void fillT5DNNBranches(LSTEvent* event, unsigned int iT3) {
687759
ana.tx->pushbackToBranch<float>("t5_t3_phi", hitObjects[0].phi());
688760
}
689761

762+
void setT3DNNBranches(LSTEvent* event) {
763+
auto const triplets = event->getTriplets<TripletsSoA>();
764+
auto const tripletsOccupancy = event->getTriplets<TripletsOccupancySoA>();
765+
auto modules = event->getModules<ModulesSoA>();
766+
auto ranges = event->getRanges();
767+
768+
for (unsigned int lowerModuleIdx = 0; lowerModuleIdx < modules.nLowerModules(); ++lowerModuleIdx) {
769+
int nTriplets = tripletsOccupancy.nTriplets()[lowerModuleIdx];
770+
for (unsigned int idx = 0; idx < nTriplets; idx++) {
771+
unsigned int tripletIndex = ranges.tripletModuleIndices()[lowerModuleIdx] + idx;
772+
773+
// Get hit indices and types
774+
std::vector<unsigned int> hit_idx = getHitsFromT3(event, tripletIndex);
775+
std::vector<unsigned int> hit_type = getHitTypesFromT3(event, tripletIndex);
776+
std::vector<unsigned int> module_idx = getModuleIdxsFromT3(event, tripletIndex);
777+
778+
// Calculate layer binary representation
779+
int layer_binary = 0;
780+
for (size_t i = 0; i < module_idx.size(); i += 2) {
781+
layer_binary |= (1 << (modules.layers()[module_idx[i]] + 6 * (modules.subdets()[module_idx[i]] == 4)));
782+
}
783+
784+
// Get matching information with percent matched
785+
float percent_matched;
786+
std::vector<int> simidx = matchedSimTrkIdxs(hit_idx, hit_type, false, &percent_matched);
787+
788+
// Fill the branches with T3-specific data
789+
ana.tx->pushbackToBranch<float>("t3_betaIn", triplets.betaIn()[tripletIndex]);
790+
ana.tx->pushbackToBranch<float>("t3_centerX", triplets.centerX()[tripletIndex]);
791+
ana.tx->pushbackToBranch<float>("t3_centerY", triplets.centerY()[tripletIndex]);
792+
ana.tx->pushbackToBranch<float>("t3_radius", triplets.radius()[tripletIndex]);
793+
ana.tx->pushbackToBranch<bool>("t3_partOfPT5", triplets.partOfPT5()[tripletIndex]);
794+
ana.tx->pushbackToBranch<bool>("t3_partOfT5", triplets.partOfT5()[tripletIndex]);
795+
ana.tx->pushbackToBranch<bool>("t3_partOfPT3", triplets.partOfPT3()[tripletIndex]);
796+
ana.tx->pushbackToBranch<int>("t3_layer_binary", layer_binary);
797+
ana.tx->pushbackToBranch<std::vector<int>>("t3_matched_simIdx", simidx);
798+
ana.tx->pushbackToBranch<float>("t3_pMatched", percent_matched);
799+
800+
// Add vertex information for matched sim tracks
801+
if (simidx.size() == 0) {
802+
// No matched sim track - set default values
803+
ana.tx->pushbackToBranch<float>("t3_sim_vxy", 0.0);
804+
ana.tx->pushbackToBranch<float>("t3_sim_vz", 0.0);
805+
} else {
806+
// Get vertex information from the first matched sim track
807+
int vtxidx = trk.sim_parentVtxIdx()[simidx[0]];
808+
float vtx_x = trk.simvtx_x()[vtxidx];
809+
float vtx_y = trk.simvtx_y()[vtxidx];
810+
float vtx_z = trk.simvtx_z()[vtxidx];
811+
812+
// Calculate transverse distance from origin
813+
float vxy = sqrt(vtx_x * vtx_x + vtx_y * vtx_y);
814+
815+
ana.tx->pushbackToBranch<float>("t3_sim_vxy", vxy);
816+
ana.tx->pushbackToBranch<float>("t3_sim_vz", vtx_z);
817+
}
818+
819+
// Fill hit-specific information
820+
fillT3DNNBranches(event, tripletIndex);
821+
}
822+
}
823+
}
824+
690825
//________________________________________________________________________________________________________________________________
691826
void setT5DNNBranches(LSTEvent* event) {
692827
auto triplets = event->getTriplets<TripletsOccupancySoA>();

‎RecoTracker/LSTCore/standalone/code/core/write_lst_ntuple.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ void createRequiredOutputBranches();
1919
void createOptionalOutputBranches();
2020
void createGnnNtupleBranches();
2121
void createT5DNNBranches();
22+
void createT3DNNBranches();
2223

2324
void fillOutputBranches(LSTEvent* event);
2425
void setOutputBranches(LSTEvent* event);
@@ -30,7 +31,9 @@ void setPixelTripletOutputBranches(LSTEvent* event);
3031
void setGnnNtupleBranches(LSTEvent* event);
3132
void setGnnNtupleMiniDoublet(LSTEvent* event, unsigned int MD);
3233
void fillT5DNNBranches(LSTEvent* event, unsigned int T3);
34+
void fillT3DNNBranches(LSTEvent* event, unsigned int iT3);
3335
void setT5DNNBranches(LSTEvent* event);
36+
void setT3DNNBranches(LSTEvent* event);
3437

3538
std::tuple<int, float, float, float, int, std::vector<int>> parseTrackCandidate(LSTEvent* event, unsigned int);
3639
std::tuple<float, float, float, std::vector<unsigned int>, std::vector<unsigned int>> parsepT5(LSTEvent* event,

0 commit comments

Comments
 (0)