Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-12-17 09:18:15

0001 #include <TSystem.h>
0002 #include <TStopwatch.h>
0003 #include <iostream>
0004 #include <vector>
0005 
0006 #include "PtCalculator.h"  // SiCaloPt::PtCalculator & friends
0007 
0008 // ---- Weights(onnx) and Scalers(json) Path ---------------------------
0009 struct DemoPaths
0010 {
0011     std::string emd_onnx          = "/sphenix/user/jzhang1/testcode4all/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/Implement/ML_Weight_Scaler/model_MLEMD.onnx"; 
0012     std::string emd_scaler_json   = "";
0013 
0014     std::string eproj_onnx        = "/sphenix/user/jzhang1/testcode4all/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/Implement/ML_Weight_Scaler/model_MLEproj.onnx"; 
0015     std::string eproj_scaler_json = "/sphenix/user/jzhang1/testcode4all/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/Implement/ML_Weight_Scaler/scaler_MLEproj.json"; 
0016 
0017     std::string combined_onnx         = "/sphenix/user/jzhang1/testcode4all/INTT-EMCAL/InttSeedingTrackDev/ML4Reco/Implement/ML_Weight_Scaler/model_MLCombined.onnx"; 
0018     std::string combined_scaler_json  = "";
0019 };
0020 
0021 // ---- turn string into optional<string> for Config ------------------------
0022 template<typename Opt>
0023 Opt make_opt(const std::string& s) 
0024 {
0025     if (s.empty()) return std::nullopt;
0026     return s;
0027 }
0028 
0029 // ---- PtCalc Tutorial -------------------------------------
0030 void PtCalcMLTutorial()
0031 {
0032     // Load PtCalc shared library
0033     gSystem->Load("libPtCalc.so");
0034     // gSystem->Load("/sphenix/u/jzhang1/install/lib/libPtCalc.so");
0035 
0036     // Use appropriate paths in your environment, the default "DemoPaths" setup is correct here for mine
0037     DemoPaths WS_Path;  
0038 
0039     // SiCaloPt::PtCalculatorConfig setting(if you want to use the ML models, Formula-method donnot need these)
0040     SiCaloPt::PtCalculatorConfig cfg;
0041     cfg.mlEMD_model_path        = make_opt<decltype(cfg.mlEMD_model_path)>(WS_Path.emd_onnx);
0042     cfg.mlEMD_scaler_json       = make_opt<decltype(cfg.mlEMD_scaler_json)>(WS_Path.emd_scaler_json);
0043     cfg.mlEproj_model_path      = make_opt<decltype(cfg.mlEproj_model_path)>(WS_Path.eproj_onnx);
0044     cfg.mlEproj_scaler_json     = make_opt<decltype(cfg.mlEproj_scaler_json)>(WS_Path.eproj_scaler_json);
0045     cfg.mlCombined_model_path   = make_opt<decltype(cfg.mlCombined_model_path)>(WS_Path.combined_onnx);
0046     cfg.mlCombined_scaler_json  = make_opt<decltype(cfg.mlCombined_scaler_json)>(WS_Path.combined_scaler_json);
0047 
0048     // PtCalculator instance
0049     SiCaloPt::PtCalculator calcTutorial(cfg);
0050     
0051     // initialize (load models and scalers for ML methods)
0052     std::string err;
0053     if (!calcTutorial.init(&err)) 
0054     {
0055         std::cout << "[init] failed: " << err << std::endl;
0056         return;
0057     }
0058     std::cout << "[init] OK\n";
0059  
0060     // ======================== EMD Formula ==========================
0061     {
0062         SiCaloPt::InputEMD in;
0063         in.EMD_Angle  = 0.025;  // delta_phi - EM Deflection angle in rad
0064         in.EMD_Eta    = 0.00;   // EMCal Cluster eta
0065         in.EMD_Radius = 93.5;   // EMCal Cluster radius in cm
0066 
0067         calcTutorial.setParCeta(0.2);   // set C_eta parameter if needed
0068         calcTutorial.setParPower(-1.0); // set Power parameter if needed
0069 
0070         auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodEMD, SiCaloPt::AnyInput{in});
0071         std::cout << "[EMD-analytic] ok=" << r.ok
0072                   << "  pt=" << r.pt_reco
0073                   << "  err=\"" << r.err << "\"\n";
0074     }
0075 
0076     // ======================== Eproj Formula ==========================
0077     {
0078         SiCaloPt::InputEproj in;
0079         in.Energy_Calo   = 1.8;   // EMCal Cluster energy in GeV
0080         in.Radius_Calo   = 93.5;  // EMCal Cluster radius in cm
0081         in.Z_Calo        = 0.0;   // EMCal Cluster z in cm
0082         in.Radius_vertex = 0.0;   // Vertex radius in cm
0083         in.Z_vertex      = 0.0;   // Vertex z in cm
0084 
0085         auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodEproj, SiCaloPt::AnyInput{in});
0086         std::cout << "[Eproj-analytic] ok=" << r.ok
0087                   << "  pt=" << r.pt_reco
0088                   << "  err=\"" << r.err << "\"\n";
0089     }
0090 
0091     // ============= ML:MLEMD (2-d input: 1/dphi_EMD, eta_track = 0 now) ===============
0092     {
0093         // 2-d input features:{ dphi_EMD, eta_track }
0094         std::vector<float> featsMLEMD = {15, 0};
0095 
0096         SiCaloPt::InputMLEMD in{featsMLEMD};
0097         auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodMLEMD, SiCaloPt::AnyInput{in});
0098         std::cout << "[MLEMD-2D] ok=" << r.ok
0099                 << "  pt=" << r.pt_reco
0100                 << "  err=\"" << r.err << "\"\n";
0101     }
0102 
0103     // ====== ML:MLEproj (7-d input: INTT 3/4 layer R,Z, INTT 5/4 layer R,Z, Calo R,Z,Energy) ======
0104     {
0105         // 7-d input features:{ INTT 3/4 layer R, INTT 3/4 layer Z, INTT 5/4 layer R, INTT 5/4 layer Z, Calo R, Calo Z, Calo Energy }
0106         std::vector<float> featsMLEproj = { 10.0,  5.0,   // INTT 3/4 layer
0107                                             15.0,  7.5,   // INTT 5/4 layer
0108                                             100.0, 50.0, 8.0 }; // Calo R,Z,Energy
0109 
0110         SiCaloPt::InputMLEproj in{featsMLEproj};
0111         auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodMLEproj, SiCaloPt::AnyInput{in});
0112         std::cout << "[MLEproj-7D] ok=" << r.ok
0113                 << "  pt=" << r.pt_reco
0114                 << "  err=\"" << r.err << "\"\n";
0115     }
0116 
0117     // ============ ML:Combined/Gate (2-d input: pt_from_MLEMD, pt_from_MLEproj) ===================
0118     {
0119         // 2-d input features:{ pt_from_MLEMD, pt_from_MLEproj }
0120         std::vector<float> featsMLCombined = {8.0, 9.5};
0121 
0122         SiCaloPt::InputMLCombined in{featsMLCombined};
0123         auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodMLCombined, SiCaloPt::AnyInput{in});
0124         std::cout << "[MLCombined] ok=" << r.ok
0125                 << "  pt=" << r.pt_reco
0126                 << "  err=\"" << r.err << "\"\n";
0127     }
0128 
0129     // // ======================== Batch example ========================
0130     // {
0131     //     std::vector<SiCaloPt::InputEproj> batch(1000);
0132     //     for (size_t i = 0; i < batch.size(); ++i) {
0133     //     auto& x = batch[i];
0134     //     x.Energy_Calo   = 1.0 + 0.001 * i;
0135     //     x.Radius_Calo   = 93.5;
0136     //     x.Z_Calo        = 0.0;
0137     //     x.Radius_vertex = 0.0;
0138     //     x.Z_vertex      = 0.0;
0139     //     }
0140 
0141     //     TStopwatch sw; sw.Start();
0142     //     double sum_pt = 0.0; size_t ok_cnt = 0;
0143     //     for (auto& x : batch) {
0144     //     auto r = calcTutorial.ComputePt(SiCaloPt::Method::MethodEproj, SiCaloPt::AnyInput{x});
0145     //     if (r.ok) { sum_pt += r.pt_reco; ++ok_cnt; }
0146     //     }
0147     //     sw.Stop();
0148 
0149     //     std::cout << "[Batch Eproj] ok=" << ok_cnt << "/" << batch.size()
0150     //             << "  avg_pt=" << (ok_cnt ? sum_pt/ok_cnt : 0.0)
0151     //             << "  time=" << sw.RealTime() << " s\n";
0152     // }
0153 }