From 0bdc88aaca47d00d71fbfd73f99b6f4930a174a0 Mon Sep 17 00:00:00 2001 From: Andrew Edmonds Date: Mon, 30 Mar 2026 10:48:02 -0500 Subject: [PATCH 1/3] Run the TrackQuality ANN directly from the onnx file with ONNX Runtime --- TrkDiag/data/TrkQual_ANN1_v2.onnx | Bin 0 -> 2529 bytes TrkDiag/src/SConscript | 6 ++- TrkDiag/src/TrackQuality_module.cc | 82 +++++++++++++++++++++++++++-- 3 files changed, 82 insertions(+), 6 deletions(-) create mode 100644 TrkDiag/data/TrkQual_ANN1_v2.onnx diff --git a/TrkDiag/data/TrkQual_ANN1_v2.onnx b/TrkDiag/data/TrkQual_ANN1_v2.onnx new file mode 100644 index 0000000000000000000000000000000000000000..7b89bcbb790123d1384a43cd0b1c56dd1a4bc71d GIT binary patch literal 2529 zcmbu>e^Aq990%}yk=Xn?5oaaRovutbp9%tgpBh^D)k!Ee67|e%8T)d{##U?Lzk7Q z)hRMbMAUsEvB7t;!l`74xm;r~%@1ca5ka0$<=Cf;MHZ9ypNkjaGgpCLS7y{}oqdD{ z8jEGSLgW>KyhjC6=v_fA4QM#mPd1?996#azIHo%L_gTJYb{BT)I3-^o$1@*~@29D| znm$g>Zd{{xJ+BhvJ*)T;O`L#eqId$L?!qptG{2(Wq5mQ3T}meLBbw+ob9Un^>KVO{ zs3+$&**qji* zeYkgz&yA}(Kv*PyqltFFf~p>9sw;xDAq)0Tslt%q9$23IHGcP420lFJFjnn&4w$t^ z!Pxy1HPEsL5B>HU%9(T=>fR2P830_(S?KP~#B6IPo{ku1p85SCtX}aZ$mn=*tS>@o z+cDgF@ia!=cnV%+K8Lf$Gax(p23l8rfk%(`;YW*4LrH5ET<@`~RSL4W;J*cbv6sPMBV&?G+5SJDXQx;ip+PX$2SCW8jIU=l|mxZQ!5g@%E zYDSuvmj|P9|MeQI-MW)ewk|=;vqJ1(J=iD>Zr;?4hZCxBRDS?1y6rfyGoR}JqybNL zJPo_nRKx17rLa4GAH{9f!OB5B3iHPn?h7dmcwANdZh?|1eg76!7uP_A&XH1ef8L-L zZmOZ$`evi>rWO3N%;QO4E`-Ai+A=lUyUuaw)F?#II*HJJ|`m!hc$ z-`ax93i*t{l#Pz&NqAxUHpY63%Wxxpm9uQ2^BgN%sS*2L=KQ=`cyEBuWXW!z`SQ?Usex z6#Wh5iidbbiq4s0#r4Q%73aR%VE=MiuzczUCA5)Uw#Jl3sLNC{4M!58wUWyS^TsI0 z|H4u6XNOM5 #include @@ -66,20 +68,59 @@ namespace mu2e std::shared_ptr mva_; + Ort::Env _env; + Ort::SessionOptions _session_options; + Ort::Session _session; + Ort::AllocatorWithDefaultOptions _allocator; + Ort::AllocatedStringPtr _input_name; + Ort::TypeInfo _type_info; + Ort::ConstTensorTypeAndShapeInfo _tensor_info; + std::vector _input_shape; + size_t _total_size; + Ort::MemoryInfo _memory_info; + Ort::AllocatedStringPtr _output_name; + + std::string print_shape(const std::vector& v) { + std::stringstream ss(""); + for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x"; + ss << v[v.size() - 1]; + return ss.str(); + } }; TrackQuality::TrackQuality(const Parameters& conf) : art::EDProducer{conf}, _kalSeedPtrTag(conf().kalSeedPtrTag()), _printMVA(conf().printMVA()), - _debug(conf().debug()) + _debug(conf().debug()), + + _env(ORT_LOGGING_LEVEL_WARNING, "ONNXInference"), + _session(_env, "ArtAnalysis/TrkDiag/data/TrkQual_ANN1_v2.onnx", _session_options), + _input_name(_session.GetInputNameAllocated(0, _allocator)), + _type_info(_session.GetInputTypeInfo(0)), + _tensor_info(_type_info.GetTensorTypeAndShapeInfo()), + _input_shape(_tensor_info.GetShape()), // Get input shape from model + _memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)), + _output_name(_session.GetOutputNameAllocated(0, _allocator)) { produces(); ConfigFileLookupPolicy configFile; mva_ = std::make_shared(configFile(conf().datFilename())); - } + + // Handle dynamic dimensions if needed + for (auto& dim : _input_shape) { + if (dim == -1) { dim = 1; } // Set dynamic dims to 1 (or your desired value) + } + + // Calculate total size + _total_size = 1; + for (auto dim : _input_shape) { + _total_size *= dim; + } + } + void TrackQuality::produce(art::Event& event ) { // create output unique_ptr mvacol(new MVAResultCollection()); @@ -89,6 +130,9 @@ namespace mu2e event.getByLabel(_kalSeedPtrTag, kalSeedPtrHandle); const auto& kalSeedPtrs = *kalSeedPtrHandle; + // Prepare input tensor + std::vector input_tensor_values(_total_size, 0.0f); // Initialize with zeros + // Go through the tracks and calculate their track qualities for (const auto& kalSeedPtr : kalSeedPtrs) { const auto& kalSeed = *kalSeedPtr; @@ -135,6 +179,11 @@ namespace mu2e features[3] = (double) nnullambig / nactive; features[4] = kalSeed.fitConsistency(); features[6] = (double)nmatactive / nactive; + input_tensor_values[0] = nactive; + input_tensor_values[1] = (double) nactive / nhits; + input_tensor_values[3] = (double) nnullambig / nactive; + input_tensor_values[4] = kalSeed.fitConsistency(); + input_tensor_values[6] = (double)nmatactive / nactive; // Now get the features that are for the entrance of the trackre bool entrance_found = false; @@ -143,6 +192,8 @@ namespace mu2e if (kinter.surfaceId() == SurfaceIdDetail::TT_Front) { // we only want the tracker entrance (sid=0) features[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); features[5] = kinter.momerr(); + input_tensor_values[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); + input_tensor_values[5] = kinter.momerr(); entrance_found = true; break; } @@ -150,18 +201,41 @@ namespace mu2e if (!entrance_found) { features[2] = -9999; features[5] = -9999; + input_tensor_values[2] = -9999; + input_tensor_values[5] = -9999; } std::vector mvaout = mva_->infer(features.data()); + Ort::Value input_tensor = Ort::Value::CreateTensor(_memory_info, + input_tensor_values.data(), + input_tensor_values.size(), + _input_shape.data(), + _input_shape.size() + ); + // Run inference + const char* input_names[] = {_input_name.get()}; + const char* output_names[] = {_output_name.get()}; + auto output_tensors = _session.Run( + Ort::RunOptions{nullptr}, + input_names, + &input_tensor, + 1, + output_names, + 1 + ); + // Get output + float* output_data = output_tensors[0].GetTensorMutableData(); + if (!entrance_found) { mvaout[0] = 0; // this is not a good track + output_data[0] = 0; } if(_debug > 0) { - printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f\n", + printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f (ORT: %.4f)\n", __func__, moduleDescription().moduleLabel().c_str(), - features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0]); + features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0], output_data[0]); } mvacol->push_back(MVAResult(mvaout[0])); From 95fe55759ead329c574f23472d2fdb2a44ab5fa4 Mon Sep 17 00:00:00 2001 From: Andrew Edmonds Date: Mon, 30 Mar 2026 11:17:23 -0500 Subject: [PATCH 2/3] Let onnx file be defined in fcl --- TrkDiag/src/TrackQuality_module.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/TrkDiag/src/TrackQuality_module.cc b/TrkDiag/src/TrackQuality_module.cc index e43a1ca..6340306 100644 --- a/TrkDiag/src/TrackQuality_module.cc +++ b/TrkDiag/src/TrackQuality_module.cc @@ -52,6 +52,7 @@ namespace mu2e fhicl::Atom kalSeedPtrTag{Name("KalSeedPtrCollection"), Comment("Input tag for KalSeedPtrCollection")}; fhicl::Atom printMVA{Name("PrintMVA"), Comment("Print the MVA used"), false}; fhicl::Atom datFilename{Name("datFilename"), Comment("Filename for the .dat file to use")}; + fhicl::Atom onnxFilename{Name("onnxFilename"), Comment("Filename for the .onnx file to use")}; fhicl::Atom debug{Name("debugLevel"), Comment("Debug printout level"), 0}; }; @@ -68,6 +69,8 @@ namespace mu2e std::shared_ptr mva_; + ConfigFileLookupPolicy _configFileLookup; + Ort::Env _env; Ort::SessionOptions _session_options; Ort::Session _session; @@ -95,7 +98,7 @@ namespace mu2e _debug(conf().debug()), _env(ORT_LOGGING_LEVEL_WARNING, "ONNXInference"), - _session(_env, "ArtAnalysis/TrkDiag/data/TrkQual_ANN1_v2.onnx", _session_options), + _session(_env, _configFileLookup(conf().onnxFilename()).c_str(), _session_options), _input_name(_session.GetInputNameAllocated(0, _allocator)), _type_info(_session.GetInputTypeInfo(0)), _tensor_info(_type_info.GetTensorTypeAndShapeInfo()), @@ -105,8 +108,7 @@ namespace mu2e { produces(); - ConfigFileLookupPolicy configFile; - mva_ = std::make_shared(configFile(conf().datFilename())); + mva_ = std::make_shared(_configFileLookup(conf().datFilename())); // Handle dynamic dimensions if needed From 58dc7e32abd53f7e649e19dd551ecc87c1cf2e73 Mon Sep 17 00:00:00 2001 From: Andrew Edmonds Date: Tue, 28 Apr 2026 15:58:50 -0500 Subject: [PATCH 3/3] Remove TMVA::SOFIE from TrkQual --- TrkDiag/src/TrackQuality_module.cc | 39 ++++++------------------------ 1 file changed, 8 insertions(+), 31 deletions(-) diff --git a/TrkDiag/src/TrackQuality_module.cc b/TrkDiag/src/TrackQuality_module.cc index 6340306..4ba84c8 100644 --- a/TrkDiag/src/TrackQuality_module.cc +++ b/TrkDiag/src/TrackQuality_module.cc @@ -21,7 +21,6 @@ // data #include "Offline/RecoDataProducts/inc/KalSeed.hh" #include "Offline/RecoDataProducts/inc/MVAResult.hh" -#include "ArtAnalysis/TrkDiag/inc/TrkQual_ANN1.hxx" // ONNXRuntime #include "onnxruntime/core/session/onnxruntime_cxx_api.h" // C++ @@ -35,10 +34,6 @@ using namespace std; using CLHEP::Hep3Vector; using CLHEP::HepVector; -namespace TMVA_SOFIE_TrkQual_ANN1 { - class Session; -} - namespace mu2e { @@ -51,7 +46,6 @@ namespace mu2e fhicl::Atom kalSeedPtrTag{Name("KalSeedPtrCollection"), Comment("Input tag for KalSeedPtrCollection")}; fhicl::Atom printMVA{Name("PrintMVA"), Comment("Print the MVA used"), false}; - fhicl::Atom datFilename{Name("datFilename"), Comment("Filename for the .dat file to use")}; fhicl::Atom onnxFilename{Name("onnxFilename"), Comment("Filename for the .onnx file to use")}; fhicl::Atom debug{Name("debugLevel"), Comment("Debug printout level"), 0}; }; @@ -67,8 +61,6 @@ namespace mu2e bool _printMVA; int _debug; - std::shared_ptr mva_; - ConfigFileLookupPolicy _configFileLookup; Ort::Env _env; @@ -108,9 +100,6 @@ namespace mu2e { produces(); - mva_ = std::make_shared(_configFileLookup(conf().datFilename())); - - // Handle dynamic dimensions if needed for (auto& dim : _input_shape) { if (dim == -1) { dim = 1; } // Set dynamic dims to 1 (or your desired value) @@ -122,7 +111,7 @@ namespace mu2e _total_size *= dim; } } - + void TrackQuality::produce(art::Event& event ) { // create output unique_ptr mvacol(new MVAResultCollection()); @@ -133,12 +122,12 @@ namespace mu2e const auto& kalSeedPtrs = *kalSeedPtrHandle; // Prepare input tensor - std::vector input_tensor_values(_total_size, 0.0f); // Initialize with zeros + std::vector features(_total_size, 0.0f); // Initialize with zeros + // Go through the tracks and calculate their track qualities for (const auto& kalSeedPtr : kalSeedPtrs) { const auto& kalSeed = *kalSeedPtr; - std::array features; // the features we trained on // fill the hit count variables int nhits = 0; int nactive = 0; int ndouble = 0; int ndactive = 0; int nnullambig = 0; @@ -181,11 +170,6 @@ namespace mu2e features[3] = (double) nnullambig / nactive; features[4] = kalSeed.fitConsistency(); features[6] = (double)nmatactive / nactive; - input_tensor_values[0] = nactive; - input_tensor_values[1] = (double) nactive / nhits; - input_tensor_values[3] = (double) nnullambig / nactive; - input_tensor_values[4] = kalSeed.fitConsistency(); - input_tensor_values[6] = (double)nmatactive / nactive; // Now get the features that are for the entrance of the trackre bool entrance_found = false; @@ -194,8 +178,6 @@ namespace mu2e if (kinter.surfaceId() == SurfaceIdDetail::TT_Front) { // we only want the tracker entrance (sid=0) features[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); features[5] = kinter.momerr(); - input_tensor_values[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); - input_tensor_values[5] = kinter.momerr(); entrance_found = true; break; } @@ -203,15 +185,11 @@ namespace mu2e if (!entrance_found) { features[2] = -9999; features[5] = -9999; - input_tensor_values[2] = -9999; - input_tensor_values[5] = -9999; } - std::vector mvaout = mva_->infer(features.data()); - Ort::Value input_tensor = Ort::Value::CreateTensor(_memory_info, - input_tensor_values.data(), - input_tensor_values.size(), + features.data(), + features.size(), _input_shape.data(), _input_shape.size() ); @@ -227,17 +205,16 @@ namespace mu2e 1 ); // Get output - float* output_data = output_tensors[0].GetTensorMutableData(); + float* mvaout = output_tensors[0].GetTensorMutableData(); if (!entrance_found) { mvaout[0] = 0; // this is not a good track - output_data[0] = 0; } if(_debug > 0) { - printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f (ORT: %.4f)\n", + printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f\n", __func__, moduleDescription().moduleLabel().c_str(), - features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0], output_data[0]); + features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0]); } mvacol->push_back(MVAResult(mvaout[0]));