diff --git a/CMakeLists.txt b/CMakeLists.txt index 87f21eb..b2fa789 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ ######################################################## # cmake file for building LCFIPlus package # @author Tomohiko Tanabe (ICEPP, The University of Tokyo) -CMAKE_MINIMUM_REQUIRED(VERSION 2.6 FATAL_ERROR) +CMAKE_MINIMUM_REQUIRED(VERSION 3.5 FATAL_ERROR) ######################################################## @@ -14,8 +14,6 @@ SET( ${PROJECT_NAME}_VERSION_MAJOR 0 ) SET( ${PROJECT_NAME}_VERSION_MINOR 11 ) SET( ${PROJECT_NAME}_VERSION_PATCH 0 ) - - ### DEPENDENCIES ############################################################ FIND_PACKAGE( ILCUTIL COMPONENTS ILCSOFT_CMAKE_MODULES REQUIRED ) @@ -23,19 +21,48 @@ FIND_PACKAGE( ILCUTIL COMPONENTS ILCSOFT_CMAKE_MODULES REQUIRED ) # load default settings from ILCSOFT_CMAKE_MODULES INCLUDE( ilcsoft_default_settings ) +# Set up C++ Standard +# ``-DCMAKE_CXX_STANDARD=`` when invoking CMake +set(CMAKE_CXX_STANDARD 17 CACHE STRING "") + +if(NOT CMAKE_CXX_STANDARD MATCHES "17|20") + message(FATAL_ERROR "Unsupported C++ standard: ${CMAKE_CXX_STANDARD}") +endif() + +# Prevent CMake falls back to the latest standard the compiler does support +set(CMAKE_CXX_STANDARD_REQUIRED ON) + + +# ONNX support option +OPTION( ENABLE_ONNX "Enable ONNX Runtime support for ML inference" OFF ) FIND_PACKAGE( Marlin 1.0 REQUIRED ) FIND_PACKAGE( MarlinUtil REQUIRED ) #FIND_PACKAGE( ROOT REQUIRED COMPONENTS Minuit2 TMVA TreePlayer ) FIND_PACKAGE( ROOT REQUIRED COMPONENTS Minuit2 TMVA TreePlayer Gui Geom Eve Minuit XMLIO RGL Ged EG MLP ) +#FIND_PACKAGE( Torch REQUIRED ) FIND_PACKAGE( LCFIVertex REQUIRED ) +# Optional ONNX dependencies +IF( ENABLE_ONNX ) + MESSAGE( STATUS "ONNX support enabled" ) + FIND_PACKAGE( onnxruntime 1.17.1 REQUIRED) + FIND_PACKAGE( nlohmann_json REQUIRED) + ADD_DEFINITIONS( -DENABLE_ONNX ) +ELSE() + MESSAGE( STATUS "ONNX support disabled" ) +ENDIF() + +#set(JSON_BuildTests OFF CACHE INTERNAL "") +#add_subdirectory(nlohmann_json) + INCLUDE_DIRECTORIES( SYSTEM ${Marlin_INCLUDE_DIRS} ) ADD_DEFINITIONS ( ${Marlin_DEFINITIONS} ) INCLUDE_DIRECTORIES( SYSTEM ${MarlinUtil_INCLUDE_DIRS} ) ADD_DEFINITIONS ( ${MarlinUtil_DEFINITIONS} ) INCLUDE_DIRECTORIES( SYSTEM ${ROOT_INCLUDE_DIRS} ) INCLUDE_DIRECTORIES( SYSTEM ${LCFIVertex_INCLUDE_DIRS} ) +#INCLUDE_DIRECTORIES( SYSTEM ${TORCH_INCLUDE_DIRS} ) # left here for backwards compatibility INCLUDE_DIRECTORIES( SYSTEM ${LCFIVertex_ROOT}/vertex_lcfi ${LCFIVertex_ROOT}/boost ) @@ -58,8 +85,20 @@ SET( ROOT_DICT_INPUT_HEADERS ${PROJECT_SOURCE_DIR}/include/TrackNtuple.h ${PROJECT_SOURCE_DIR}/include/VertexMassRecovery.h ${PROJECT_SOURCE_DIR}/include/VertexNtuple.h - ${PROJECT_SOURCE_DIR}/include/LinkDef.h + ${PROJECT_SOURCE_DIR}/include/vertexfinderdnn.h + ${PROJECT_SOURCE_DIR}/include/DNNProvider2.h ) + +# Add ONNX-related headers if enabled +IF( ENABLE_ONNX ) + LIST( APPEND ROOT_DICT_INPUT_HEADERS + ${PROJECT_SOURCE_DIR}/include/MLInputGenerator.h + ${PROJECT_SOURCE_DIR}/include/MLMakeNtuple.h + ${PROJECT_SOURCE_DIR}/include/MLInferenceWeaver.h + ) +ENDIF() + +LIST( APPEND ROOT_DICT_INPUT_HEADERS ${PROJECT_SOURCE_DIR}/include/LinkDef.h ) SET( ROOT_DICT_INCLUDE_DIRS ${Marlin_INCLUDE_DIRS} ${LCIO_INCLUDE_DIRS}) GEN_ROOT_DICT_SOURCES( dict.cc ) INSTALL(FILES ${ROOT_DICT_OUTPUT_DIR}/dict_rdict.pcm DESTINATION lib) @@ -79,11 +118,13 @@ ENDIF() # include directories INCLUDE_DIRECTORIES( "./include" ) +#LINK_DIRECTORIES( "./lib" ) # definitions to pass to the compiler ADD_DEFINITIONS( "-Wno-effc++ -Wno-shadow" ) ADD_DEFINITIONS( "-Wno-long-long" ) ADD_DEFINITIONS( "-Wno-strict-aliasing" ) # avoid warnings in dict.cc + # shut up warnings in boost #ADD_DEFINITIONS( "-Wno-unused-local-typedefs" ) @@ -91,9 +132,39 @@ ADD_DEFINITIONS( "-Wno-strict-aliasing" ) # avoid warnings in dict.cc SET( libname ${PROJECT_NAME} ) AUX_SOURCE_DIRECTORY( ./src library_sources ) LIST( REMOVE_ITEM library_sources ${binary_sources} ) + +# Remove ONNX-related source files if ONNX is disabled +IF( NOT ENABLE_ONNX ) + LIST( REMOVE_ITEM library_sources + ./src/ONNXRuntime.cc + ./src/MLInferenceWeaver.cc + ./src/MLInputGenerator.cc + ./src/MLMakeNtuple.cc + ./src/WeaverInterface.cc + ) +ENDIF() + ADD_SHARED_LIBRARY( ${libname} ${library_sources} ) INSTALL_SHARED_LIBRARY( ${libname} DESTINATION lib ) -TARGET_LINK_LIBRARIES( ${libname} ${Marlin_LIBRARIES} ${ROOT_LIBRARIES} ${ROOT_COMPONENT_LIBRARIES} ${MarlinUtil_LIBRARIES} ${LCFIVertex_LIBRARIES} ) + +# Link libraries +SET( LINK_LIBS + ${Marlin_LIBRARIES} + ${ROOT_LIBRARIES} + ${ROOT_COMPONENT_LIBRARIES} + ${MarlinUtil_LIBRARIES} + ${LCFIVertex_LIBRARIES} +) + +# Add ONNX libraries if enabled +IF( ENABLE_ONNX ) + LIST( APPEND LINK_LIBS + nlohmann_json::nlohmann_json + onnxruntime::onnxruntime + ) +ENDIF() + +TARGET_LINK_LIBRARIES( ${libname} ${LINK_LIBS} ) ### EXECUTABLE MODE ######################################################## @@ -108,7 +179,7 @@ LIST(APPEND binary_sources ${ROOT_DICT_OUTPUT_SOURCES} ) ADD_EXECUTABLE( lcfiplus_bin EXCLUDE_FROM_ALL ${binary_sources} ) SET_TARGET_PROPERTIES( lcfiplus_bin PROPERTIES COMPILE_FLAGS "-DBUILD_EVE" OUTPUT_NAME lcfiplus ) -TARGET_LINK_LIBRARIES( lcfiplus_bin ${libname} ${ROOT_LIBRARIES} ${ROOT_COMPONENT_LIBRARIES} ) +TARGET_LINK_LIBRARIES( lcfiplus_bin ${libname} torch ${ROOT_LIBRARIES} ${ROOT_COMPONENT_LIBRARIES} ) # display some variables and write them to cache DISPLAY_STD_VARIABLES() diff --git a/doc/ReleaseNotes.md b/doc/ReleaseNotes.md index 2b3ec46..32e3ae2 100644 --- a/doc/ReleaseNotes.md +++ b/doc/ReleaseNotes.md @@ -1,3 +1,40 @@ +# Development branch additions (merged 2025-12-06) + +* 2025-12-05 SUEHARA Taikan + - Implement backward compatibility for MC-PFO assignment + - Add default parameter to InitMCPPFOCollections for backward compatibility + - Use simple method (first element) when Track-MC relation is not available + - Use improved method (weight-based, multi-track support) when available + - Maintain full backward compatibility with upstream v00-11 + +* 2025-12-03 SUEHARA Taikan + - Add optional ONNX support with backward compatibility + - Enable ONNX Runtime support as optional feature (ENABLE_ONNX CMake option, default: OFF) + - Make onnxruntime and nlohmann_json optional dependencies + - Conditionally compile ONNX-related source files + - Full backward compatibility when ENABLE_ONNX=OFF + +* SUEHARA Taikan and collaborators (2024-2025) + - Machine Learning and ONNX integration + - Add MLInputGenerator, MLMakeNtuple, MLInferenceWeaver for ML-based flavor tagging + - Add WeaverInterface and ONNXRuntime for ONNX model inference + - Add DNNProvider2 for DNN-based vertex finding + - Implement event-based classification with jets + - Add dEdx support for particle identification + + - Flavor tagging improvements + - Improve PFA-track assignment and track-MC assignment + - Implement true jet flavor assignment from MC + - Add MC-to-jet assignment algorithm (AssignJetsToMC) + - Bugfixes on MC flavor assignment + - Add sorted track and neutral accessors (getAllTracksSorted, getNeutralsSorted) + + - Code quality and compatibility + - Update C++ standard to C++17 with CMake 3.5+ requirement + - Compatibility fixes for key4hep environment and onnxruntime + - Various bugfixes in weaver output and neutral PF candidate masking + - Add event-based input support + # v00-11 * 2025-02-24 Thomas Madlener ([PR#73](https://github.com/LCFIPlus/LCFIPlus/pull/73)) diff --git a/examples/steering/ml_inference_test.xml b/examples/steering/ml_inference_test.xml new file mode 100644 index 0000000..f49b170 --- /dev/null +++ b/examples/steering/ml_inference_test.xml @@ -0,0 +1,81 @@ + + + + + + + + + + + + + + + + + + + + + + + /group/ilc/users/rtagami/correction_dEdx/bb/output/n000.d_dstm_14986_0.slcio + + + + + MESSAGE0 + + + + + + ${CompactFile} + + + + + + MLInferenceWeaver + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + ml_inference_test_output.slcio + + WRITE_NEW + + false + + 1996488704 + + + diff --git a/examples/steering/ml_training_data_collection.xml b/examples/steering/ml_training_data_collection.xml new file mode 100644 index 0000000..0fee3a4 --- /dev/null +++ b/examples/steering/ml_training_data_collection.xml @@ -0,0 +1,64 @@ + + + + + + + + + + + + + + + + + + + + + +test_nnbb.slcio + + + + + MESSAGE0 + + + + + + ${CompactFile} + + + + + + MLMakeNtuple + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/include/DNNProvider2.h b/include/DNNProvider2.h new file mode 100644 index 0000000..8141dc4 --- /dev/null +++ b/include/DNNProvider2.h @@ -0,0 +1,225 @@ +// DNNProvider2.h + +#ifndef dnnprovider2_h +#define dnnprovider2_h 1 + +#include "lcfiplus.h" +#include "TTree.h" +#include "TFile.h" +#include "TVector3.h" + +#include + +namespace lcfiplus { + +class DNNProvider2 : public Algorithm { + public: + DNNProvider2() {} + virtual ~DNNProvider2() {} + + void init(Parameters* param); + void process(); + void end(); + ClassDef(DNNProvider2,1); + private: + TTree* _ntp; + TFile* _file; + JetVec* _jets; //! + string _jetname; + int _mcIsB; + int _mcIsC; + int _mcIsU; + int _mcIsD; + int _mcIsS; + int _mcIsG; + + // FCC functions + float calc_dxy(float, float, float, TVector3, TVector3, int); + float calc_dz(float, float, float, TVector3, TVector3, int); + float calc_sip2d(float, float, float, float); + float calc_sip3d(float, float, float, TVector3); + float calc_jetDist(float, float, float, TVector3, TVector3); + void calc_thetaphi(TVector3, TVector3, float&, float&); + + struct DNNData { + float jet_px; + float jet_py; + float jet_pz; + float jet_e; + float jet_mass; + int jet_ntracks; + int jet_nneutrals; + + float jet_theta; + float jet_phi; + + // for charged hadrons + std::vector px; + std::vector py; + std::vector pz; + std::vector e; + std::vector efrac; + std::vector erel_log; + std::vector dtheta; + std::vector dphi; + std::vector dtheta_ilc; + std::vector dphi_ilc; + + std::vector dEdx; + + // covariant matrix of tracks + std::vector cov_d0; + std::vector cov_z0; + std::vector cov_phi; + std::vector cov_omega; + std::vector cov_tanlambda; + + std::vector cov_d0_z0; + std::vector cov_d0_phi; + std::vector cov_d0_omega; + std::vector cov_d0_tanlambda; + + std::vector cov_z0_phi; + std::vector cov_z0_omega; + std::vector cov_z0_tanlambda; + + std::vector cov_phi_omega; + std::vector cov_phi_tanlambda; + std::vector cov_omega_tanlambda; + + std::vector d0; + std::vector d0sig; + std::vector z0; + std::vector z0sig; + std::vector ip3d; + std::vector ip3dsig; + + std::vector dxy; + std::vector dz; + std::vector ip2d_fcc; + std::vector ip2dsig_fcc; + std::vector ip3d_fcc; + std::vector ip3dsig_fcc; + std::vector jetdist_fcc; + std::vector jetdistsig_fcc; + + std::vector charge; + std::vector ismuon; + std::vector iselectron; + std::vector isphoton; + std::vector ispion; + std::vector iskaon; + std::vector isproton; + std::vector ischargedhadron; + std::vector isneutralhadron; + std::vector iskaon0; + std::vector pdg_pfa; + std::vector mcpid; + std::vector mcp_pdg; + std::vector K_pdg_pfa; //20240203 + + // for PI3 (20240203) + std::vector proton_K; + std::vector pion_K; + std::vector proton_Klike; + std::vector pion_Klike; + std::vector electron_dEdxdistance; + std::vector muon_dEdxdistance; + std::vector kaon_dEdxdistance; + std::vector pion_dEdxdistance; + std::vector proton_dEdxdistance; + + std::vector ismuonlike; + std::vector iselectronlike; + std::vector ispionlike; + std::vector iskaonlike; + std::vector isprotonlike; + + std::vector ismuonprob; + std::vector iselectronprob; + std::vector ispionprob; + std::vector iskaonprob; + std::vector isprotonprob; + + + int mc_b; + int mc_c; + int mc_u; + int mc_d; + int mc_s; + int mc_g; + int mc_q; + + // for neutral hadrons + std::vector neu_px; + std::vector neu_py; + std::vector neu_pz; + std::vector neu_e; + std::vector neu_efrac; + std::vector neu_erel_log; + std::vector neu_dtheta; + std::vector neu_dphi; + std::vector neu_dtheta_ilc; + std::vector neu_dphi_ilc; + + // covariant matrix of tracks + std::vector neu_cov_d0; + std::vector neu_cov_z0; + std::vector neu_cov_phi; + std::vector neu_cov_omega; + std::vector neu_cov_tanlambda; + + std::vector neu_cov_d0_z0; + std::vector neu_cov_d0_phi; + std::vector neu_cov_d0_omega; + std::vector neu_cov_d0_tanlambda; + + std::vector neu_cov_z0_phi; + std::vector neu_cov_z0_omega; + std::vector neu_cov_z0_tanlambda; + + std::vector neu_cov_phi_omega; + std::vector neu_cov_phi_tanlambda; + std::vector neu_cov_omega_tanlambda; + + std::vector neu_d0; + std::vector neu_d0sig; + std::vector neu_z0; + std::vector neu_z0sig; + std::vector neu_ip3d; + std::vector neu_ip3dsig; + + std::vector neu_dxy; + std::vector neu_dz; + std::vector neu_ip2d_fcc; + std::vector neu_ip2dsig_fcc; + std::vector neu_ip3d_fcc; + std::vector neu_ip3dsig_fcc; + std::vector neu_jetdist_fcc; + std::vector neu_jetdistsig_fcc; + + std::vector neu_charge; + std::vector neu_ismuon; + std::vector neu_iselectron; + std::vector neu_isphoton; + std::vector neu_ispion; + std::vector neu_iskaon; + std::vector neu_isproton; + std::vector neu_ischargedhadron; + std::vector neu_isneutralhadron; + std::vector neu_iskaon0; + std::vector neu_pdg_pfa; + std::vector neu_mcpid; + std::vector neu_mcp_pdg; + std::vector neu_K_pdg_pfa; //20240201 + + + + + }; + DNNData _data; +}; + +} + +#endif diff --git a/include/EventStore.h b/include/EventStore.h index f06dd60..5bf5edf 100644 --- a/include/EventStore.h +++ b/include/EventStore.h @@ -67,6 +67,10 @@ class EventStore { } bool IsExist(const char* name, const char* classname)const; + // Flag management + bool AddFlags(const char* name, int flags); + bool RemoveFlags(const char* name, int flags); + // Object retrieval const char* GetClassName(const char* name, int idx = 0)const; void* GetObject(const char* name, const char* classname = "")const; // CAUTION: no type check diff --git a/include/LCIOStorer.h b/include/LCIOStorer.h index 47cc970..840b16b 100644 --- a/include/LCIOStorer.h +++ b/include/LCIOStorer.h @@ -31,7 +31,7 @@ class LCIOStorer : public TObject, public EventStoreObserver { // LCIO -> lcfiplus // register basic collections to EventStore - void InitMCPPFOCollections(const char* pfoColName, const char* mcColName, const char* mcpfoColName); + void InitMCPPFOCollections(const char* pfoColName, const char* mcColName, const char* mcpfoColName, const char *mctrkColName = ""); void InitPFOCollections(const char* pfoColName); // initialize misc collections @@ -112,6 +112,7 @@ class LCIOStorer : public TObject, public EventStoreObserver { map *> _importMCCSCols; map *, vector *> >_importPFOCols; map > _importMCPFOLinkCols; // Link, + map > _importMCTrkLinkCols; // Link (trk), map *> _importVertexCols; map *> _importJetCols; diff --git a/include/LcfiplusProcessor.h b/include/LcfiplusProcessor.h index 17b615b..d7bcc67 100644 --- a/include/LcfiplusProcessor.h +++ b/include/LcfiplusProcessor.h @@ -66,6 +66,7 @@ class LcfiplusProcessor : public Processor, public lcfiplus::EventStoreObserver std::string _pfoCollectionName; std::string _mcpCollectionName; std::string _mcpfoRelationName; + std::string _mctrkRelationName; std::vector _algonames; std::vector _algos; diff --git a/include/LinkDef.h b/include/LinkDef.h index 51502bc..05b6828 100644 --- a/include/LinkDef.h +++ b/include/LinkDef.h @@ -46,8 +46,14 @@ #pragma link C++ class lcfiplus::ZHHAlgo; #pragma link C++ class lcfiplus::TestAlgo; #pragma link C++ class lcfiplus::FlavtagReader; +#pragma link C++ class lcfiplus::WeaverReader; #pragma link C++ class lcfiplus::VertexAnalysis; +#pragma link C++ class lcfiplus::TrackPairTree; #pragma link C++ class lcfiplus::TestAlgoV0; +#pragma link C++ class lcfiplus::VertexFinderDNNccbar; +#pragma link C++ class lcfiplus::DNNProvider2; +#pragma link C++ class lcfiplus::MLMakeNtuple; +#pragma link C++ class lcfiplus::MLInferenceWeaver; #endif diff --git a/include/MLInferenceWeaver.h b/include/MLInferenceWeaver.h new file mode 100644 index 0000000..6b950ab --- /dev/null +++ b/include/MLInferenceWeaver.h @@ -0,0 +1,75 @@ +// MLInferenceWeaver.h + +#ifndef MLInferenceWeaver_h +#define MLInferenceWeaver_h 1 + +#include "ROOT/RVec.hxx" +#include "lcfiplus.h" + +namespace rv = ROOT::VecOps; +namespace lcfiplus { + +class WeaverInterface; + +class MLInferenceWeaver : public Algorithm { + public: + MLInferenceWeaver() {} + virtual ~MLInferenceWeaver() {} + + void init(Parameters* param); + void process(); + void end(); + + private: + void processJet(); // for jet classification mode + void processEvent(); // for event classification mode + std::string _jetCollectionName; + std::string _updateJetCollectionName; + std::string _jsonFileName; + std::string _onnxFileName; + std::string _outputParameterName; + bool _eventClassification; + std::vector _outputVariables; + rv::RVec _variables; //! + WeaverInterface* _weaver; //! + std::vector* _outputJets; //! + + struct PreprocessParams { + struct VarInfo { + VarInfo() {} + VarInfo(float imedian, + float inorm_factor, + float ireplace_inf_value, + float ilower_bound, + float iupper_bound, + float ipad) + : center(imedian), + norm_factor(inorm_factor), + replace_inf_value(ireplace_inf_value), + lower_bound(ilower_bound), + upper_bound(iupper_bound), + pad(ipad) {} + + float center{0.}; + float norm_factor{1.}; + float replace_inf_value{0.}; + float lower_bound{-5.}; + float upper_bound{5.}; + float pad{0.}; + }; + std::string name; + size_t min_length{0}, max_length{0}; + std::vector var_names; + std::unordered_map var_info_map; + VarInfo info(const std::string& name) const { return var_info_map.at(name); } + void dumpVars() const; + }; + std::unordered_map prep_info_map_; //! + void parseJSON(const string& json_filename); + + ClassDef(MLInferenceWeaver,1); +}; + +} + +#endif diff --git a/include/MLInputGenerator.h b/include/MLInputGenerator.h new file mode 100644 index 0000000..522bff0 --- /dev/null +++ b/include/MLInputGenerator.h @@ -0,0 +1,60 @@ +// MLInputGenerator.h +#ifndef mlinputgenerator_h +#define mlinputgenerator_h 1 + +#include "lcfiplus.h" +#include "TTree.h" +#include "TFile.h" +#include "TVector3.h" + +#include +#include +#include +#include + +namespace lcfiplus { + +// Type alias for the variant of calculation functions +using CalcInputVariant = std::variant< + std::function, + std::function, + std::function, + std::function, + std::function, + std::function, + std::function, + std::function +>; + +// Type alias for the map of calculation functions +using CalcInputMap = std::map; + +// Static class for ML input generation +class MLInputGenerator { +public: + // Initialize the calculation function map + static void init(); + + // Get const reference to the calculation input map + static const CalcInputMap& getCalcInput() { return _calcInput; } + + // FCC functions + static float calc_dxy(float, float, float, TVector3, TVector3, int); + static float calc_dz(float, float, float, TVector3, TVector3, int); + static float calc_sip2d(float, float, float, float); + static float calc_sip3d(float, float, float, TVector3); + static float calc_jetDist(float, float, float, TVector3, TVector3); + static void calc_thetaphi(TVector3, TVector3, float&, float&); + +private: + // Private static members + static CalcInputMap _calcInput; + static bool _initialized; + + // Delete constructor to prevent instantiation + MLInputGenerator() = delete; +}; + +} + +#endif diff --git a/include/MLMakeNtuple.h b/include/MLMakeNtuple.h new file mode 100644 index 0000000..e8dc717 --- /dev/null +++ b/include/MLMakeNtuple.h @@ -0,0 +1,61 @@ +// MakeNtuple.h + +// Relation to DNNProvider2: +// It does the same thing but cleaner implementation +// +// Relation to MakeNtuple + FTManager: +// It is a modernized implementation for input variable output + +#ifndef MLMakeNtuple_h +#define MLMakeNtuple_h 1 + +class TFile; +class TTree; + +#include "lcfiplus.h" + +namespace lcfiplus { + +class MLMakeNtuple : public Algorithm { + public: + MLMakeNtuple() {} + virtual ~MLMakeNtuple() {} + + void init(Parameters* param); + void process(); + void processEvent(); + void processEventNoJets(); + void processJets(); + void end(); + + private: + string _jetname; + TFile* _file; + TTree* _tree; + int _label; + int _labelKeep; + int _outEvent; + int _outEventNoJets; + + struct MLData { + // methods + MLData() {}; + float& newData(const string& key); + vector& newDataVec(const string& key); + + void setData(const string& key, const float& val); + void addDataVec(const string& key, const float& val); + + void resetData(); + + // private members + map _mapData; + map > _mapDataVec; + } _data; + + ClassDef(MLMakeNtuple,1); +}; + +} + +#endif diff --git a/include/ONNXRuntime.h b/include/ONNXRuntime.h new file mode 100644 index 0000000..3565357 --- /dev/null +++ b/include/ONNXRuntime.h @@ -0,0 +1,45 @@ +// Code from https://github.com/HEP-FCC/FCCAnalyses + +#ifndef ONNXRuntime_ONNXRuntime_h +#define ONNXRuntime_ONNXRuntime_h + +//#include "onnxruntime/onnxruntime_cxx_api.h" +#include "onnxruntime_cxx_api.h" + +#include +#include +#include +#include + +namespace lcfiplus { + +class ONNXRuntime { +public: + explicit ONNXRuntime(const std::string& = "", const std::vector& = {}); + virtual ~ONNXRuntime(); + + template + using Tensor = std::vector>; + + ONNXRuntime(const ONNXRuntime&) = delete; + ONNXRuntime& operator=(const ONNXRuntime&) = delete; + + const std::vector& inputNames() const { return input_names_; } + + template + Tensor run(Tensor&, const Tensor& = {}, unsigned long long = 1ull) const; + +private: + size_t variablePos(const std::string&) const; + + std::unique_ptr env_; + std::unique_ptr session_; + Ort::MemoryInfo memoryInfo_; + + std::vector input_node_strings_, output_node_strings_; + std::vector input_names_; + std::map> input_node_dims_, output_node_dims_; +}; + +} +#endif diff --git a/include/VertexFinderTearDown.h b/include/VertexFinderTearDown.h index b77cf09..e0f56fb 100644 --- a/include/VertexFinderTearDown.h +++ b/include/VertexFinderTearDown.h @@ -50,10 +50,9 @@ class VertexFinderTearDown { const Track* worstTrack = resultVertex->getWorstTrack(); if (fixedTracks && find(fixedTracks->begin(), fixedTracks->end(), worstTrack) != fixedTracks->end()) { // sort Chi2Tracks - vector > vpair; const map& mpair = resultVertex->getTracksChi2Map(); - vpair.resize(mpair.size()); - partial_sort_copy(mpair.begin(), mpair.end(), vpair.begin(), vpair.end(), SortTracksByChi2()); + vector > vpair(mpair.begin(), mpair.end()); + std::sort(vpair.begin(), vpair.end(), SortTracksByChi2()); unsigned int nworst = 1; do { diff --git a/include/WeaverInterface.h b/include/WeaverInterface.h new file mode 100644 index 0000000..1cfeb4f --- /dev/null +++ b/include/WeaverInterface.h @@ -0,0 +1,83 @@ +// Code from https://github.com/HEP-FCC/FCCAnalyses + +#ifndef ONNXRuntime_WeaverInterface_h +#define ONNXRuntime_WeaverInterface_h + +#include "ONNXRuntime.h" +#include "ROOT/RVec.hxx" + +//#define DUMP_WEAVER_INPUT 1 + +namespace rv = ROOT::VecOps; + +namespace lcfiplus { + +class WeaverInterface { +public: + using ConstituentVars = rv::RVec; + + /// Initialise an inference model from Weaver output ONNX/JSON files and + /// a list of variables to be provided for each event/jet + explicit WeaverInterface(const std::string& onnx_filename = "", + const std::string& json_filename = "", + const rv::RVec& vars = {}); + + /// Run inference given a list of jet constituents variables + rv::RVec run(const rv::RVec&); + +private: + struct PreprocessParams { + struct VarInfo { + VarInfo() {} + VarInfo(float imedian, + float inorm_factor, + float ireplace_inf_value, + float ilower_bound, + float iupper_bound, + float ipad) + : center(imedian), + norm_factor(inorm_factor), + replace_inf_value(ireplace_inf_value), + lower_bound(ilower_bound), + upper_bound(iupper_bound), + pad(ipad) {} + + float center{0.}; + float norm_factor{1.}; + float replace_inf_value{0.}; + float lower_bound{-5.}; + float upper_bound{5.}; + float pad{0.}; + }; + std::string name; + size_t min_length{0}, max_length{0}; + std::vector var_names; + std::unordered_map var_info_map; + VarInfo info(const std::string& name) const { return var_info_map.at(name); } + void dumpVars() const; + }; + std::vector center_norm_pad(const rv::RVec& input, + float center, + float scale, + size_t min_length, + size_t max_length, + float pad_value = 0, + float replace_inf_value = 0, + float min = 0, + float max = -1); + size_t variablePos(const std::string&) const; + + std::unique_ptr onnx_; + std::vector variables_names_; + ONNXRuntime::Tensor input_shapes_; + std::vector input_sizes_; + std::unordered_map prep_info_map_; + ONNXRuntime::Tensor data_; +#ifdef DUMP_WEAVER_INPUT + void writeToFile(const rv::RVec< rv::RVec >& vars); + void writeToFile(const std::vector< std::vector >& vars); +#endif +}; + +} +#endif diff --git a/include/algoEtc.h b/include/algoEtc.h index 973fa46..0ffb29d 100644 --- a/include/algoEtc.h +++ b/include/algoEtc.h @@ -22,6 +22,7 @@ bool SimpleSecMuonFinder(const Track* tr, double d0sigth, double z0sigth, double bool SimpleSecElectronFinder(const Track* tr, double d0sigth, double z0sigth, double maxpos, double emin, double minfracecal, double minecalpertrackenergy, double maxecalpertrackenergy, const Vertex* ip = 0); +void AssignJetsToMC(JetVec &jets, vector&mcs_out); } } diff --git a/include/lcfiplus.h b/include/lcfiplus.h index 38deb74..fd84ab8 100644 --- a/include/lcfiplus.h +++ b/include/lcfiplus.h @@ -350,6 +350,8 @@ class Event : public EventStore { } // utility functions for MCParticles + bool isMCExist() const { return IsExist(getDefaultMCParticles());} + const MCParticle* getMCParticle(int id) const; const MCParticle* getMCParticle(const Track* trk) const; @@ -421,6 +423,13 @@ class Track : public TLorentzVector {//, protected TrackData {//, public EventPo _charge = charge; } + double getdEdx() const { + return _dEdx; + } + void setdEdx(double dEdx) { + _dEdx = dEdx; + } + double getD0() const { return _par[tpar::d0]; } @@ -569,6 +578,7 @@ class Track : public TLorentzVector {//, protected TrackData {//, public EventPo const lcfiplus::MCParticle* _mcp; int _pdg; double _charge; + double _dEdx; // track parameter double _par[tpar::parN]; @@ -711,7 +721,7 @@ class MCParticle : public TLorentzVector { int getFlavor() const; const MCParticle* getColorString()const; - const MCColorSinglet* getColorSinglet(const vector* pcs)const; + const MCColorSinglet* getColorSinglet(const vector * pcs)const; bool isStableTrack() const; bool isStable() const; @@ -1007,6 +1017,9 @@ class Jet : public TLorentzVector { const vector& getNeutrals() const { return _neutrals; } + // get all neutrals sorted by energy + vector getNeutralsSorted() const; + const vector& getVertices() const { return _vertices; } @@ -1018,6 +1031,8 @@ class Jet : public TLorentzVector { if withoutV0 is true, skip tracks which are identified as v0 (default: false) */ vector getAllTracks(bool withoutV0=false) const; + // get all tracks sorted by energy + vector getAllTracksSorted(bool withoutV0=false) const; // methods void add(const Jet& jet); diff --git a/include/testproc.h b/include/testproc.h index 100e516..0027155 100644 --- a/include/testproc.h +++ b/include/testproc.h @@ -231,6 +231,56 @@ class TestAlgo : public Algorithm { }; +class TrackPairTree : public Algorithm { + public: + TrackPairTree() {} + virtual ~TrackPairTree() {} + + void init(Parameters* param); + void process(); + void end(); + + ClassDef(TrackPairTree,1); + + private: + int _nev; + + TFile* _file; + TTree* _trackTree; + TTree* _vtxTree; + + struct dataTrackTree{ + int nev; + int njet; + int idx; // index of the track in the event + float d0; + float phi; + float omega; + float z0; + float tanLambda; + float covMatrix[15]; // 1+2+3+4+5, triangle matrix, order d0d0, d0phi, phiphi, ..., tanltanl + float chi2; + int ndf; + }; + struct dataVtxTree{ + int nev; + int njet; + int idx; // index of the vtx in the event + int tracks[2]; // index of the tracks + float pos[3]; // 3d position of the reconstructed vertex + float covMatrix[6]; // 1+2+3, should be xx, xy, yy, xz, yz, zz + float chi2; + float prob; // probability, not independent of chi2 + int trueVtx; + }; + + dataTrackTree _dataTrack; + dataVtxTree _dataVtx; + + string _privtxname; + string _jetname; +}; + class VertexAnalysis : public Algorithm { public: VertexAnalysis() {} @@ -250,6 +300,7 @@ class VertexAnalysis : public Algorithm { string _privtxname; string _secvtxname; + string _jetname; }; class FlavtagReader : public Algorithm { @@ -277,6 +328,9 @@ class FlavtagReader : public Algorithm { string _jetname; bool _bbhh; + bool _isWeaver; + bool _is6cat; + VertexVec* _vertices; //! VertexVec* _v0vertices; //! JetVec* _jets; //! @@ -293,6 +347,27 @@ class FlavtagReader : public Algorithm { }; +class WeaverReader : public Algorithm { + public: + WeaverReader() {} + virtual ~WeaverReader() {} + + void init(Parameters* param); + void process(); + void end(); + + ClassDef(WeaverReader,1); + + private: + TNtupleD* _nt; + int _nev; + TFile* _file; + + string _jetname; + JetVec* _jets; //! + +}; + class TestAlgoV0 : public Algorithm { public: TestAlgoV0() {} diff --git a/include/vertexfinderdnn.h b/include/vertexfinderdnn.h new file mode 100644 index 0000000..bcf5a16 --- /dev/null +++ b/include/vertexfinderdnn.h @@ -0,0 +1,128 @@ +// testproc.h + +#ifndef vertexfinderdnn_h +#define vertexfinderdnn_h 1 + +#include "lcfiplus.h" +#include "LcfiplusProcessor.h" +#include "geometry.h" +#include "TNtuple.h" +#include "TNtupleD.h" +#include "TH2D.h" +#include + +namespace lcfiplus { + +class VertexFinderDNNccbar : public Algorithm { + public: + VertexFinderDNNccbar() {} + virtual ~VertexFinderDNNccbar() {} + void init(Parameters* param); + void process(); + void end(); + ClassDef(VertexFinderDNNccbar,1); + private: + TTree* _ntp; + TFile* _file; + Helix* _hel1; + Helix* _hel2; + int nEvt; + int ntr1Trk; + int ntr2Trk; + int tr1sscid = -999; + int tr1sscpdg = -999; + int tr1ssc = -999; + int tr2sscid = -999; + int tr2sscpdg = -999; + int tr2ssc = -999; + + TVector3* xyz; + //enum par : int; + //enum tpar : int; + struct TracksData { + double tr1d0 = -999; + double tr1z0 = -999; + double tr1phi = -999; + double tr1omega = -999; + double tr1tanlam = -999; + double tr1x = -999; + double tr1y = -999; + double tr1z = -999; + int tr1charge = -999; + double tr1covmatrixd0d0 = -999; + double tr1covmatrixd0z0 = -999; + double tr1covmatrixd0ph = -999; + double tr1covmatrixd0om = -999; + double tr1covmatrixd0tl = -999; + double tr1covmatrixz0z0 = -999; + double tr1covmatrixz0ph = -999; + double tr1covmatrixz0om = -999; + double tr1covmatrixz0tl = -999; + double tr1covmatrixphph = -999; + double tr1covmatrixphom = -999; + double tr1covmatrixphtl = -999; + double tr1covmatrixomom = -999; + double tr1covmatrixomtl = -999; + double tr1covmatrixtltl = -999; + double tr2d0 = -999; + double tr2z0 = -999; + double tr2phi = -999; + double tr2omega = -999; + double tr2tanlam = -999; + double tr2x = -999; + double tr2y = -999; + double tr2z = -999; + int tr2charge = -999; + double tr2covmatrixd0d0 = -999; + double tr2covmatrixd0z0 = -999; + double tr2covmatrixd0ph = -999; + double tr2covmatrixd0om = -999; + double tr2covmatrixd0tl = -999; + double tr2covmatrixz0z0 = -999; + double tr2covmatrixz0ph = -999; + double tr2covmatrixz0om = -999; + double tr2covmatrixz0tl = -999; + double tr2covmatrixphph = -999; + double tr2covmatrixphom = -999; + double tr2covmatrixphtl = -999; + double tr2covmatrixomom = -999; + double tr2covmatrixomtl = -999; + double tr2covmatrixtltl = -999; + double tr1mcx = -999; + double tr1mcy = -999; + double tr1mcz = -999; + int tr1id = -999; + int tr1pdg = -999; + double tr2mcx = -999; + double tr2mcy = -999; + double tr2mcz = -999; + int tr2id = -999; + int tr2pdg = -999; + int cosemistable = -999; + int cosemistablec = -999; + int connect = -999; + double tr1tlvx = -999; + double tr1tlvy = -999; + double tr1tlvz = -999; + double tr2tlvx = -999; + double tr2tlvy = -999; + double tr2tlvz = -999; + double mass = -999; + double l0mass = -999; + double minE = -999; + double chi2 = -999; + double vchi2 = -999; + double vprob = -999; + double vposx = -999; + double vposy = -999; + double vposz = -999; + int lcfiplustag = -999; + double tr1xyz[10][3]; + double tr2xyz[10][3]; + }; + TracksData _data; +}; + +} + +#endif diff --git a/memo.md b/memo.md new file mode 100644 index 0000000..f98bca0 --- /dev/null +++ b/memo.md @@ -0,0 +1,84 @@ +# Installing additional dependencies for LCFIPlus + +## JSON parser by nlohmann +This JSON parser is used by the weaver interface code taken from LCCAnalyses. + +``` +git clone https://github.com/nlohmann/json.git nlohmann_json +``` + +## ONNX Runtime C/C++ library + +The ONNX Runtime library is needed to perform the inference using the trained weights output from weaver. +The following instructions were adapted from [this link](https://stackoverflow.com/questions/63420533/setting-up-onnx-runtime-on-ubuntu-20-04-c-api). + +### Download and install library +Here we assume local home directory installation. +``` +mkdir /tmp/onnxInstall +cd /tmp/onnxInstall +wget -O onnx_archive.nupkg https://www.nuget.org/api/v2/package/Microsoft.ML.OnnxRuntime/1.19.2 +unzip onnx_archive.nupkg +cp runtimes/linux-x64/native/libonnxruntime.so ~/.local/lib/ +mkdir -p ~/.local/include/onnxruntime +cp -r build/native/include/*.h ~/.local/include/onnxruntime/ +``` + +### Create CMake files + +Additional cmake files are needed for cmake to find onnxruntime as a package. + +``` +mkdir -p ~/.local/share/cmake/onnxruntime +cd ~/.local/share/cmake/onnxruntime +cat <<'EOF' > onnxruntime-config.cmake +# Custom cmake config file by jcarius to enable find_package(onnxruntime) without modifying LIBRARY_PATH and LD_LIBRARY_PATH +# +# This will define the following variables: +# onnxruntime_FOUND -- True if the system has the onnxruntime library +# onnxruntime_INCLUDE_DIRS -- The include directories for onnxruntime +# onnxruntime_LIBRARIES -- Libraries to link against +# onnxruntime_CXX_FLAGS -- Additional (required) compiler flags + +include(FindPackageHandleStandardArgs) + +# Assume we are in /share/cmake/onnxruntime/onnxruntimeConfig.cmake +get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH) +get_filename_component(onnxruntime_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE) + +set(onnxruntime_INCLUDE_DIRS ${onnxruntime_INSTALL_PREFIX}/include) +set(onnxruntime_LIBRARIES onnxruntime) +set(onnxruntime_CXX_FLAGS "") # no flags needed + + +find_library(onnxruntime_LIBRARY onnxruntime + PATHS "${onnxruntime_INSTALL_PREFIX}/lib" +) + +add_library(onnxruntime SHARED IMPORTED) +set_property(TARGET onnxruntime PROPERTY IMPORTED_LOCATION "${onnxruntime_LIBRARY}") +set_property(TARGET onnxruntime PROPERTY INTERFACE_INCLUDE_DIRECTORIES "${onnxruntime_INCLUDE_DIRS}") +set_property(TARGET onnxruntime PROPERTY INTERFACE_COMPILE_OPTIONS "${onnxruntime_CXX_FLAGS}") + +find_package_handle_standard_args(onnxruntime DEFAULT_MSG onnxruntime_LIBRARY onnxruntime_INCLUDE_DIRS) +EOF +cat <<'EOF' > onnxruntime-version.cmake +# Custom cmake version file by jcarius + +set(PACKAGE_VERSION "1.19.2") + +# Check whether the requested PACKAGE_FIND_VERSION is compatible +if("${PACKAGE_VERSION}" VERSION_LESS "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_COMPATIBLE FALSE) +else() + set(PACKAGE_VERSION_COMPATIBLE TRUE) + if("${PACKAGE_VERSION}" VERSION_EQUAL "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_EXACT TRUE) + endif() +endif() +``` + +## Add object file to MARLIN_DLL +``` +export MARLIN_DLL=$MARLIN_DLL:$HOME/.local/lib/libonnxruntime.so +``` diff --git a/src/DNNProvider2.cc b/src/DNNProvider2.cc new file mode 100644 index 0000000..906dbe9 --- /dev/null +++ b/src/DNNProvider2.cc @@ -0,0 +1,904 @@ +#include + +#include "TFile.h" +#include "TNtuple.h" +#include "TNtupleD.h" +#include "TSystem.h" +#include "TPad.h" +#include "TStyle.h" + +#include "lcfiplus.h" +#include "process.h" +#include "DNNProvider2.h" +#include "VertexSelector.h" +#include "algoEtc.h" +#include "VertexFinderSuehara.h" +#include "VertexFitterSimple.h" + +#include +#include + +using namespace lcfiplus; + +namespace lcfiplus { + +void DNNProvider2::init(Parameters* param) { + Algorithm::init(param); + string filename = param->get("DNNProvider2.FileName",string("DNNProvider2.root")); + _jetname = param->get("DNNProvider2.JetCollectionName",string("RefinedJets")); + string privtx = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + Event::Instance()->setDefaultPrimaryVertex(privtx.c_str()); + + _mcIsB = param->get("DNNProvider2.MC_IsB",int(0)); + _mcIsC = param->get("DNNProvider2.MC_IsC",int(0)); + _mcIsU = param->get("DNNProvider2.MC_IsU",int(0)); + _mcIsD = param->get("DNNProvider2.MC_IsD",int(0)); + _mcIsS = param->get("DNNProvider2.MC_IsS",int(0)); + _mcIsG = param->get("DNNProvider2.MC_IsG",int(0)); + + _jets = 0; + + if(_mcIsB + _mcIsC + _mcIsU + _mcIsD + _mcIsS + _mcIsG != 1){ + cout << "Label parameter is wrong! Put 1 on one of MC_Is(BCQ)." << endl; + } + + cout << filename << endl; + + _file = new TFile(filename.c_str(),"RECREATE"); + _ntp = new TTree("tree","tree"); + + DNNData & d = _data; + // jet variables + _ntp->Branch("jet_px",&d.jet_px,"jet_px/F"); + _ntp->Branch("jet_py",&d.jet_py,"jet_py/F"); + _ntp->Branch("jet_pz",&d.jet_pz,"jet_pz/F"); + _ntp->Branch("jet_mass",&d.jet_mass,"jet_mass/F"); + _ntp->Branch("jet_ntracks",&d.jet_ntracks,"jet_ntracks/I"); + _ntp->Branch("jet_nneutrals",&d.jet_nneutrals,"jet_nneutrals/I"); + + _ntp->Branch("jet_phi",&d.jet_phi,"jet_phi/F"); + _ntp->Branch("jet_theta",&d.jet_theta,"jet_theta/F"); + + // for charges + // particle kinematics + _ntp->Branch("pfcand_px",&d.px); + _ntp->Branch("pfcand_py",&d.py); + _ntp->Branch("pfcand_pz",&d.pz); + _ntp->Branch("pfcand_e",&d.e); + _ntp->Branch("pfcand_efrac",&d.efrac); + _ntp->Branch("pfcand_erel_log",&d.erel_log); + _ntp->Branch("pfcand_thetarel",&d.dtheta); + _ntp->Branch("pfcand_phirel",&d.dphi); + _ntp->Branch("pfcand_thetarel_ilc",&d.dtheta_ilc); + _ntp->Branch("pfcand_phirel_ilc",&d.dphi_ilc); + + // track errors + _ntp->Branch("pfcand_dptdpt",&d.cov_omega); + _ntp->Branch("pfcand_detadeta",&d.cov_tanlambda); + _ntp->Branch("pfcand_dphidphi",&d.cov_phi); + _ntp->Branch("pfcand_dxydxy",&d.cov_d0); + _ntp->Branch("pfcand_dzdz",&d.cov_z0); + _ntp->Branch("pfcand_dxydz",&d.cov_d0_z0); + _ntp->Branch("pfcand_dphidxy",&d.cov_d0_phi); + _ntp->Branch("pfcand_dlambdadz",&d.cov_z0_tanlambda); + _ntp->Branch("pfcand_dxyc",&d.cov_d0_omega); + _ntp->Branch("pfcand_dxyctgtheta",&d.cov_d0_tanlambda); + _ntp->Branch("pfcand_phic",&d.cov_phi_omega); + _ntp->Branch("pfcand_phidz",&d.cov_z0_phi); + _ntp->Branch("pfcand_phictgtheta",&d.cov_phi_tanlambda); + _ntp->Branch("pfcand_cdz",&d.cov_z0_omega); + _ntp->Branch("pfcand_cctgtheta",&d.cov_omega_tanlambda); + + // particle displacements + _ntp->Branch("d0",&d.d0); + _ntp->Branch("d0sig",&d.d0sig); + _ntp->Branch("z0",&d.z0); + _ntp->Branch("z0sig",&d.z0sig); + _ntp->Branch("ip3d",&d.ip3d); + _ntp->Branch("ip3dsig",&d.ip3dsig); + + _ntp->Branch("dEdx",&d.dEdx); + + _ntp->Branch("pfcand_dxy",&d.dxy); + _ntp->Branch("pfcand_dz",&d.dz); + _ntp->Branch("pfcand_btagSip2dVal",&d.ip2d_fcc); + _ntp->Branch("pfcand_btagSip2dSig",&d.ip2dsig_fcc); + _ntp->Branch("pfcand_btagSip3dVal",&d.ip3d_fcc); + _ntp->Branch("pfcand_btagSip3dSig",&d.ip3dsig_fcc); + _ntp->Branch("pfcand_btagJetDistVal",&d.jetdist_fcc); + _ntp->Branch("pfcand_btagJetDistSig",&d.jetdistsig_fcc); + + // particle ID + _ntp->Branch("pfcand_charge",&d.charge); + _ntp->Branch("pfcand_isMu",&d.ismuon); + _ntp->Branch("pfcand_isEl",&d.iselectron); + _ntp->Branch("pfcand_isGamma",&d.isphoton); + _ntp->Branch("pfcand_isChargedHad",&d.ischargedhadron); + _ntp->Branch("pfcand_isNeutralHad",&d.isneutralhadron); + _ntp->Branch("pfcand_type",&d.pdg_pfa); + _ntp->Branch("pfcand_mcpid",&d.mcpid); + _ntp->Branch("pfcand_mcp_pdg",&d.mcp_pdg); + _ntp->Branch("pfcand_Ktype",&d.K_pdg_pfa); //20240201 + _ntp->Branch("pfcand_isPion",&d.ispion); + _ntp->Branch("pfcand_isKaon",&d.iskaon); + _ntp->Branch("pfcand_isProton",&d.isproton); + // _ntp->Branch("pfcand_isKaon0", &d.iskaon0); + + // for PI3 (20240203) + _ntp->Branch("pfcand_proton_K",&d.proton_K); + _ntp->Branch("pfcand_pion_K",&d.pion_K); + _ntp->Branch("pfcand_proton_Klike",&d.proton_Klike); + _ntp->Branch("pfcand_pion_Klike",&d.pion_Klike); + _ntp->Branch("pfcand_dEdxEl",&d.electron_dEdxdistance); + _ntp->Branch("pfcand_dEdxMu",&d.muon_dEdxdistance); + _ntp->Branch("pfcand_dEdxPion",&d.pion_dEdxdistance); + _ntp->Branch("pfcand_dEdxKaon",&d.kaon_dEdxdistance); + _ntp->Branch("pfcand_dEdxProton",&d.proton_dEdxdistance); + _ntp->Branch("pfcand_iselectronlike",&d.iselectronlike); + _ntp->Branch("pfcand_ismuonlike",&d.ismuonlike); + _ntp->Branch("pfcand_ispionlike",&d.ispionlike); + _ntp->Branch("pfcand_iskaonlike",&d.iskaonlike); + _ntp->Branch("pfcand_isprotonlike",&d.isprotonlike); + + _ntp->Branch("pfcand_iselectronprob",&d.iselectronprob); + _ntp->Branch("pfcand_ismuonprob",&d.ismuonprob); + _ntp->Branch("pfcand_ispionprob",&d.ispionprob); + _ntp->Branch("pfcand_iskaonprob",&d.iskaonprob); + _ntp->Branch("pfcand_isprotonprob",&d.isprotonprob); + + + + // for neutrals + // particle kinematics + _ntp->Branch("neu_pfcand_px",&d.neu_px); + _ntp->Branch("neu_pfcand_py",&d.neu_py); + _ntp->Branch("neu_pfcand_pz",&d.neu_pz); + _ntp->Branch("neu_pfcand_e",&d.neu_e); + _ntp->Branch("neu_pfcand_efrac",&d.neu_efrac); + _ntp->Branch("neu_pfcand_erel_log",&d.neu_erel_log); + _ntp->Branch("neu_pfcand_thetarel",&d.neu_dtheta); + _ntp->Branch("neu_pfcand_phirel",&d.neu_dphi); + _ntp->Branch("neu_pfcand_thetarel_ilc",&d.neu_dtheta_ilc); + _ntp->Branch("neu_pfcand_phirel_ilc",&d.neu_dphi_ilc); + + // track errors + _ntp->Branch("neu_pfcand_dptdpt",&d.neu_cov_omega); + _ntp->Branch("neu_pfcand_detadeta",&d.neu_cov_tanlambda); + _ntp->Branch("neu_pfcand_dphidphi",&d.neu_cov_phi); + _ntp->Branch("neu_pfcand_dxydxy",&d.neu_cov_d0); + _ntp->Branch("neu_pfcand_dzdz",&d.neu_cov_z0); + _ntp->Branch("neu_pfcand_dxydz",&d.neu_cov_d0_z0); + _ntp->Branch("neu_pfcand_dphidxy",&d.neu_cov_d0_phi); + _ntp->Branch("neu_pfcand_dlambdadz",&d.neu_cov_z0_tanlambda); + _ntp->Branch("neu_pfcand_dxyc",&d.neu_cov_d0_omega); + _ntp->Branch("neu_pfcand_dxyctgtheta",&d.neu_cov_d0_tanlambda); + _ntp->Branch("neu_pfcand_phic",&d.neu_cov_phi_omega); + _ntp->Branch("neu_pfcand_phidz",&d.neu_cov_z0_phi); + _ntp->Branch("neu_pfcand_phictgtheta",&d.neu_cov_phi_tanlambda); + _ntp->Branch("neu_pfcand_cdz",&d.neu_cov_z0_omega); + _ntp->Branch("neu_pfcand_cctgtheta",&d.neu_cov_omega_tanlambda); + + // particle displacements + _ntp->Branch("neu_d0",&d.neu_d0); + _ntp->Branch("neu_d0sig",&d.neu_d0sig); + _ntp->Branch("neu_z0",&d.neu_z0); + _ntp->Branch("neu_z0sig",&d.neu_z0sig); + _ntp->Branch("neu_ip3d",&d.neu_ip3d); + _ntp->Branch("neu_ip3dsig",&d.neu_ip3dsig); + + _ntp->Branch("neu_pfcand_dxy",&d.neu_dxy); + _ntp->Branch("neu_pfcand_dz",&d.neu_dz); + _ntp->Branch("neu_pfcand_btagSip2dVal",&d.neu_ip2d_fcc); + _ntp->Branch("neu_pfcand_btagSip2dSig",&d.neu_ip2dsig_fcc); + _ntp->Branch("neu_pfcand_btagSip3dVal",&d.neu_ip3d_fcc); + _ntp->Branch("neu_pfcand_btagSip3dSig",&d.neu_ip3dsig_fcc); + _ntp->Branch("neu_pfcand_btagJetDistVal",&d.neu_jetdist_fcc); + _ntp->Branch("neu_pfcand_btagJetDistSig",&d.neu_jetdistsig_fcc); + + // particle ID + _ntp->Branch("neu_pfcand_charge",&d.neu_charge); + _ntp->Branch("neu_pfcand_isMu",&d.neu_ismuon); + _ntp->Branch("neu_pfcand_isEl",&d.neu_iselectron); + _ntp->Branch("neu_pfcand_isGamma",&d.neu_isphoton); + _ntp->Branch("neu_pfcand_isChargedHad",&d.neu_ischargedhadron); + _ntp->Branch("neu_pfcand_isNeutralHad",&d.neu_isneutralhadron); + _ntp->Branch("neu_pfcand_type",&d.neu_pdg_pfa); + _ntp->Branch("neu_pfcand_mcpid",&d.neu_mcpid); + _ntp->Branch("neu_pfcand_mcp_pdg",&d.neu_mcp_pdg); + _ntp->Branch("neu_pfcand_Ktype",&d.neu_K_pdg_pfa); //20240201 + _ntp->Branch("neu_pfcand_isPion",&d.neu_ispion); + _ntp->Branch("neu_pfcand_isKaon",&d.neu_iskaon); + _ntp->Branch("neu_pfcand_isProton",&d.neu_isproton); + // _ntp->Branch("neu_pfcand_isKaon0", &d.neu_iskaon0); + + + // label + _ntp->Branch("mc_b",&d.mc_b,"mc_b/I"); + _ntp->Branch("mc_c",&d.mc_c,"mc_c/I"); + _ntp->Branch("mc_u",&d.mc_u,"mc_u/I"); + _ntp->Branch("mc_d",&d.mc_d,"mc_d/I"); + _ntp->Branch("mc_s",&d.mc_s,"mc_s/I"); + _ntp->Branch("mc_g",&d.mc_g,"mc_g/I"); + _ntp->Branch("mc_q",&d.mc_q,"mc_q/I"); // usdg + +} + +void DNNProvider2::process() { + if (!_jets) { + Event::Instance()->Get(_jetname.c_str(), _jets); + } + + const JetVec& jets = *_jets; + const Vertex *privtx = Event::Instance()->getPrimaryVertex(); + + DNNData &d = _data; + + for (unsigned int j=0; j < jets.size(); ++j) { + const Jet* jet = jets[j]; + + + + _data = DNNData(); + + d.jet_px = jet->Px(); + d.jet_py = jet->Py(); + d.jet_pz = jet->Pz(); + d.jet_e = jet->E(); + d.jet_mass = jet->M(); + TrackVec &tracks = jet->getAllTracks(); + NeutralVec &neutrals = jet->getNeutrals(); + // MCParticleVec &mcps = jet->getMCParticle(); + + d.jet_theta = jet->Theta(); + d.jet_phi = jet->Phi(); + + d.jet_ntracks = tracks.size(); + d.jet_nneutrals = neutrals.size(); + + //float jet_theta = jet->Theta(); + //float jet_phi = jet->Phi(); + + d.mc_b = _mcIsB; + d.mc_c = _mcIsC; + d.mc_u = _mcIsU; + d.mc_d = _mcIsD; + d.mc_s = _mcIsS; + d.mc_g = _mcIsG; + d.mc_q = _mcIsU || _mcIsD || _mcIsS || _mcIsG; + + // probably order of tracks/netural does not matter... + //int nall = d.jet_ntracks + d.jet_nneutrals; + int ntra = d.jet_ntracks; + int nneu = d.jet_nneutrals; + if (ntra==0) continue; + if (nneu==0) continue; + //for charges + d.px.resize(ntra); + d.py.resize(ntra); + d.pz.resize(ntra); + d.e.resize(ntra); + d.efrac.resize(ntra); + d.erel_log.resize(ntra); + d.dtheta.resize(ntra); + d.dphi.resize(ntra); + d.dtheta_ilc.resize(ntra); + d.dphi_ilc.resize(ntra); + + d.cov_d0.resize(ntra); + d.cov_z0.resize(ntra); + d.cov_phi.resize(ntra); + d.cov_omega.resize(ntra); + d.cov_tanlambda.resize(ntra); + + d.cov_d0_z0.resize(ntra); + d.cov_d0_phi.resize(ntra); + d.cov_d0_omega.resize(ntra); + d.cov_d0_tanlambda.resize(ntra); + + d.cov_z0_phi.resize(ntra); + d.cov_z0_omega.resize(ntra); + d.cov_z0_tanlambda.resize(ntra); + + d.cov_phi_omega.resize(ntra); + d.cov_phi_tanlambda.resize(ntra); + d.cov_omega_tanlambda.resize(ntra); + + d.d0.resize(ntra); + d.d0sig.resize(ntra); + d.z0.resize(ntra); + d.z0sig.resize(ntra); + d.ip3d.resize(ntra); + d.ip3dsig.resize(ntra); + + d.dEdx.resize(ntra); + + d.dxy.resize(ntra); + d.dz.resize(ntra); + d.ip2d_fcc.resize(ntra); + d.ip2dsig_fcc.resize(ntra); + d.ip3d_fcc.resize(ntra); + d.ip3dsig_fcc.resize(ntra); + d.jetdist_fcc.resize(ntra); + d.jetdistsig_fcc.resize(ntra); + + d.charge.resize(ntra); + d.ismuon.resize(ntra); + d.iselectron.resize(ntra); + d.isphoton.resize(ntra); + d.ischargedhadron.resize(ntra); + d.isneutralhadron.resize(ntra); + d.ispion.resize(ntra); + d.iskaon.resize(ntra); + // d.iskaon0.resize(ntra); + d.isproton.resize(ntra); + d.pdg_pfa.resize(ntra); + d.mcpid.resize(ntra); + d.mcp_pdg.resize(ntra); + d.K_pdg_pfa.resize(ntra); + + d.proton_K.resize(ntra); + d.pion_K.resize(ntra); + d.proton_Klike.resize(ntra); + d.pion_Klike.resize(ntra); + + d.electron_dEdxdistance.resize(ntra); + d.muon_dEdxdistance.resize(ntra); + d.pion_dEdxdistance.resize(ntra); + d.kaon_dEdxdistance.resize(ntra); + d.proton_dEdxdistance.resize(ntra); + + d.iselectronlike.resize(ntra); + d.ismuonlike.resize(ntra); + d.ispionlike.resize(ntra); + d.iskaonlike.resize(ntra); + d.isprotonlike.resize(ntra); + + d.iselectronprob.resize(ntra); + d.ismuonprob.resize(ntra); + d.ispionprob.resize(ntra); + d.iskaonprob.resize(ntra); + d.isprotonprob.resize(ntra); + + + // for neutrals + d.neu_px.resize(nneu); + d.neu_py.resize(nneu); + d.neu_pz.resize(nneu); + d.neu_e.resize(nneu); + d.neu_efrac.resize(nneu); + d.neu_erel_log.resize(nneu); + d.neu_dtheta.resize(nneu); + d.neu_dphi.resize(nneu); + d.neu_dtheta_ilc.resize(nneu); + d.neu_dphi_ilc.resize(nneu); + + d.neu_cov_d0.resize(nneu); + d.neu_cov_z0.resize(nneu); + d.neu_cov_phi.resize(nneu); + d.neu_cov_omega.resize(nneu); + d.neu_cov_tanlambda.resize(nneu); + + d.neu_cov_d0_z0.resize(nneu); + d.neu_cov_d0_phi.resize(nneu); + d.neu_cov_d0_omega.resize(nneu); + d.neu_cov_d0_tanlambda.resize(nneu); + + d.neu_cov_z0_phi.resize(nneu); + d.neu_cov_z0_omega.resize(nneu); + d.neu_cov_z0_tanlambda.resize(nneu); + + d.neu_cov_phi_omega.resize(nneu); + d.neu_cov_phi_tanlambda.resize(nneu); + d.neu_cov_omega_tanlambda.resize(nneu); + + d.neu_d0.resize(nneu); + d.neu_d0sig.resize(nneu); + d.neu_z0.resize(nneu); + d.neu_z0sig.resize(nneu); + d.neu_ip3d.resize(nneu); + d.neu_ip3dsig.resize(nneu); + + d.neu_dxy.resize(nneu); + d.neu_dz.resize(nneu); + d.neu_ip2d_fcc.resize(nneu); + d.neu_ip2dsig_fcc.resize(nneu); + d.neu_ip3d_fcc.resize(nneu); + d.neu_ip3dsig_fcc.resize(nneu); + d.neu_jetdist_fcc.resize(nneu); + d.neu_jetdistsig_fcc.resize(nneu); + + d.neu_charge.resize(nneu); + d.neu_ismuon.resize(nneu); + d.neu_iselectron.resize(nneu); + d.neu_isphoton.resize(nneu); + d.neu_ischargedhadron.resize(nneu); + d.neu_isneutralhadron.resize(nneu); + d.neu_ispion.resize(nneu); + d.neu_iskaon.resize(nneu); + // d.neu_iskaon0.resize(nneu); + d.neu_isproton.resize(nneu); + d.neu_pdg_pfa.resize(nneu); + d.neu_mcpid.resize(nneu); + d.neu_mcp_pdg.resize(nneu); + d.neu_K_pdg_pfa.resize(nneu); + + + vector > order_tr; + order_tr.resize(ntra); + vector > order_n; + order_n.resize(nneu); + + int i; + + for(i=0;i(tr->E(), i); + } + for(i=0;i(neu->E(), i); + } + + // sort by energy + // std::sort(order.begin(), order.end(), [](std::paira, std::pair b){ + // return a.first > b.first; + // }); + + + + std::sort(order_tr.begin(), order_tr.end(), [](std::paira, std::pair b){ + return a.first > b.first; + }); + std::sort(order_n.begin(), order_n.end(), [](std::paira, std::pair b){ + return a.first > b.first; + }); + + + + for(i=0;i= d.jet_ntracks) continue; + //cout << i << " " << order[i].second << " " << d.jet_ntracks << " " << nall << endl; + if(ntra==0){ + cout << "passedT" << endl; + continue; + } + + const Track *tr = tracks[order_tr[i].second]; + // const Track *tr = tracks[i]; + d.px[i] = tr->Px(); + d.py[i] = tr->Py(); + d.pz[i] = tr->Pz(); + d.e[i] = tr->E(); + d.efrac[i] = tr->E() / jet->E(); + d.erel_log[i] = log10(d.efrac[i]); + d.dtheta_ilc[i] = tr->Theta() - jet->Theta(); + d.dphi_ilc[i] = tr->Phi() - jet->Phi(); + if(d.dphi_ilc[i] < -TMath::Pi())d.dphi_ilc[i] += TMath::Pi() * 2; + if(d.dphi_ilc[i] > TMath::Pi())d.dphi_ilc[i] -= TMath::Pi() * 2; + calc_thetaphi(jet->Vect(), tr->Vect(), d.dtheta[i], d.dphi[i]); + + // track covmatrix + d.cov_d0[i] = tr->getCovMatrix()[tpar::d0d0]; + d.cov_z0[i] = tr->getCovMatrix()[tpar::z0z0]; + d.cov_phi[i] = tr->getCovMatrix()[tpar::phph]; + d.cov_omega[i] = tr->getCovMatrix()[tpar::omom]; + d.cov_tanlambda[i] = tr->getCovMatrix()[tpar::tdtd]; + + d.cov_d0_z0[i] = tr->getCovMatrix()[tpar::d0z0]; + d.cov_d0_phi[i] = tr->getCovMatrix()[tpar::d0ph]; + d.cov_d0_omega[i] = tr->getCovMatrix()[tpar::d0om]; + d.cov_d0_tanlambda[i] = tr->getCovMatrix()[tpar::d0td]; + + d.cov_z0_phi[i] = tr->getCovMatrix()[tpar::z0ph]; + d.cov_z0_omega[i] = tr->getCovMatrix()[tpar::z0om]; + d.cov_z0_tanlambda[i] = tr->getCovMatrix()[tpar::z0td]; + + d.cov_phi_omega[i] = tr->getCovMatrix()[tpar::phom]; + d.cov_phi_tanlambda[i] = tr->getCovMatrix()[tpar::phtd]; + d.cov_omega_tanlambda[i] = tr->getCovMatrix()[tpar::omtd]; + + d.d0[i] = tr->getD0(); + d.d0sig[i] = tr->getD0() / sqrt(tr->getCovMatrix()[tpar::cov::d0d0]); + d.z0[i] = tr->getZ0(); + d.z0sig[i] = tr->getZ0() / sqrt(tr->getCovMatrix()[tpar::cov::z0z0]); + + d.ip3d[i] = sqrt(tr->getD0() * tr->getD0() + tr->getZ0() * tr->getZ0()); + d.ip3dsig[i] = d.ip3d[i] / sqrt(tr->getCovMatrix()[tpar::cov::d0d0] + tr->getCovMatrix()[tpar::cov::z0z0] + 2 * tr->getCovMatrix()[tpar::cov::d0z0]); + + d.dEdx[i] = tr->getdEdx(); + + d.dxy[i] = calc_dxy(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), privtx->getPos(), tr->getCharge()); + d.dz[i] = calc_dz(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), privtx->getPos(), tr->getCharge()); + d.ip2d_fcc[i] = calc_sip2d(tr->getD0(), tr->getPhi(), jet->Px(), jet->Py()); + d.ip2dsig_fcc[i] = d.ip2d_fcc[i] / sqrt(d.cov_d0[i]); + d.ip3d_fcc[i] = calc_sip3d(tr->getD0(), tr->getZ0(), tr->getPhi(), jet->Vect()); + d.ip3dsig_fcc[i] = d.ip3d_fcc[i] / sqrt(d.cov_d0[i] + d.cov_z0[i]); + d.jetdist_fcc[i] = calc_jetDist(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), jet->Vect()); + d.jetdistsig_fcc[i] = d.jetdist_fcc[i] / sqrt(d.cov_d0[i] + d.cov_z0[i]); + + + d.charge[i] = tr->getCharge(); + + // // get ParticleID 0 or 1 + // d.isneutralhadron[i] = 0.0; + // std::vector charged(5); + // charged.clear(); + // d.ismuon[i] = 0; + // d.iselectron[i] = 0; + // d.iskaon[i] = 0; + // d.ispion[i] = 0; + // d.isproton[i] = 0; + // d.isphoton[i] = 0; + // charged.push_back(tr->getParticleIDProbability("muonProbability")); + // charged.push_back (tr->getParticleIDProbability("electronProbability")); + // charged.push_back(tr->getParticleIDProbability("kaonProbability")); + // charged.push_back(tr->getParticleIDProbability("pionProbability")); + // charged.push_back(tr->getParticleIDProbability("protonProbability")); + + // //std::vector::iterator iter = std::max_element(charged.begin(), charged.end()); + // auto iter = std::max_element(charged.begin(), charged.end()); + // int index = std::distance(charged.begin(), iter); + // if (index==0) d.ismuon[i] = 1; + // if (index==1) d.iselectron[i] = 1; + // if (index==2) d.iskaon[i] = 1; + // if (index==3) d.ispion[i] = 1; + // if (index==4) d.isproton[i] = 1; + + // pdg etc + d.pdg_pfa[i] = tr->getPDG(); + if (tr->getMcp()) d.mcp_pdg[i] = (tr->getMcp())->getPDG(); + else d.mcp_pdg[i] = 0; + d.K_pdg_pfa[i] = 0; + if (d.pdg_pfa[i]==321) d.K_pdg_pfa[i] = 1; + if (d.pdg_pfa[i]==310) d.K_pdg_pfa[i] = 1; + + // new PID 2024.06.13 + d.isneutralhadron[i] = 0.0; + d.ismuonprob[i] = tr->getParticleIDProbability("muonProbability"); + d.iselectronprob[i] = tr->getParticleIDProbability("electronProbability"); + d.iskaonprob[i] = tr->getParticleIDProbability("kaonProbability"); + d.ispionprob[i] = tr->getParticleIDProbability("pionProbability"); + d.isprotonprob[i] = tr->getParticleIDProbability("protonProbability"); + + // dEdx check !! if not test, need to comment out !! + float pfcandp = pow(d.px[i]*d.px[i]+d.py[i]*d.py[i]+d.pz[i]*d.pz[i],0.5); + if ((d.pdg_pfa[i]==321) && (pfcandp<0.4) && (0.0000001 charged(5); + charged.clear(); + d.ismuon[i] = 0; + d.iselectron[i] = 0; + d.iskaon[i] = 0; + d.ispion[i] = 0; + d.isproton[i] = 0; + d.isphoton[i] = 0; + charged.push_back(d.ismuonprob[i]); + charged.push_back(d.iselectronprob[i]); + charged.push_back(d.iskaonprob[i]); + charged.push_back(d.ispionprob[i]); + charged.push_back(d.isprotonprob[i]); + auto iter = std::max_element(charged.begin(), charged.end()); + int index = std::distance(charged.begin(), iter); + if ((d.ismuonprob[i]==0) && (d.iselectronprob[i]==0) && (d.iskaonprob[i]==0) && (d.ispionprob[i]==0) && (d.isprotonprob[i]==0)) { + continue; + } + else if (index==0) d.ismuon[i] = 1; + else if (index==1) d.iselectron[i] = 1; + else if (index==2) d.iskaon[i] = 1; + else if (index==3) d.ispion[i] = 1; + else if (index==4) d.isproton[i] = 1; + + + + + + d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + + //d.mcpid[i] = (tr->getMcp())->getId(); + // if (tr->getMcp()) d.mcp_pdg[i] = (tr->getMcp())->getPDG(); + // else d.mcp_pdg[i] = 0; + // d.K_pdg_pfa[i] = 0; + // if (d.pdg_pfa[i]==321) d.K_pdg_pfa[i] = 1; + // if (d.pdg_pfa[i]==310) d.K_pdg_pfa[i] = 1; + + d.pion_K[i] = d.ispion[i] - d.iskaon[i]; + d.proton_K[i] = d.isproton[i] - d.iskaon[i]; + // d.pion_Klike[i] = tr->getParticleIDProbability("pionLikelihood") - tr->getParticleIDProbability("kaonLikelihood"); + // d.proton_Klike[i] = tr->getParticleIDProbability("protonLikelihood") - tr->getParticleIDProbability("kaonLikelihood"); + + + // d.electron_dEdxdistance[i] = tr->getParticleIDProbability("electron_dEdxdistance"); + // d.muon_dEdxdistance[i] = tr->getParticleIDProbability("muon_dEdxdist ance"); + // d.pion_dEdxdistance[i] = tr->getParticleIDProbability("pion_dEdxdistance"); + // d.kaon_dEdxdistance[i] = tr->getParticleIDProbability("kaon_dEdxdistance"); + // d.proton_dEdxdistance[i] = tr->getParticleIDProbability("proton_dEdxdistance"); + + // d.iselectronlike[i] = tr->getParticleIDProbability("electronLikelihood"); + // d.ismuonlike[i] = tr->getParticleIDProbability("muonLikelihood"); + // d.ispionlike[i] = tr->getParticleIDProbability("pionLikelihood"); + // d.iskaonlike[i] = tr->getParticleIDProbability("kaonLikelihood"); + // d.isprotonlike[i] = tr->getParticleIDProbability("protonLikelihood"); + + // d.iselectronprob[i] = tr->getParticleIDProbability("electronProbability"); + // d.ismuonprob[i] = tr->getParticleIDProbability("muonProbability"); + // d.ispionprob[i] = tr->getParticleIDProbability("pionProbability"); + // d.iskaonprob[i] = tr->getParticleIDProbability("kaonProbability"); + // d.isprotonprob[i] = tr->getParticleIDProbability("protonProbability"); + + + // PDG ID + + // d.pdg_pfa[i] = tr->getPDG(); + // int pidnum = tr->getPDG(); + // d.ismuon[i] = 0; + // d.iselectron[i] = 0; + // d.iskaon[i] = 0; + // d.ispion[i] = 0; + // d.isproton[i] = 0; + // d.isphoton[i] = 0; + // d.isneutralhadron[i] = 0; + // d.iskaon0[i] = 0; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // if (pidnum==11){ + // d.iselectron[i] = 1; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // } + // if (pidnum==13){ + // d.ismuon[i] = 1; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // } + // if (pidnum==2112){ + // d.isneutralhadron[i] = 1; + // } + // if (pidnum==22){ + // d.isphoton[i] = 1; + // } + // if (pidnum==211){ + // d.ispion[i] = 1; + // } + // if (pidnum==321){ + // d.iskaon[i] = 1; + // } + // if (pidnum==310){ + // d.isneutralhadron[i] = 1; + // d.iskaon0[i] = 1; + // } + // if (pidnum==2212){ + // d.isproton[i] = 1; + // } + + + // original version + // d.ismuon[i] = algoEtc::SimpleSecMuonFinder(tr, 5,5,5, -0.1, 0.2, 0.8, 1.5, 4, 0.5, privtx); + // d.iselectron[i] = algoEtc::SimpleSecElectronFinder(tr, 5,5,5,5,0.98,0.9, 1.15, privtx); + // d.isphoton[i] = 0; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // d.isneutralhadron[i] = 0; + // d.pdg_pfa[i] = tr->getPDG(); + + } + // for neutrals + for(i=0;iPx(); + d.neu_py[i] = neu->Py(); + d.neu_pz[i] = neu->Pz(); + d.neu_e[i] = neu->E(); + d.neu_efrac[i] = neu->E() / jet->E(); + d.neu_erel_log[i] = log10(d.neu_efrac[i]); + d.neu_dtheta_ilc[i] = neu->Theta() - jet->Theta(); + d.neu_dphi_ilc[i] = neu->Phi() - jet->Phi(); + if(d.neu_dphi_ilc[i] < -TMath::Pi())d.neu_dphi_ilc[i] += TMath::Pi() * 2; + if(d.neu_dphi_ilc[i] > TMath::Pi())d.neu_dphi_ilc[i] -= TMath::Pi() * 2; + calc_thetaphi(jet->Vect(), neu->Vect(), d.neu_dtheta[i], d.neu_dphi[i]); + + d.neu_cov_d0[i] = -9; + d.neu_cov_z0[i] = -9; + d.neu_cov_phi[i] = -9; + d.neu_cov_omega[i] = -9; + d.neu_cov_tanlambda[i] = -9; + + d.neu_cov_d0_z0[i] = -9; + d.neu_cov_d0_phi[i] = -9; + d.neu_cov_d0_omega[i] = -9; + d.neu_cov_d0_tanlambda[i] = -9; + + d.neu_cov_z0_phi[i] = -9; + d.neu_cov_z0_omega[i] = -9; + d.neu_cov_z0_tanlambda[i] = -9; + + d.neu_cov_phi_omega[i] = -9; + d.neu_cov_phi_tanlambda[i] = -9; + d.neu_cov_omega_tanlambda[i] = -9; + + d.neu_dxy[i] = -9; + d.neu_dz[i] = -9; + d.neu_ip2d_fcc[i] = -9; + d.neu_ip2dsig_fcc[i] = -9; + d.neu_ip3d_fcc[i] = -9; + d.neu_ip3dsig_fcc[i] = -9; + d.neu_jetdist_fcc[i] = -9; + d.neu_jetdistsig_fcc[i] = -9; + + d.neu_charge[i] = 0; + d.neu_ismuon[i] = 0; + d.neu_iselectron[i] = 0; + // simple photon finder + double ecaldep = neu->getCaloEdep()[tpar::ecal]; + double hcaldep = neu->getCaloEdep()[tpar::hcal]; + d.neu_isphoton[i] = (ecaldep / (ecaldep + hcaldep) > 0.98); + d.neu_ischargedhadron[i] = 0; + d.neu_isneutralhadron[i] = !d.isphoton[i]; + d.neu_ispion[i] = 0; + d.neu_iskaon[i] = 0; + d.neu_isproton[i] = 0; + // d.neu_iskaon0[i] = 0; + + d.neu_pdg_pfa[i] = neu->getPDG(); + + //d.neu_mcpid[i] = (neu->getMcp())->getId(); + if(neu->getMcp()){ + d.neu_mcp_pdg[i] = (neu->getMcp())->getPDG();} + else d.neu_mcp_pdg[i] = 0; + d.neu_K_pdg_pfa[i] = 0; + if (d.neu_pdg_pfa[i]==321) d.neu_K_pdg_pfa[i] = 1; + if (d.neu_pdg_pfa[i]==310) d.neu_K_pdg_pfa[i] = 1; + + // d.neu_ismuon[i] = neu->getParticleIDProbability("muonProbability"); + // d.neu_iselectron[i] = neu->getParticleIDProbability("electronProbability"); + // d.neu_ispion[i] = neu->getParticleIDProbability("pionProbability"); + // d.neu_iskaon[i] = neu->getParticleIDProbability("kaonProbability"); + // d.neu_isproton[i] = neu->getParticleIDProbability("protonProbability"); + + + // PDG ID + // d.pdg_pfa[i] = neu->getPDG(); + // int pidnum = neu->getPDG(); + // d.ismuon[i] = 0; + // d.iselectron[i] = 0; + // d.iskaon[i] = 0; + // d.ispion[i] = 0; + // d.isproton[i] = 0; + // d.isphoton[i] = 0; + // d.isneutralhadron[i] = 0; + // d.iskaon0[i] = 0; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // if (pidnum==11){ + // d.iselectron[i] = 1; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // } + // if (pidnum==13){ + // d.ismuon[i] = 1; + // d.ischargedhadron[i] = !(d.ismuon[i] || d.iselectron[i]); + // } + // if (pidnum==2112){ + // d.isneutralhadron[i] = 1; + // } + // if (pidnum==22){ + // d.isphoton[i] = 1; + // } + // if (pidnum==211){ + // d.ispion[i] = 1; + // } + // if (pidnum==321){ + // d.iskaon[i] = 1; + // } + // if (pidnum==310){ + // d.isneutralhadron[i] = 1; + // d.iskaon0[i] = 1; + // } + // if (pidnum==2212){ + // d.isproton[i] = 1; + // } + + + } + + _ntp->Fill(); + } +} + +void DNNProvider2::end() { + _file->Write(); + _file->Close(); +} + +// copied from FCCANalyses/analyzers/dataframe/src/ReconstructedParticle2Track.cc + +float DNNProvider2::calc_dxy(float D0_wrt0, float Z0_wrt0, float phi0_wrt0, TVector3 p, TVector3 privtx, int charge){ + double Bz = Globals::Instance()->getBField(); + double cSpeed = 2.99792458e8 * 1e-9; + + TVector3 X( - D0_wrt0 * TMath::Sin(phi0_wrt0) , D0_wrt0 * TMath::Cos(phi0_wrt0) , Z0_wrt0); + TVector3 x = X - privtx; + //std::cout<<"vertex: "< 0) { + double T = TMath::Sqrt(pt * pt - 2 * a * cross + a * a * r2); + if (pt < 10.0) D = (T - pt) / a; + else D = (-2 * cross + a * r2) / (T + pt); + } + return D; +} + +float DNNProvider2::calc_dz(float D0_wrt0, float Z0_wrt0, float phi0_wrt0, TVector3 p, TVector3 privtx, int charge){ + double Bz = Globals::Instance()->getBField(); + double cSpeed = 2.99792458e8 * 1e-9; + + TVector3 X( - D0_wrt0 * TMath::Sin(phi0_wrt0) , D0_wrt0 * TMath::Cos(phi0_wrt0) , Z0_wrt0); + TVector3 x = X - privtx; + + double a = - charge * Bz * cSpeed; + double pt = p.Pt(); + double C = a/(2 * pt); + double r2 = x(0) * x(0) + x(1) * x(1); + double cross = x(0) * p(1) - x(1) * p(0); + double T = TMath::Sqrt(pt * pt - 2 * a * cross + a * a * r2); + double D; + if (pt < 10.0) D = (T - pt) / a; + else D = (-2 * cross + a * r2) / (T + pt); + double B = C * TMath::Sqrt(TMath::Max(r2 - D * D, 0.0) / (1 + 2 * C * D)); + if ( TMath::Abs(B) > 1.) B = TMath::Sign(1, B); + double st = TMath::ASin(B) / C; + double ct = p(2) / pt; + double z0; + double dot = x(0) * p(0) + x(1) * p(1); + if (dot > 0.0) z0 = x(2) - ct * st; + else z0 = x(2) + ct * st; + + return z0; +} + +float DNNProvider2::calc_sip2d(float D0, float phi0, float jetpx, float jetpy) +{ + TVector2 p(jetpx, jetpy); + + TVector2 d0(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0)); + return TMath::Sign(1, d0 * p) * fabs(D0); +} + +float DNNProvider2::calc_sip3d(float D0, float Z0, float phi0, TVector3 p_jet) +{ + TVector3 d(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0), Z0); + return TMath::Sign(1, d * p_jet) * fabs(sqrt(D0 * D0 + Z0 * Z0)); +} + +float DNNProvider2::calc_jetDist(float D0, float Z0, float phi0, TVector3 p_ct, TVector3 p_jet) +{ + TVector3 d(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0), Z0); + TVector3 r_jet(0.0, 0.0, 0.0); + TVector3 n = p_ct.Cross(p_jet).Unit(); // What if they are parallel? + return n.Dot(d - r_jet); +} + +void DNNProvider2::calc_thetaphi(TVector3 jet, TVector3 part, float &theta, float &phi){ + part.RotateZ(-jet.Phi()); + part.RotateY(-jet.Theta()); + + theta = part.Theta(); + phi = part.Phi(); +} + +} + diff --git a/src/EventStore.cc b/src/EventStore.cc index e1a0bfa..a24facb 100644 --- a/src/EventStore.cc +++ b/src/EventStore.cc @@ -44,6 +44,54 @@ bool EventStore::IsExist(const char* name, const char* classname)const { return false; } +// add flags to collection +bool EventStore::AddFlags(const char* name, int flags) { + // check if collection exists + if (!IsExist(name)) { + cerr << "EventStore::AddFlags: Collection '" << name << "' does not exist" << endl; + return false; + } + + // find all entries with the specified name and add flags + pair::iterator, multimap::iterator> range; + range = _objectMap.equal_range(name); + + if (range.first == _objectMap.end()) { + return false; + } + + // add flags using OR operation + for (multimap::iterator it = range.first; it != range.second; ++it) { + it->second.flag |= flags; + } + + return true; +} + +// remove flags from collection +bool EventStore::RemoveFlags(const char* name, int flags) { + // check if collection exists + if (!IsExist(name)) { + cerr << "EventStore::RemoveFlags: Collection '" << name << "' does not exist" << endl; + return false; + } + + // find all entries with the specified name and remove flags + pair::iterator, multimap::iterator> range; + range = _objectMap.equal_range(name); + + if (range.first == _objectMap.end()) { + return false; + } + + // remove flags using AND NOT operation + for (multimap::iterator it = range.first; it != range.second; ++it) { + it->second.flag &= ~flags; + } + + return true; +} + // obtain class name of the named buffer const char* EventStore::GetClassName(const char* name, int idx)const { int count = _objectMap.count(name); diff --git a/src/FlavorTag.cc b/src/FlavorTag.cc index 5caf343..d42f433 100644 --- a/src/FlavorTag.cc +++ b/src/FlavorTag.cc @@ -1919,6 +1919,10 @@ void FlavorTag::init(Parameters* param) { FTManager& mgr = FTManager::getInstance(); mgr.initVars(); + // Mark the jet collection as PERSIST so it will be written to LCIO + // This is necessary for jets read from LCIO that will have ParticleID added + Event::Instance()->AddFlags(_jetcolname.c_str(), EventStore::PERSIST); + } void FlavorTag::process() { diff --git a/src/LCIOStorer.cc b/src/LCIOStorer.cc index 85b7e17..93fe143 100644 --- a/src/LCIOStorer.cc +++ b/src/LCIOStorer.cc @@ -101,7 +101,7 @@ LCIOStorer::~LCIOStorer() { if (_writer)_writer->close(); } -void LCIOStorer::InitMCPPFOCollections(const char* pfoColName, const char* mcColName, const char* mcpfoColName) { +void LCIOStorer::InitMCPPFOCollections(const char* pfoColName, const char* mcColName, const char* mcpfoColName, const char *mctrkColName) { Event* event = Event::Instance(); vector* pTracks; @@ -130,6 +130,11 @@ void LCIOStorer::InitMCPPFOCollections(const char* pfoColName, const char* mcCol _importMCPFOLinkCols[mcpfoColName] = make_pair(mcColName, pfoColName); } + // mctrack relation - keep PFOs instead of tracks + if (_importMCTrkLinkCols.count(mctrkColName) == 0 && strlen(mctrkColName) > 0) { + _importMCTrkLinkCols[mctrkColName] = make_pair(mcColName, pfoColName); + } + cout << "LCIOStorer initialized with MCParticle collection." << endl; } @@ -374,6 +379,16 @@ void LCIOStorer::SetEvent(lcio::LCEvent* evt) { } //cout << navs.size() << " relation found." << endl; + // looking for LCRelation for tracks + vector navTrks; + //map >::iterator itRelCol; + for (itRelCol = _importMCTrkLinkCols.begin(); itRelCol != _importMCTrkLinkCols.end(); itRelCol ++) { + if (itRelCol->second.second == itPfoCol->first) { + navTrks.push_back(new LCRelationNavigator(evt->getCollection(itRelCol->first))); + } + } + //cout << navTrks.size() << " relation found." << endl; + // sorted PFO list vector pfo_list; for (int i=0; igetNumberOfElements(); ++i) { @@ -391,13 +406,63 @@ void LCIOStorer::SetEvent(lcio::LCEvent* evt) { lcio::MCParticle* mcp = NULL; lcfiplus::MCParticle* mcpf = NULL; - // looking for MCParticle - for (unsigned int n=0; ngetRelatedToObjects(pfo).size()) { - mcp = dynamic_cast(navs[n]->getRelatedToObjects(pfo)[0]); // TODO [0] OK? - mcpf = _mcpLCIORel2[mcp]; - break; - } + // looking for MCParticle for tracks (if Track-MC relation is available) + if(navTrks.size() > 0 && pfo->getCharge() != 0 && pfo->getTracks().size() > 0){ + int tridx =0; // which track in PFO will be used + + if(pfo->getTracks().size() > 1){ + double maxomega = 10000.; + + //cout << "PFO has " << (unsigned int)(pfo->getTracks().size()) << " tracks." << endl; + // selecting track with largest Pt, can also be smallest hit radius (but not taken because of uncertainty) + for(unsigned int n=0;ngetTracks().size();n++){ + lcio::Track *trk = pfo->getTracks()[n]; + if(maxomega > fabs(trk->getOmega())){ + maxomega = fabs(trk->getOmega()); + tridx = n; + } + //cout << " Track " << n << ", pt = " << 0.00105 / trk->getOmega() << " ,rhitmin = " << trk->getRadiusOfInnermostHit() << endl;; + } + //cout << " PFO energy: " << pfo->getEnergy() << " , pt = " << sqrt(pow(pfo->getMomentum()[0],2) + pow(pfo->getMomentum()[1],2)) << endl; + //cout << " Track " << tridx << " is selected." << endl; + } + for (unsigned int n=0; ngetTracks()[tridx]; + double relmax = 0; + for(unsigned int r=0;rgetRelatedToObjects(trk).size();r++){ + double weight = navTrks[n]->getRelatedToWeights(trk)[r]; + if(weight < relmax)continue; + relmax = weight; + mcp = dynamic_cast(navTrks[n]->getRelatedToObjects(trk)[r]); + mcpf = _mcpLCIORel2[mcp]; + break; + } + } + } + + // looking for MCParticle from PFO-MC relation + if(!mcp){ + for (unsigned int n=0; ngetRelatedToObjects(pfo).size()) { + mcp = dynamic_cast(navs[n]->getRelatedToObjects(pfo)[0]); + mcpf = _mcpLCIORel2[mcp]; + break; + } + } else { + // Use weight-based method if Track-MC relation is available + double relmax = 0; + for(unsigned int r=0;rgetRelatedToObjects(pfo).size();r++){ + double weight = navs[n]->getRelatedToWeights(pfo)[r]; + if(weight < relmax)continue; + relmax = weight; + mcp = dynamic_cast(navs[n]->getRelatedToObjects(pfo)[r]); + mcpf = _mcpLCIORel2[mcp]; + break; + } + } + } } // find clusters @@ -432,6 +497,8 @@ void LCIOStorer::SetEvent(lcio::LCEvent* evt) { track->setId(trkIdCounter++); // start from 0. 110927 suehara track->setMcp(mcpf); + + track->setPDG(pfo->getType()); // added for SGV, 250908 suehara track->setCharge(pfo->getCharge()); @@ -518,6 +585,8 @@ void LCIOStorer::SetEvent(lcio::LCEvent* evt) { track->setChi2(trk->getChi2()); track->setNdf(trk->getNdf()); + track->setdEdx(trk->getdEdx()); + //cout << "dEdx:" << trk->getdEdx() << endl; // store detector hit numbers int nhits[lcfiplus::tpar::hitN]; @@ -934,6 +1003,8 @@ void LCIOStorer::WriteVertices(VertexVec* pvvtx, const char* newName, const char } void LCIOStorer::WriteJets(const char* jetName, const char* newName, bool writeVertex, const char* vtxName, const char* relName) { + //cout << "WriteJets" << endl; + if (!_event) { cerr << "LCIOStorer::ConvertJet: LCIO event has not been initialized." << endl; return ; @@ -949,10 +1020,6 @@ void LCIOStorer::WriteJets(const char* jetName, const char* newName, bool writeV // _jetLCIORel.clear(); // _jetLCIORel.resize(pvjet->size()); // map has no resize() - // make collection - lcio::LCCollectionVec* col = new lcio::LCCollectionVec(lcio::LCIO::RECONSTRUCTEDPARTICLE); - lcio::LCRelationNavigator rel(lcio::LCIO::RECONSTRUCTEDPARTICLE, lcio::LCIO::VERTEX); - // naming if (!newName)newName = jetName; string vn = string(newName) + "_vtx"; @@ -960,7 +1027,40 @@ void LCIOStorer::WriteJets(const char* jetName, const char* newName, bool writeV string rn = string(newName) + "_rel"; if (!relName)relName = rn.c_str(); - //TODO: check existence of collections + // check if collection already exists in LCIO + lcio::LCCollection* existingCol = nullptr; + bool isExistingCollection = false; + try { + existingCol = _event->getCollection(newName); + isExistingCollection = true; + } catch (lcio::DataNotAvailableException& e) { + // collection doesn't exist, will create new one + isExistingCollection = false; + } + + // if collection exists, just add ParticleIDs to existing objects + if (isExistingCollection && existingCol) { + cout << "LCIOStorer::WriteJets: Collection '" << newName << "' already exists in LCIO. Adding ParticleIDs to existing jets." << endl; + + for (unsigned int n=0; nsize(); n++) { + const lcfiplus::Jet* flajet = (*pvjet)[n]; + lcio::ReconstructedParticle* lciojet = _jetLCIORel[flajet]; + + if (lciojet) { + // add PID stuffs to existing jet + WriteAllPIDs(existingCol, lciojet, flajet); + } else { + cerr << "LCIOStorer::WriteJets: Warning - no LCIO jet found for lcfiplus jet " << n << endl; + } + } + + // collection already exists, no need to add it again + return; + } + + // create new collection for new jets + lcio::LCCollectionVec* col = new lcio::LCCollectionVec(lcio::LCIO::RECONSTRUCTEDPARTICLE); + lcio::LCRelationNavigator rel(lcio::LCIO::RECONSTRUCTEDPARTICLE, lcio::LCIO::VERTEX); // first reserve vertex collection if (writeVertex) { @@ -1093,6 +1193,7 @@ void LCIOStorer::WriteAllPIDs(lcio::LCCollection* lciocol, lcio::ReconstructedPa map::const_iterator it; for (it = parammap.begin(); it != parammap.end(); it++) { + //cout << "Writing PID " << it->first << endl; WritePID(lciocol, lciojet, lcfijet, it->first.c_str()); } diff --git a/src/LcfiplusProcessor.cc b/src/LcfiplusProcessor.cc index 37278fa..9cd5a68 100644 --- a/src/LcfiplusProcessor.cc +++ b/src/LcfiplusProcessor.cc @@ -42,8 +42,10 @@ LcfiplusProcessor::LcfiplusProcessor() : Processor("LcfiplusProcessor") { _pfoCollectionName, std::string("")); registerInputCollection(LCIO::MCPARTICLE, "MCPCollection" , "MC particle collection", _mcpCollectionName, std::string("")); - registerInputCollection(LCIO::LCRELATION, "MCPFORelation", "Relation between MC and PFO particles", + registerInputCollection(LCIO::LCRELATION, "MCPFORelation", "Relation between MC and PFO particles (MCParticle to ReconstructedParticle)", _mcpfoRelationName, std::string("")); + registerInputCollection(LCIO::LCRELATION, "MCTrackRelation", "Relation between MC and tracks (MCParticle to Track), usually better in terms of assignment of tracks", + _mctrkRelationName, std::string("")); registerProcessorParameter("Algorithms", "LCFIPlus algorithms to run", _algonames, vector()); registerProcessorParameter("ReadSubdetectorEnergies", "Read subdetector energies (ILD)", _readSubdetectorEnergies, int(1)); @@ -162,7 +164,7 @@ void LcfiplusProcessor::init() { // load basic collection if (_useMcp) - _lcio->InitMCPPFOCollections(_pfoCollectionName.c_str(), _mcpCollectionName.c_str(), _mcpfoRelationName.c_str()); + _lcio->InitMCPPFOCollections(_pfoCollectionName.c_str(), _mcpCollectionName.c_str(), _mcpfoRelationName.c_str(), _mctrkRelationName.c_str()); else _lcio->InitPFOCollections(_pfoCollectionName.c_str()); diff --git a/src/MLInferenceWeaver.cc b/src/MLInferenceWeaver.cc new file mode 100644 index 0000000..fb5bb90 --- /dev/null +++ b/src/MLInferenceWeaver.cc @@ -0,0 +1,275 @@ +#include +#include +#include +#include +#include +#include + +#include "lcfiplus.h" +#include "process.h" +#include "MLInputGenerator.h" +#include "MLInferenceWeaver.h" +#include "WeaverInterface.h" +#include "nlohmann/json.hpp" + +using namespace lcfiplus; + +void MLInferenceWeaver::parseJSON(const string& json_filename) { + std::ifstream json_file(json_filename); + std::vector input_names; + std::vector output_names; + try { + const auto json = nlohmann::json::parse(json_file); + + // process output names + json.at("output_names").get_to(output_names); + for (const auto& output : output_names) { + _outputVariables.emplace_back(output); + } + + // process input names + json.at("input_names").get_to(input_names); + for (const auto& input : input_names) { + + // skip if string ends with _mask + const string ending = "_mask"; + if (input.compare(input.size()-ending.size(), ending.size(), ending) == 0) { + continue; + } + + const auto& group_params = json.at(input); + auto& info = prep_info_map_[input]; + info.name = input; + // retrieve the variables names + group_params.at("var_names").get_to(info.var_names); + for (const auto& v : info.var_names) { + _variables.emplace_back(v); + } + } + } catch (const nlohmann::json::exception& exc) { + throw std::runtime_error("Failed to parse input JSON file '" + json_filename + "'.\n" + exc.what()); + } +} + +void MLInferenceWeaver::init(Parameters* param) { + Algorithm::init(param); + _jetCollectionName = param->get("MLInferenceWeaver.JetCollectionName",string("RefinedJets")); + _updateJetCollectionName = param->get("MLInferenceWeaver.UpdateJetCollectionName",string("")); + string privtx = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + Event::Instance()->setDefaultPrimaryVertex(privtx.c_str()); + _jsonFileName = param->get("MLInferenceWeaver.JsonFileName",string("preprocess.json")); + _onnxFileName = param->get("MLInferenceWeaver.OnnxFileName",string("test.onnx")); + _outputParameterName = param->get("MLInferenceWeaver.OutputParameterName",string("weaver")); + + // event classification (true) vs. jet classification (false; default) + _eventClassification = param->get("MLInferenceWeaver.EventClassification",false); + + parseJSON(_jsonFileName); + _weaver = new WeaverInterface(_onnxFileName, _jsonFileName, _variables); + MLInputGenerator::init(); + + // If UpdateJetCollectionName is specified, create a new collection + if (!_updateJetCollectionName.empty()) { + Event::Instance()->Register(_updateJetCollectionName.c_str(), _outputJets, EventStore::PERSIST); + cout << "MLInferenceWeaver: Creating new jet collection '" << _updateJetCollectionName << "'" << endl; + } else { + // Mark the jet collection as PERSIST so it will be written to LCIO + // This is necessary for jets read from LCIO that will have ParticleID added + Event::Instance()->AddFlags(_jetCollectionName.c_str(), EventStore::PERSIST); + _outputJets = 0; + } +} + +void MLInferenceWeaver::process() { + if (_eventClassification) { + processEvent(); + } else { + processJet(); + } +} + +void MLInferenceWeaver::processJet() { + + Event* event = Event::Instance(); + const Vertex* privtx = Event::Instance()->getPrimaryVertex(); + JetVec* jetsPtr(0); + bool success = event->Get(_jetCollectionName.c_str(), jetsPtr); + if (!success) { + cout << "jets could not be found" << endl; + return; + } + JetVec& jets = *jetsPtr; + + // output vector to hold the weights for every jet + rv::RVec< rv::RVec > out; + + // loop over jets + for (unsigned int njet = 0; njet < jets.size(); ++njet) { + const Jet* jet = jets[njet]; + TrackVec tracks = jet->getAllTracksSorted(true); + NeutralVec neutrals = jet->getNeutralsSorted(); + + // prepare input + rv::RVec< rv::RVec > vars; + for (size_t i=0; i<_variables.size(); ++i) { + + // hold computed variable for all the candidates in a jet + rv::RVec cand_vars; + + // find function provided by MLInputGenerator; first check if it exists + const auto& name = _variables[i]; + const auto& calcInputMap = MLInputGenerator::getCalcInput(); + if (calcInputMap.find(name) == calcInputMap.end()) { + cerr << "MLInferenceWeaver: no function to compute input variable " << name << endl; + exit(1); + } + const auto& func = calcInputMap.at(name); + + if (std::holds_alternative >(func)) { + auto f = std::get >(func); + cand_vars.push_back( f(jet) ); + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& tr: tracks) { + cand_vars.push_back( f(tr) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& tr: tracks) { + cand_vars.push_back( f(tr,jet) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& tr: tracks) { + cand_vars.push_back( f(tr,privtx) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& neu: neutrals) { + cand_vars.push_back( f(neu) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& neu: neutrals) { + cand_vars.push_back( f(neu,jet) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& neu: neutrals) { + cand_vars.push_back( f(neu,privtx) ); + } + } + + vars.push_back(cand_vars); + } + + rv::RVec res = _weaver->run(vars); + out.emplace_back(res); + + assert(res.size() == _outputVariables.size()); + + Parameters param; + for (size_t i=0; iaddParam( _outputParameterName.data(), param ); + // Add to output collection + _outputJets->push_back(newJet); + } else { + // Add ParticleID to existing jet + jet->addParam( _outputParameterName.data(), param ); + } + } + +} + +void MLInferenceWeaver::processEvent() { + + Event* event = Event::Instance(); + const Vertex* privtx = Event::Instance()->getPrimaryVertex(); + + // output vector to hold the weights for every jet + rv::RVec< rv::RVec > out; + + // loop over jets + TrackVec tracks = event->getTracks(); + NeutralVec neutrals = event->getNeutrals(); + + // prepare input + rv::RVec< rv::RVec > vars; + for (size_t i=0; i<_variables.size(); ++i) { + + // hold computed variable for all the candidates in a jet + rv::RVec cand_vars; + + // find function provided by MLInputGenerator; first check if it exists + const auto& name = _variables[i]; + const auto& calcInputMap = MLInputGenerator::getCalcInput(); + if (calcInputMap.find(name) == calcInputMap.end()) { + cerr << "MLInferenceWeaver: no function to compute input variable " << name << endl; + exit(1); + } + const auto& func = calcInputMap.at(name); + + if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& tr: tracks) { + cand_vars.push_back( f(tr) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& tr: tracks) { + cand_vars.push_back( f(tr,privtx) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& neu: neutrals) { + cand_vars.push_back( f(neu) ); + } + } + else if (std::holds_alternative >(func)) { + auto f = std::get >(func); + for (const auto& neu: neutrals) { + cand_vars.push_back( f(neu,privtx) ); + } + } + + vars.push_back(cand_vars); + } + + rv::RVec res = _weaver->run(vars); + out.emplace_back(res); + + assert(res.size() == _outputVariables.size()); + + Parameters param; + for (size_t i=0; iaddParam( _outputParameterName.data(), param ); + +} + +void MLInferenceWeaver::end() { + // do nothing +} diff --git a/src/MLInputGenerator.cc b/src/MLInputGenerator.cc new file mode 100644 index 0000000..e486c31 --- /dev/null +++ b/src/MLInputGenerator.cc @@ -0,0 +1,438 @@ +#include + +#include "TFile.h" +#include "TNtuple.h" +#include "TNtupleD.h" +#include "TSystem.h" +#include "TPad.h" +#include "TStyle.h" + +#include "lcfiplus.h" +#include "process.h" +#include "MLInputGenerator.h" +#include "VertexSelector.h" +#include "algoEtc.h" +#include "VertexFinderSuehara.h" +#include "VertexFitterSimple.h" + +#include +#include + +#include +#include +#include + +using namespace lcfiplus; +using namespace std; + +namespace lcfiplus { + +// Static member definitions +CalcInputMap MLInputGenerator::_calcInput; +bool MLInputGenerator::_initialized = false; + +void MLInputGenerator::init(){ + + if (_initialized) { + return; + } + _initialized = true; + //cout << "MLInputGenerator: preparing functions for input variables" << endl; + + //const string _jet_prefix = "jet_"; + //const string _trk_prefix = "tr_"; + //const string _neu_prefix = "neu_"; + + const string _event_prefix = "event_"; + const string _jet_prefix = "jet_"; + const string _trk_prefix = "pfcand_"; + const string _neu_prefix = "neu_pfcand_"; + const string _trk_prefix0 = ""; + const string _neu_prefix0 = "neu_"; + + // event inputs + // number of bs from Higgs decay + _calcInput[_event_prefix+"nbh"] = [](const Event* ev) + { + int nbh = 0; + if(strlen(ev->getDefaultMCParticles())==0)return (double)0.; + + MCParticleVec &mcps = ev->getMCParticles(); + for(unsigned int nmc = 0; nmc < mcps.size();nmc++){ + if(mcps[nmc]->getPDG() == 25 && mcps[nmc]->getDaughters().size()>=1){ // higgs + if(fabs(mcps[nmc]->getDaughters()[0]->getPDG()) == 5) nbh += 2; + } + } + return (double)nbh; + }; + + // jet inputs + _calcInput[_jet_prefix+"px"] = [](const Jet* jet){ return jet->Px(); }; + _calcInput[_jet_prefix+"py"] = [](const Jet* jet){ return jet->Py(); }; + _calcInput[_jet_prefix+"pz"] = [](const Jet* jet){ return jet->Pz(); }; + _calcInput[_jet_prefix+"e"] = [](const Jet* jet){ return jet->E(); }; + _calcInput[_jet_prefix+"mass"] = [](const Jet* jet){ return jet->M(); }; + _calcInput[_jet_prefix+"ntracks"] = [](const Jet* jet){ return jet->getAllTracks().size(); }; + _calcInput[_jet_prefix+"nneutrals"] = [](const Jet* jet){ return jet->getNeutrals().size(); }; + _calcInput[_jet_prefix+"phi"] = [](const Jet* jet){ return jet->Phi(); }; + _calcInput[_jet_prefix+"theta"] = [](const Jet* jet){ return jet->Theta(); }; + + // track inputs + // particle kinematics + _calcInput[_trk_prefix+"px"] = [](const Track* tr){ return tr->Px(); }; + _calcInput[_trk_prefix+"py"] = [](const Track* tr){ return tr->Py(); }; + _calcInput[_trk_prefix+"pz"] = [](const Track* tr){ return tr->Pz(); }; + _calcInput[_trk_prefix+"e"] = [](const Track* tr){ return tr->E(); }; + _calcInput[_trk_prefix+"efrac"] = [](const Track* tr, const Jet* jet){ return tr->E() / jet->E(); }; + _calcInput[_trk_prefix+"erel_log"] = [](const Track* tr, const Jet* jet){ return log10( static_cast( tr->E() / jet->E() ) ); }; + _calcInput[_trk_prefix+"thetarel"] = [](const Track* tr, const Jet* jet){ + float theta, phi; + calc_thetaphi( jet->Vect(), tr->Vect(), theta, phi); + return theta; + }; + _calcInput[_trk_prefix+"phirel"] = [](const Track* tr, const Jet* jet){ + float theta, phi; + calc_thetaphi( jet->Vect(), tr->Vect(), theta, phi); + return phi; + }; + _calcInput[_trk_prefix+"thetarel_ilc"] = [](const Track* tr, const Jet* jet){ return tr->Theta() - jet->Theta(); }; + _calcInput[_trk_prefix+"phirel_ilc"] = [](const Track* tr, const Jet* jet){ + float ret = tr->Phi() - jet->Phi(); + if(ret < -TMath::Pi()) ret += TMath::Pi() * 2; + if(ret > TMath::Pi()) ret -= TMath::Pi() * 2; + return ret; + }; + + // track errors + _calcInput[_trk_prefix+"dptdpt"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::omom]; }; + _calcInput[_trk_prefix+"detadeta"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::tdtd]; }; + _calcInput[_trk_prefix+"dphidphi"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::phph]; }; + _calcInput[_trk_prefix+"dxydxy"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::d0d0]; }; + _calcInput[_trk_prefix+"dzdz"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::z0z0]; }; + _calcInput[_trk_prefix+"dxydz"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::d0z0]; }; + _calcInput[_trk_prefix+"dphidxy"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::d0ph]; }; + _calcInput[_trk_prefix+"dlambdadz"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::z0td]; }; + _calcInput[_trk_prefix+"dxyc"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::d0om]; }; + _calcInput[_trk_prefix+"dxyctgtheta"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::d0td]; }; + _calcInput[_trk_prefix+"phic"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::phom]; }; + _calcInput[_trk_prefix+"phidz"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::z0ph]; }; + _calcInput[_trk_prefix+"phictgtheta"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::phtd]; }; + _calcInput[_trk_prefix+"cdz"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::z0om]; }; + _calcInput[_trk_prefix+"cctgtheta"] = [](const Track* tr){ return tr->getCovMatrix()[tpar::omtd]; }; + + // particle displacements + _calcInput[_trk_prefix0+"d0"] = [](const Track* tr){ return tr->getD0(); }; + _calcInput[_trk_prefix0+"d0sig"] = [](const Track* tr){ return tr->getD0() / sqrt(tr->getCovMatrix()[tpar::cov::d0d0]); }; + _calcInput[_trk_prefix0+"z0"] = [](const Track* tr){ return tr->getZ0(); }; + _calcInput[_trk_prefix0+"z0sig"] = [](const Track* tr){ return tr->getZ0() / sqrt(tr->getCovMatrix()[tpar::cov::z0z0]); }; + + _calcInput[_trk_prefix0+"ip3d"] = [](const Track* tr){ + return sqrt(tr->getD0() * tr->getD0() + tr->getZ0() * tr->getZ0()); + }; + + _calcInput[_trk_prefix0+"ip3dsig"] = [](const Track* tr){ + float ip3d = sqrt(tr->getD0() * tr->getD0() + tr->getZ0() * tr->getZ0()); + return ip3d / sqrt(tr->getCovMatrix()[tpar::cov::d0d0] + tr->getCovMatrix()[tpar::cov::z0z0] + 2 * tr->getCovMatrix()[tpar::cov::d0z0]); + }; + + _calcInput[_trk_prefix+"dxy"] = [](const Track* tr, const Vertex* privtx){ + return calc_dxy(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), privtx->getPos(), tr->getCharge()); + }; + + _calcInput[_trk_prefix+"dz"] = [](const Track* tr, const Vertex* privtx){ + return calc_dz(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), privtx->getPos(), tr->getCharge()); + }; + + _calcInput[_trk_prefix+"btagSip2dVal"] = [](const Track* tr, const Jet* jet){ + return calc_sip2d(tr->getD0(), tr->getPhi(), jet->Px(), jet->Py()); + }; + + _calcInput[_trk_prefix+"btagSip2dSig"] = [](const Track* tr, const Jet* jet){ + return static_cast( calc_sip2d(tr->getD0(), tr->getPhi(), jet->Px(), jet->Py()) )/ sqrt( tr->getCovMatrix()[tpar::d0d0] ); + }; + + _calcInput[_trk_prefix+"btagSip3dVal"] = [](const Track* tr, const Jet* jet){ + return calc_sip3d(tr->getD0(), tr->getZ0(), tr->getPhi(), jet->Vect()); + }; + + _calcInput[_trk_prefix+"btagSip3dSig"] = [](const Track* tr, const Jet* jet){ + return static_cast( calc_sip3d(tr->getD0(), tr->getZ0(), tr->getPhi(), jet->Vect()) )/ sqrt( tr->getCovMatrix()[tpar::d0d0] + tr->getCovMatrix()[tpar::z0z0] ); + }; + + _calcInput[_trk_prefix+"btagJetDistVal"] = [](const Track* tr, const Jet* jet){ + return calc_jetDist(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), jet->Vect()); + }; + + _calcInput[_trk_prefix+"btagJetDistSig"] = [](const Track* tr, const Jet* jet){ + return static_cast( calc_jetDist(tr->getD0(), tr->getZ0(), tr->getPhi(), tr->Vect(), jet->Vect()) )/ sqrt( tr->getCovMatrix()[tpar::d0d0] + tr->getCovMatrix()[tpar::z0z0] ); + }; + + // particle ID + auto trk_pid_index = [](const Track* tr) { + std::vector charged(5); + charged.clear(); + // charged.push_back( tr->getParticleIDProbability("electronProbability") ); + // charged.push_back( tr->getParticleIDProbability("muonProbability") ); + // charged.push_back( tr->getParticleIDProbability("pionProbability") ); + // charged.push_back( tr->getParticleIDProbability("kaonProbability") ); + // charged.push_back( tr->getParticleIDProbability("protonProbability") ); + charged.push_back( tr->getParticleIDProbability("11-ness") ); + charged.push_back( tr->getParticleIDProbability("13-ness") ); + charged.push_back( tr->getParticleIDProbability("211-ness") ); + charged.push_back( tr->getParticleIDProbability("321-ness") ); + charged.push_back( tr->getParticleIDProbability("2212-ness") ); + auto iter = std::max_element(charged.begin(), charged.end()); + int index = std::distance(charged.begin(), iter); + return index; + }; + + _calcInput[_trk_prefix0+"dEdx"] = [](const Track* tr){ return tr->getdEdx(); }; + _calcInput[_trk_prefix+"charge"] = [](const Track* tr){ return tr->getCharge(); }; + _calcInput[_trk_prefix+"isChargedHad"] = [trk_pid_index](const Track* tr){ + return !(trk_pid_index(tr) == 0 || trk_pid_index(tr) == 1); + }; + _calcInput[_trk_prefix+"isNeutralHad"] = [](const Track*){ return 0.; }; + _calcInput[_trk_prefix+"type"] = [](const Track* tr){ return tr->getPDG(); }; + _calcInput[_trk_prefix+"mcpid"] = [](const Track* tr){ return (tr->getMcp()) ? tr->getMcp()->getId() : 0.; }; + _calcInput[_trk_prefix+"mcp_pdg"] = [](const Track* tr){ return (tr->getMcp()) ? tr->getMcp()->getPDG() : 0.; }; + _calcInput[_trk_prefix+"Ktype"] = [](const Track* tr){ + auto pdg = abs(tr->getPDG()); + if (pdg==321 || pdg==310) return 1; + return 0; + }; + + _calcInput[_trk_prefix+"isEl"] = [trk_pid_index](const Track* tr){ return trk_pid_index(tr) == 0; }; + _calcInput[_trk_prefix+"isMu"] = [trk_pid_index](const Track* tr){ return trk_pid_index(tr) == 1; }; + _calcInput[_trk_prefix+"isPion"] = [trk_pid_index](const Track* tr){ return trk_pid_index(tr) == 2; }; + _calcInput[_trk_prefix+"isKaon"] = [trk_pid_index](const Track* tr){ return trk_pid_index(tr) == 3; }; + _calcInput[_trk_prefix+"isProton"] = [trk_pid_index](const Track* tr){ return trk_pid_index(tr) == 4; }; + + // for PI3 (20240203) + _calcInput[_trk_prefix+"proton_K"] = [](const Track* tr){ + //auto isproton = tr->getParticleIDProbability("protonProbability"); + //auto iskaon = tr->getParticleIDProbability("kaonProbability"); + auto isproton = tr->getParticleIDProbability("2212-ness"); + auto iskaon = tr->getParticleIDProbability("321-ness"); + return isproton-iskaon; + }; + _calcInput[_trk_prefix+"pion_K"] = [](const Track* tr){ + //auto ispion = tr->getParticleIDProbability("pionProbability"); + //auto iskaon = tr->getParticleIDProbability("kaonProbability"); + auto ispion = tr->getParticleIDProbability("211-ness"); + auto iskaon = tr->getParticleIDProbability("321-ness"); + return ispion-iskaon; + }; + _calcInput[_trk_prefix+"proton_Klike"] = [](const Track* tr){ + auto isproton = tr->getParticleIDProbability("protonLikelihood"); + auto iskaon = tr->getParticleIDProbability("kaonLikelihood"); + return isproton-iskaon; + }; + _calcInput[_trk_prefix+"pion_Klike"] = [](const Track* tr){ + auto ispion = tr->getParticleIDProbability("pionLikelihood"); + auto iskaon = tr->getParticleIDProbability("kaonLikelihood"); + return ispion-iskaon; + }; + _calcInput[_trk_prefix+"dEdxEl"] = [](const Track* tr){ return tr->getParticleIDProbability("electron_dEdxdistance"); }; + _calcInput[_trk_prefix+"dEdxMu"] = [](const Track* tr){ return tr->getParticleIDProbability("muon_dEdxdistance"); }; + _calcInput[_trk_prefix+"dEdxPion"] = [](const Track* tr){ return tr->getParticleIDProbability("pion_dEdxdistance"); }; + _calcInput[_trk_prefix+"dEdxKaon"] = [](const Track* tr){ return tr->getParticleIDProbability("kaon_dEdxdistance"); }; + _calcInput[_trk_prefix+"dEdxProton"] = [](const Track* tr){ return tr->getParticleIDProbability("proton_dEdxdistance"); }; + _calcInput[_trk_prefix+"iselectronlike"] = [](const Track* tr){ return tr->getParticleIDProbability("electronLikelihood"); }; + _calcInput[_trk_prefix+"ismuonlike"] = [](const Track* tr){ return tr->getParticleIDProbability("muonLikelihood"); }; + _calcInput[_trk_prefix+"ispionlike"] = [](const Track* tr){ return tr->getParticleIDProbability("pionLikelihood"); }; + _calcInput[_trk_prefix+"iskaonlike"] = [](const Track* tr){ return tr->getParticleIDProbability("kaonLikelihood"); }; + _calcInput[_trk_prefix+"isprotonlike"] = [](const Track* tr){ return tr->getParticleIDProbability("protonLikelihood"); }; + // _calcInput[_trk_prefix+"iselectronprob"] = [](const Track* tr){ return tr->getParticleIDProbability("electronProbability"); }; + // _calcInput[_trk_prefix+"ismuonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("muonProbability"); }; + // _calcInput[_trk_prefix+"ispionprob"] = [](const Track* tr){ return tr->getParticleIDProbability("pionProbability"); }; + // _calcInput[_trk_prefix+"iskaonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("kaonProbability"); }; + // _calcInput[_trk_prefix+"isprotonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("protonProbability"); }; + _calcInput[_trk_prefix+"iselectronprob"] = [](const Track* tr){ return tr->getParticleIDProbability("11-ness"); }; + _calcInput[_trk_prefix+"ismuonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("13-ness"); }; + _calcInput[_trk_prefix+"ispionprob"] = [](const Track* tr){ return tr->getParticleIDProbability("211-ness"); }; + _calcInput[_trk_prefix+"iskaonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("321-ness"); }; + _calcInput[_trk_prefix+"isprotonprob"] = [](const Track* tr){ return tr->getParticleIDProbability("2212-ness"); }; + _calcInput[_trk_prefix+"isGamma"] = [](const Track* ){ return 0.; }; + + // for neutrals + // particle kinematics + _calcInput[_neu_prefix+"px"] = [](const Neutral* neu){ return neu->Px(); }; + _calcInput[_neu_prefix+"py"] = [](const Neutral* neu){ return neu->Py(); }; + _calcInput[_neu_prefix+"pz"] = [](const Neutral* neu){ return neu->Pz(); }; + _calcInput[_neu_prefix+"e"] = [](const Neutral* neu){ return neu->E(); }; + _calcInput[_neu_prefix+"efrac"] = [](const Neutral* neu, const Jet* jet){ return neu->E() / jet->E(); }; + _calcInput[_neu_prefix+"erel_log"] = [](const Neutral* neu, const Jet* jet){ return log10(neu->E() / jet->E()); }; + _calcInput[_neu_prefix+"thetarel"] = [](const Neutral* neu, const Jet* jet){ + float theta, phi; + calc_thetaphi(jet->Vect(), neu->Vect(), theta, phi); + return theta; + }; + _calcInput[_neu_prefix+"phirel"] = [](const Neutral* neu, const Jet* jet){ + float theta, phi; + calc_thetaphi(jet->Vect(), neu->Vect(), theta, phi); + return phi; + }; + _calcInput[_neu_prefix+"thetarel_ilc"] = [](const Neutral* neu, const Jet* jet){ return neu->Theta() - jet->Theta(); }; + _calcInput[_neu_prefix+"phirel_ilc"] = [](const Neutral* neu, const Jet* jet){ + auto ret = neu->Phi() - jet->Phi(); + if(ret < -TMath::Pi()) ret += TMath::Pi() * 2; + if(ret > TMath::Pi()) ret -= TMath::Pi() * 2; + return ret; + }; + + // track errors -- why are these implemented for neutrals? + _calcInput[_neu_prefix+"dptdpt"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"detadeta"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dphidphi"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dxydxy"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dzdz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dxydz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dphidxy"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dlambdadz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dxyc"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dxyctgtheta"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"phic"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"phidz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"phictgtheta"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"cdz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"cctgtheta"] = [](const Neutral*){ return -9; }; + + // particle displacements + _calcInput[_neu_prefix0+"d0"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix0+"d0sig"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix0+"z0"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix0+"z0sig"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix0+"ip3d"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix0+"ip3dsig"] = [](const Neutral*){ return 0; }; + + _calcInput[_neu_prefix+"dxy"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"dz"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagSip2dVal"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagSip2dSig"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagSip3dVal"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagSip3dSig"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagJetDistVal"] = [](const Neutral*){ return -9; }; + _calcInput[_neu_prefix+"btagJetDistSig"] = [](const Neutral*){ return -9; }; + + // particle ID + _calcInput[_neu_prefix+"charge"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isMu"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isEl"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isGamma"] = [](const Neutral* neu){ + // simple photon finder + double ecaldep = neu->getCaloEdep()[tpar::ecal]; + double hcaldep = neu->getCaloEdep()[tpar::hcal]; + return (ecaldep / (ecaldep + hcaldep) > 0.98); + }; + _calcInput[_neu_prefix+"isChargedHad"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isNeutralHad"] = [](const Neutral* neu){ + // FIXME: match bug of current training + return 1; + /* + // simple photon finder + double ecaldep = neu->getCaloEdep()[tpar::ecal]; + double hcaldep = neu->getCaloEdep()[tpar::hcal]; + return (ecaldep / (ecaldep + hcaldep) <= 0.98); + */ + }; + _calcInput[_neu_prefix+"type"] = [](const Neutral* neu){ return neu->getPDG(); }; + _calcInput[_neu_prefix+"mcpid"] = [](const Neutral* neu){ return (neu->getMcp()) ? neu->getMcp()->getId() : 0.; }; + _calcInput[_neu_prefix+"mcp_pdg"] = [](const Neutral* neu){ return (neu->getMcp()) ? neu->getMcp()->getPDG() : 0.; }; + + _calcInput[_neu_prefix+"Ktype"] = [](const Neutral* neu){ + auto pdg = abs(neu->getPDG()); + return ( (pdg==321 || pdg==310) ); + }; + _calcInput[_neu_prefix+"isPion"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isKaon"] = [](const Neutral*){ return 0; }; + _calcInput[_neu_prefix+"isProton"] = [](const Neutral*){ return 0; }; + +} + +#if 0 // copied from DNNProvider2.cc +// label +_ntp->Branch("mc_b",&d.mc_b,"mc_b/I"); +_ntp->Branch("mc_c",&d.mc_c,"mc_c/I"); +_ntp->Branch("mc_u",&d.mc_u,"mc_u/I"); +_ntp->Branch("mc_d",&d.mc_d,"mc_d/I"); +_ntp->Branch("mc_s",&d.mc_s,"mc_s/I"); +_ntp->Branch("mc_g",&d.mc_g,"mc_g/I"); +_ntp->Branch("mc_q",&d.mc_q,"mc_q/I"); // usdg +#endif + +// copied from FCCANalyses/analyzers/dataframe/src/ReconstructedParticle2Track.cc +float MLInputGenerator::calc_dxy(float D0_wrt0, float Z0_wrt0, float phi0_wrt0, TVector3 p, TVector3 privtx, int charge){ + double Bz = Globals::Instance()->getBField(); + double cSpeed = 2.99792458e8 * 1e-9; + + TVector3 X( - D0_wrt0 * TMath::Sin(phi0_wrt0) , D0_wrt0 * TMath::Cos(phi0_wrt0) , Z0_wrt0); + TVector3 x = X - privtx; + //std::cout<<"vertex: "< 0) { + double T = TMath::Sqrt(pt * pt - 2 * a * cross + a * a * r2); + if (pt < 10.0) D = (T - pt) / a; + else D = (-2 * cross + a * r2) / (T + pt); + } + return D; +} + +float MLInputGenerator::calc_dz(float D0_wrt0, float Z0_wrt0, float phi0_wrt0, TVector3 p, TVector3 privtx, int charge){ + double Bz = Globals::Instance()->getBField(); + double cSpeed = 2.99792458e8 * 1e-9; + + TVector3 X( - D0_wrt0 * TMath::Sin(phi0_wrt0) , D0_wrt0 * TMath::Cos(phi0_wrt0) , Z0_wrt0); + TVector3 x = X - privtx; + + double a = - charge * Bz * cSpeed; + double pt = p.Pt(); + double C = a/(2 * pt); + double r2 = x(0) * x(0) + x(1) * x(1); + double cross = x(0) * p(1) - x(1) * p(0); + double T = TMath::Sqrt(pt * pt - 2 * a * cross + a * a * r2); + double D; + if (pt < 10.0) D = (T - pt) / a; + else D = (-2 * cross + a * r2) / (T + pt); + double B = C * TMath::Sqrt(TMath::Max(r2 - D * D, 0.0) / (1 + 2 * C * D)); + if ( TMath::Abs(B) > 1.) B = TMath::Sign(1, B); + double st = TMath::ASin(B) / C; + double ct = p(2) / pt; + double z0; + double dot = x(0) * p(0) + x(1) * p(1); + if (dot > 0.0) z0 = x(2) - ct * st; + else z0 = x(2) + ct * st; + + return z0; +} + +float MLInputGenerator::calc_sip2d(float D0, float phi0, float jetpx, float jetpy){ + TVector2 p(jetpx, jetpy); + TVector2 d0(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0)); + return TMath::Sign(1, d0 * p) * fabs(D0); +} + +float MLInputGenerator::calc_sip3d(float D0, float Z0, float phi0, TVector3 p_jet){ + TVector3 d(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0), Z0); + return TMath::Sign(1, d * p_jet) * fabs(sqrt(D0 * D0 + Z0 * Z0)); +} + +float MLInputGenerator::calc_jetDist(float D0, float Z0, float phi0, TVector3 p_ct, TVector3 p_jet){ + TVector3 d(-D0 * TMath::Sin(phi0), D0 * TMath::Cos(phi0), Z0); + TVector3 r_jet(0.0, 0.0, 0.0); + TVector3 n = p_ct.Cross(p_jet).Unit(); // What if they are parallel? + return n.Dot(d - r_jet); +} + +void MLInputGenerator::calc_thetaphi(TVector3 jet, TVector3 part, float &theta, float &phi){ + part.RotateZ(-jet.Phi()); + part.RotateY(-jet.Theta()); + theta = part.Theta(); + phi = part.Phi(); +} + +} diff --git a/src/MLMakeNtuple.cc b/src/MLMakeNtuple.cc new file mode 100644 index 0000000..46c7ba4 --- /dev/null +++ b/src/MLMakeNtuple.cc @@ -0,0 +1,420 @@ +#include "MLInputGenerator.h" +#include "MLMakeNtuple.h" + +#include +#include "EventStore.h" +#include "lcfiplus.h" +#include "JetFinder.h" +#include "algoEtc.h" + +#include "TROOT.h" +#include +#include +#include + +using namespace lcfiplus; + +namespace lcfiplus { + +float& MLMakeNtuple::MLData::newData(const string& key) { + auto const& inserted = _mapData.insert({key, float()}); + auto const& iter = inserted.first; + return iter->second; +} + +vector& MLMakeNtuple::MLData::newDataVec(const string& key) { + auto const& inserted = _mapDataVec.insert({key, vector()}); + auto const& iter = inserted.first; + return iter->second; +} + +void MLMakeNtuple::MLData::setData(const string& key, const float& val) { + _mapData[key] = val; +} + +void MLMakeNtuple::MLData::addDataVec(const string& key, const float& val) { + _mapDataVec[key].push_back(val); +} + +void MLMakeNtuple::MLData::resetData() { + for (auto& d: _mapData) d.second = 0; + for (auto& d: _mapDataVec) d.second.clear(); +} + +void MLMakeNtuple::init(Parameters* param) { + Algorithm::init(param); + _jetname = param->get("MLMakeNtuple.JetCollectionName",string("RefinedJets")); + string privtx = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + Event::Instance()->setDefaultPrimaryVertex(privtx.c_str()); + + string outputFilename = param->get("MLMakeNtuple.OutputRootFileName",string("ml_flavtag.root")); + string treeName = param->get("MLMakeNtuple.TTreeName",string("ntp")); + _labelKeep = param->get("MLMakeNtuple.Label",int(0)); + _outEvent = param->get("MLMakeNtuple.EventClassification",int(0)); + _outEventNoJets = param->get("MLMakeNtuple.EventClassificationNoJets",int(0)); + + // Validate that only one of EventClassification or EventClassificationNoJets is enabled + if (_outEvent && _outEventNoJets) { + throw Exception("Ambiguous configuration: EventClassification and EventClassificationNoJets cannot both be enabled. Please set only one of them to 1."); + } + + //cout << "MLMakeNtuple: Ntuple file set to " << outputFilename << endl; + _file = new TFile(outputFilename.c_str(),"RECREATE"); + + //cout << "MLMakeNtuple: setting TTree name '" << treeName.c_str() << "'" << endl; + _tree = new TTree(treeName.c_str(),"ML input data for flavor tagging"); + + MLInputGenerator::init(); + + _tree->Branch("label",&_label,"Label/I"); + + // TODO: implement labels +#if 0 + // label + _tree->Branch("mc_b",&d.mc_b,"mc_b/I"); + _tree->Branch("mc_c",&d.mc_c,"mc_c/I"); + _tree->Branch("mc_u",&d.mc_u,"mc_u/I"); + _tree->Branch("mc_d",&d.mc_d,"mc_d/I"); + _tree->Branch("mc_s",&d.mc_s,"mc_s/I"); + _tree->Branch("mc_g",&d.mc_g,"mc_g/I"); + _tree->Branch("mc_q",&d.mc_q,"mc_q/I"); // usdg +#endif + + if (_outEvent){ + string key = "jetIndex"; + _tree->Branch( key.c_str(), &_data.newDataVec(key) ); + + key = "neu_jetIndex"; + _tree->Branch( key.c_str(), &_data.newDataVec(key) ); + } + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + + if(_outEventNoJets){ + if (std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second)){ + _tree->Branch( key.c_str(), &_data.newDataVec(key) ); + } + } + + if (_outEvent){ + if (std::holds_alternative >(v.second)) { + string buf = key + "/F"; + _tree->Branch( key.c_str(), &_data.newData(key), buf.c_str() ); + }else if (std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + || std::holds_alternative >(v.second) + ) + { + _tree->Branch( key.c_str(), &_data.newDataVec(key) ); + } + + } else { + + if (std::holds_alternative >(v.second) + || std::holds_alternative >(v.second)) { + string buf = key + "/F"; + _tree->Branch( key.c_str(), &_data.newData(key), buf.c_str() ); + } else { + _tree->Branch( key.c_str(), &_data.newDataVec(key) ); + } + } + + } +} + +//template struct overloaded : Ts... { using Ts::operator()...; }; +//template overloaded(Ts...) -> overloaded; + +void MLMakeNtuple::process() { + // Configuration validation is now done in init() + + if(_outEvent) + processEvent(); + else if (_outEventNoJets) + processEventNoJets(); + else + processJets(); +} + +void MLMakeNtuple::processEventNoJets(){ + Event* event = Event::Instance(); + const Vertex* privtx = Event::Instance()->getPrimaryVertex(); + _label = _labelKeep; + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + double ret = f(Event::Instance()); + _data.setData(key,ret); + } + } + + // get tracks and neutrals + TrackVec &tracks_orig = event->getTracks(); + NeutralVec &neutrals_orig = event->getNeutrals(); + + vector tracks(tracks_orig.begin(), tracks_orig.end()); + vector neutrals(neutrals_orig.begin(), neutrals_orig.end()); + + std::sort(tracks.begin(), tracks.end(), [](const Track *a, const Track *b){ + return a->E() > b->E(); + }); + std::sort(neutrals.begin(), neutrals.end(), [](const Neutral *a, const Neutral *b){ + return a->E() > b->E(); + }); + + _data.resetData(); + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr,privtx); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu,privtx); + _data.addDataVec(key,ret); + } + } + } + + _tree->Fill(); + +} + +void MLMakeNtuple::processEvent(){ + + Event* event = Event::Instance(); + const Vertex* privtx = Event::Instance()->getPrimaryVertex(); + JetVec* jetsPtr(0); + bool success = event->Get(_jetname.c_str(), jetsPtr); + if (!success) { + cout << "jets could not be found" << endl; + return; + } + JetVec& jets = *jetsPtr; + + _data.resetData(); + _label = _labelKeep; + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + double ret = f(Event::Instance()); + _data.setData(key,ret); + } + } + + for (unsigned int njet = 0; njet < jets.size(); ++njet) { + + const Jet* jet = jets[njet]; + TrackVec tracks = jet->getAllTracksSorted(true); + NeutralVec neutrals = jet->getNeutralsSorted(); + + for (auto tr: tracks) { + (void) tr; + _data.addDataVec("jetIndex",(float)njet); + } + for (auto neu: neutrals) { + (void) neu; + _data.addDataVec("neu_jetIndex",(float)njet); + } + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + + /* + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + double ret = f(jet); + _data.setData(key,ret); + } + */ + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr,jet); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr,privtx); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu,jet); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu,privtx); + _data.addDataVec(key,ret); + } + } + } + } + + _tree->Fill(); + +} + +void MLMakeNtuple::processJets(){ + Event* event = Event::Instance(); + const Vertex* privtx = Event::Instance()->getPrimaryVertex(); + JetVec* jetsPtr(0); + bool success = event->Get(_jetname.c_str(), jetsPtr); + if (!success) { + cout << "jets could not be found" << endl; + return; + } + JetVec& jets = *jetsPtr; + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + double ret = f(Event::Instance()); + _data.setData(key,ret); + } + } + + vector mcpJets; + if(_labelKeep == 0 && Event::Instance()->isMCExist()){ + algoEtc::AssignJetsToMC(jets, mcpJets); + } + for (unsigned int njet = 0; njet < jets.size(); ++njet) { + _label = _labelKeep; + if(_label == 0 && Event::Instance()->isMCExist()){ + _label = (mcpJets[njet] ? mcpJets[njet]->getPDG() : 0); + } + + _data.resetData(); + const Jet* jet = jets[njet]; + TrackVec tracks = jet->getAllTracksSorted(true); + NeutralVec neutrals = jet->getNeutralsSorted(); + + for (const auto& v: MLInputGenerator::getCalcInput()) { + const auto& key = v.first; + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + double ret = f(jet); + _data.setData(key,ret); + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr,jet); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto tr: tracks) { + double ret = f(tr,privtx); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu,jet); + _data.addDataVec(key,ret); + } + } + + if (std::holds_alternative >(v.second)) { + auto f = std::get >(v.second); + for (auto neu: neutrals) { + double ret = f(neu,privtx); + _data.addDataVec(key,ret); + } + } + } + _tree->Fill(); + } +} + +void MLMakeNtuple::end() { + _file->Write(); + _file->Close(); +} + +} diff --git a/src/ONNXRuntime.cc b/src/ONNXRuntime.cc new file mode 100644 index 0000000..9d65cbd --- /dev/null +++ b/src/ONNXRuntime.cc @@ -0,0 +1,124 @@ +// Code from https://github.com/HEP-FCC/FCCAnalyses + +#include "ONNXRuntime.h" + +#include +#include + +using namespace lcfiplus; + +ONNXRuntime::ONNXRuntime(const std::string& model_path, const std::vector& input_names) + : env_(new Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "onnx_runtime")), + memoryInfo_(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault)), + input_names_(input_names) { + if (model_path.empty()) + throw std::runtime_error("Path to ONNX model cannot be empty!"); + Ort::SessionOptions options; + options.SetIntraOpNumThreads(1); + session_ = std::make_unique(*env_, model_path.c_str(), options); + + Ort::AllocatorWithDefaultOptions allocator; +#if ORT_API_VERSION < 13 + // Before 1.13 we have to roll our own unique_ptr wrapper here + auto allocDeleter = [&allocator](char* p) { allocator.Free(p); }; + using AllocatedStringPtr = std::unique_ptr; +#endif + + // get input names and shapes + input_node_dims_.clear(); + for (size_t i = 0; i < session_->GetInputCount(); ++i) { +#if ORT_API_VERSION < 13 + input_node_strings_.emplace_back(AllocatedStringPtr(session_->GetInputName(i, allocator), allocDeleter).release()); +#else + input_node_strings_.emplace_back(session_->GetInputNameAllocated(i, allocator).release()); +#endif + + const auto& input_name = input_node_strings_[i]; + // get input shapes + const auto nodeInfo = session_->GetInputTypeInfo(i); + input_node_dims_[input_name] = nodeInfo.GetTensorTypeAndShapeInfo().GetShape(); + } + + // get output names and shapes + output_node_dims_.clear(); + for (size_t i = 0; i < session_->GetOutputCount(); ++i) { +#if ORT_API_VERSION < 13 + output_node_strings_.emplace_back(AllocatedStringPtr(session_->GetOutputName(i, allocator), allocDeleter).release()); +#else + output_node_strings_.emplace_back(session_->GetOutputNameAllocated(i, allocator).release()); +#endif + + // get output node names + const auto& output_name = output_node_strings_[i]; + // get output node types + const auto nodeInfo = session_->GetOutputTypeInfo(i); + output_node_dims_[output_name] = nodeInfo.GetTensorTypeAndShapeInfo().GetShape(); + // the 0th dim depends on the batch size + output_node_dims_[output_name].at(0) = -1; + } +} + +ONNXRuntime::~ONNXRuntime() {} + +template +ONNXRuntime::Tensor ONNXRuntime::run(Tensor& input, + const Tensor& input_shapes, + unsigned long long batch_size) const { + std::vector tensors_in; + for (const auto& name : input_node_strings_) { + auto input_pos = variablePos(name); + auto value = input.begin() + input_pos; + std::vector input_dims; + if (input_shapes.empty()) { + input_dims = input_node_dims_.at(name); + input_dims[0] = batch_size; + } else { + input_dims = input_shapes[input_pos]; + // rely on the given input_shapes to set the batch size + if (input_dims[0] != batch_size) + throw std::runtime_error("The first element of `input_shapes` (" + std::to_string(input_dims[0]) + + ") does not match the given `batch_size` (" + std::to_string(batch_size) + ")"); + } + auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies()); + if (expected_len != (int64_t)value->size()) + throw std::runtime_error("Input array '" + std::string(name) + "' has a wrong size of " + std::to_string(value->size()) + + ", expected " + std::to_string(expected_len)); + auto input_tensor = Ort::Value::CreateTensor(memoryInfo_, const_cast(value->data()), value->size(), input_dims.data(), input_dims.size()); + if (!input_tensor.IsTensor()) + throw std::runtime_error("Failed to create an input tensor for variable '" + std::string(name) + "'."); + + tensors_in.emplace_back(std::move(input_tensor)); + } + + // run the inference + auto output_tensors = session_->Run(Ort::RunOptions{nullptr}, input_node_strings_.data(), tensors_in.data(), tensors_in.size(), output_node_strings_.data(), output_node_strings_.size()); + + // convert output tensor to values + Tensor outputs; + size_t i = 0; + for (auto& output_tensor : output_tensors) { + if (!output_tensor.IsTensor()) + throw std::runtime_error("(at least) inference output " + std::to_string(i) + " is not a tensor."); + // get output shape + auto tensor_info = output_tensor.GetTensorTypeAndShapeInfo(); + auto length = tensor_info.GetElementCount(); + + auto floatarr = output_tensor.GetTensorMutableData(); + outputs.emplace_back(floatarr, floatarr + length); + ++i; + } + if (outputs.size() != session_->GetOutputCount()) + throw std::runtime_error("Number of outputs differ from the expected one: got " + std::to_string(outputs.size()) + + ", expected " + std::to_string(session_->GetOutputCount())); + + return outputs; +} + +size_t ONNXRuntime::variablePos(const std::string& name) const { + auto iter = std::find(input_names_.begin(), input_names_.end(), name); + if (iter == input_names_.end()) + throw std::runtime_error("Input variable '" + name + " is not provided"); + return iter - input_names_.begin(); +} + +template ONNXRuntime::Tensor ONNXRuntime::run(Tensor&, const Tensor&, unsigned long long) const; \ No newline at end of file diff --git a/src/WeaverInterface.cc b/src/WeaverInterface.cc new file mode 100644 index 0000000..335f209 --- /dev/null +++ b/src/WeaverInterface.cc @@ -0,0 +1,204 @@ +// Code from https://github.com/HEP-FCC/FCCAnalyses + +#include "WeaverInterface.h" + +#include "nlohmann/json.hpp" +#include +#include + +//#define DUMP_WEAVER_INPUT 1 + +using namespace lcfiplus; + +WeaverInterface::WeaverInterface(const std::string& onnx_filename, + const std::string& json_filename, + const rv::RVec& vars) + : variables_names_(vars.begin(), vars.end()) { + if (onnx_filename.empty()) + throw std::runtime_error("ONNX modeld input file not specified!"); + if (json_filename.empty()) + throw std::runtime_error("JSON preprocessed input file not specified!"); + + // the preprocessing JSON was found ; extract the variables listing and all useful information + std::ifstream json_file(json_filename); + std::vector input_names; + try { + const auto json = nlohmann::json::parse(json_file); + json.at("input_names").get_to(input_names); + for (const auto& input : input_names) { + const auto& group_params = json.at(input); + auto& info = prep_info_map_[input]; + info.name = input; + // retrieve the variables names + group_params.at("var_names").get_to(info.var_names); + // retrieve the shapes for all variables + if (group_params.contains("var_length")) { + info.min_length = info.max_length = group_params.at("var_length"); + input_shapes_.push_back({1, (int64_t)info.var_names.size(), (int64_t)info.min_length}); + } else { + info.min_length = group_params.at("min_length"); + info.max_length = group_params.at("max_length"); + input_shapes_.push_back({1, (int64_t)info.var_names.size(), -1}); + } + // for all variables, retrieve the allowed range + const auto& var_info_params = group_params.at("var_infos"); + for (const auto& name : info.var_names) { + const auto& var_params = var_info_params.at(name); + info.var_info_map[name] = + PreprocessParams::VarInfo(var_params.at("median"), + var_params.at("norm_factor"), + var_params.at("replace_inf_value"), + var_params.at("lower_bound"), + var_params.at("upper_bound"), + var_params.contains("pad") ? (double)var_params.at("pad") : 0.); + } + // create data storage with a fixed size vector initialised with 0's + const auto& len = input_sizes_.emplace_back(info.max_length * info.var_names.size()); + data_.emplace_back(len, 0); + } + } catch (const nlohmann::json::exception& exc) { + throw std::runtime_error("Failed to parse input JSON file '" + json_filename + "'.\n" + exc.what()); + } + onnx_ = std::make_unique(onnx_filename, input_names); +} + +std::vector WeaverInterface::center_norm_pad(const rv::RVec& input, + float center, + float scale, + size_t min_length, + size_t max_length, + float pad_value, + float replace_inf_value, + float min, + float max) { + if (min > pad_value || pad_value > max) + throw std::runtime_error("Pad value not within (min, max) range"); + if (min_length > max_length) + throw std::runtime_error("Variable length mismatch (min_length >= max_length)"); + + auto ensure_finitude = [](const float in, const float replace_val) { + if (!std::isfinite(in)) + return replace_val; + return in; + }; + + size_t target_length = std::clamp(input.size(), min_length, max_length); + std::vector out(target_length, pad_value); + for (size_t i = 0; i < input.size() && i < target_length; ++i) + out[i] = std::clamp((ensure_finitude(input[i], replace_inf_value) - center) * scale, min, max); + return out; +} + +size_t WeaverInterface::variablePos(const std::string& var_name) const { + auto var_it = std::find(variables_names_.begin(), variables_names_.end(), var_name); + if (var_it == variables_names_.end()) + throw std::runtime_error("Unable to find variable with name '" + var_name + + "' in the list of registered variables"); + return var_it - variables_names_.begin(); +} + +#ifdef DUMP_WEAVER_INPUT +void WeaverInterface::writeToFile(const rv::RVec< rv::RVec >& vars) { + static size_t cnt(0); + ++cnt; + char buf[1024]; + sprintf(buf,"pre%05ld.txt",cnt); + std::ofstream out(buf); + for (size_t i=0; i& row = vars[i]; + //const auto& var_name = params.var_names.at(i); + for (size_t j=0; j >& vars) { + static size_t cnt(0); + ++cnt; + + char buf[1024]; + sprintf(buf,"norm%05ld.txt",cnt); + std::ofstream out(buf); + + //out << "input_sizes_.size()=" << input_sizes_.size() << std::endl; + //out << "vars.size()=" << vars.size() << std::endl; + for (size_t i=0; iinputNames()[i]; + const auto& params = prep_info_map_.at(name); + out << "BLOCK[" << i << "]: " << name; + const std::vector& row = vars[i]; + //out << "row.size()=" << row.size() << std::endl; + for (size_t j=0; j WeaverInterface::run(const rv::RVec& constituents) { + size_t i = 0; + for (const auto& name : onnx_->inputNames()) { + const auto& params = prep_info_map_.at(name); + auto& values = data_[i]; + values.resize(input_sizes_.at(i)); + std::fill(values.begin(), values.end(), 0); + size_t it_pos = 0; + ConstituentVars jc; + for (size_t j = 0; j < params.var_names.size(); ++j) { // transform and add the proper amount of padding + const auto& var_name = params.var_names.at(j); + //if (std::find(variables_names_.begin(), variables_names_.end(), "pfcand_mask") == variables_names_.end()) + // jc = ConstituentVars(constituents.at(0).size(), 1.f); + + //if (var_name.find("pfcand_mask") != std::string::npos) + //jc = ConstituentVars(constituents.at(0).size(), 1.f); + + if (var_name == "pfcand_mask") { + jc = ConstituentVars(constituents.at(0).size(), 1.f); + } else if (var_name == "neu_pfcand_mask") { + jc = ConstituentVars(constituents.at(constituents.size()-1).size(), 1.f); + } else { + jc = constituents.at(variablePos(var_name)); + } + const auto& var_info = params.info(var_name); + auto val = center_norm_pad(jc, + var_info.center, + var_info.norm_factor, + params.min_length, + params.max_length, + var_info.pad, + var_info.replace_inf_value, + var_info.lower_bound, + var_info.upper_bound); + std::copy(val.begin(), val.end(), values.begin() + it_pos); + it_pos += val.size(); + } + values.resize(it_pos); + ++i; + } +#ifdef DUMP_WEAVER_INPUT + writeToFile(constituents); + writeToFile(data_); +#endif + return onnx_->run(data_, input_shapes_)[0]; +} + +void WeaverInterface::PreprocessParams::dumpVars() const { + std::cout << "List of variables for preprocessing parameter '" << name + << "': " << rv::RVec(var_names.begin(), var_names.end()) << "." << std::endl; +} \ No newline at end of file diff --git a/src/algoEtc.cc b/src/algoEtc.cc index b78463c..55dabd8 100644 --- a/src/algoEtc.cc +++ b/src/algoEtc.cc @@ -371,6 +371,106 @@ bool SimpleSecElectronFinder(const Track* tr, double d0sigth, double z0sigth, do } +void AssignJetsToMC(JetVec &jets, vector&mcs_out){ + bool debug = false; + + map > mapMcJet; + map mapJetMc; + + MCColorSingletVec &mccsv = Event::Instance()->getMCColorSinglets(); + + for(JetVecIte it = jets.begin(); it != jets.end(); it++){ + mapJetMc[*it] = 0; // initialization in case cannot find proper MC + + map mapMccs; + TrackVec vtr = (*it)->getAllTracks(); + for(TrackVecIte itt = vtr.begin();itt != vtr.end();itt++){ + if(!(*itt)->getMcp())continue; + const MCColorSinglet *mccs = (*itt)->getMcp()->getColorSinglet(&mccsv); + if(!mccs)continue; + if(mapMccs.find(mccs) != mapMccs.end()) + mapMccs[mccs] += (*itt)->E(); + else + mapMccs[mccs] = (*itt)->E(); + } + NeutralVec vnt = (*it)->getNeutrals(); + for(NeutralVecIte itn = vnt.begin(); itn != vnt.end(); itn++){ + if(!(*itn)->getMcp())continue; + const MCColorSinglet *mccs = (*itn)->getMcp()->getColorSinglet(&mccsv); + if(!mccs)continue; + + if(mapMccs.find(mccs) != mapMccs.end()) + mapMccs[mccs] += (*itn)->E(); + else + mapMccs[mccs] = (*itn)->E(); + } + + // find ColorSinglet assignment + const MCColorSinglet *mccsf = 0; + double mccsd = 0.; + + if(debug) cout << "MCCS mapsize: " << mapMccs.size() << endl; + for(map::const_iterator itmccs = mapMccs.begin(); itmccs != mapMccs.end(); itmccs ++){ + if(debug) cout << "MCCS assignment: energy = " << itmccs->second << ", MCCS: " << (long long)itmccs->first << endl; + if(itmccs->second > mccsd){ + mccsf = itmccs->first; + mccsd = itmccs->second; + } + } + + if(debug) + cout << "ColorSinglet found: Jet E = " << (*it)->E() << ", MCCS E = " << (mccsf->getMcp() ? mccsf->getMcp()->E() : 0) << endl; + + if(!mccsf){ + cout << "ColorSinglet could not found. Skipping the jet." << endl; + mapJetMc[*it] = 0; + continue; + } + + // find initial based on angle + // record energy fraction + const MCParticle *mcq = 0; + double mcqd = 100; + MCParticleVec &initials = mccsf->_initials; + for(MCParticleVecIte itin = initials.begin(); itin != initials.end(); itin ++){ + double angle = (*itin)->Angle((*it)->Vect()); + if(debug){ + cout << "Checking MCP: pdg = " << (*itin)->getPDG() << ", angle = " << angle << endl; + cout << "Jet direction: (" << (*it)->Px() << ", " << (*it)->Py() << ", " << (*it)->Pz() << ")" << endl; + cout << "MCP direction: (" << (*itin)->Px() << ", " << (*itin)->Py() << ", " << (*itin)->Pz() << ")" << endl; + } + if(angle < mcqd){ + mcq = *itin; + mcqd = angle; + } + } + if(debug) + cout << "Initial MCP found: pdg = " << mcq->getPDG() << ", angle = " << mcqd << endl; + auto itmap = mapMcJet.find(mcq); + if(itmap == mapMcJet.end() || itmap->second.second < (*it)->E()){ // not registered or larger jet energy + if(itmap != mapMcJet.end()){ // removing previous assignment + if(debug) + cout << "Jet with energy " << itmap->second.first->E() << " is removed with MCP." << endl; + mapJetMc[itmap->second.first] = 0; + } + mapMcJet[mcq] = make_pair((*it), (*it)->E()); + mapJetMc[*it] = mcq; + if(debug) + cout << "Jet with energy " << (*it)->E() << " is associated with MCP." << endl; + } + } + + // extract map to vector + for(unsigned int ij = 0; ij < jets.size(); ij++){ + mcs_out.push_back(mapJetMc[jets[ij]]); + if(debug && mapJetMc[jets[ij]]) + cout << "Jet " << ij << " assigned to MC PDG " << mapJetMc[jets[ij]]->getPDG() << endl; + } + + if(debug) + cout << "AssignJetsToMC finished." << endl; +} + } } diff --git a/src/flavtag.cc b/src/flavtag.cc index b1209a5..f82bb5b 100644 --- a/src/flavtag.cc +++ b/src/flavtag.cc @@ -89,7 +89,7 @@ void FTManager::openFile(const char* filename) { void FTManager::closeFile() { if (_file) { cout << "FTManager: closing file" << endl; - _tree->Write(); + _file->Write(); _file->Close(); _file = 0; } else { diff --git a/src/lcfiplus.cc b/src/lcfiplus.cc index 29ca155..5b420fb 100644 --- a/src/lcfiplus.cc +++ b/src/lcfiplus.cc @@ -299,7 +299,7 @@ double Track::getY() const { } double Track::getZ() const { - return _par[tpar::z0] + _par[tpar::om]*_flt; + return _par[tpar::z0] + _par[tpar::td]*_flt; } void Track::setCovMatrix(double* mycov) { @@ -1270,6 +1270,51 @@ vector Jet::getAllTracks(bool withoutV0) const { return ret; } +vector Jet::getAllTracksSorted(bool withoutV0) const { + vector tracks_unsorted = getAllTracks(withoutV0); + + // order particles by energy + vector > order_tr; + order_tr.resize(tracks_unsorted.size()); + for(size_t i=0; i(tracks_unsorted[i]->E(), i); + } + std::sort(order_tr.begin(),order_tr.end(),[](std::paira,std::pair b){ + return a.first > b.first; + }); + + vector tracks; + tracks.resize(tracks_unsorted.size()); + + for (size_t i=0; i Jet::getNeutralsSorted() const { + vector neutrals_unsorted = getNeutrals(); + + // order particles by energy + vector > order_n; + order_n.resize(neutrals_unsorted.size()); + for(size_t i=0; i(neutrals_unsorted[i]->E(), i); + } + std::sort(order_n.begin(),order_n.end(),[](std::paira,std::pair b){ + return a.first > b.first; + }); + + vector neutrals; + neutrals.resize(neutrals_unsorted.size()); + + for (size_t i=0; i Jet::getVerticesForFT() const { vector ret; diff --git a/src/testproc.cc b/src/testproc.cc index 29c0d76..d21a581 100644 --- a/src/testproc.cc +++ b/src/testproc.cc @@ -993,14 +993,21 @@ void TestAlgo::end() { void FlavtagReader::init(Parameters* param) { Algorithm::init(param); - string filename = param->get("FileName",string("test.root")); - _jetname = param->get("JetCollectionName",string("Durham_2Jets")); + string filename = param->get("FlavtagReader.FileName",string("test.root")); + _jetname = param->get("FlavtagReader.JetCollectionName",string("Durham_2Jets")); string primvtxcolname = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); Event::Instance()->setDefaultPrimaryVertex(primvtxcolname.c_str()); + _is6cat = param->get("FlavtagReader.SixCategories",int(1)); + _isWeaver = param->get("FlavtagReader.Weaver",int(0)); _file = new TFile(filename.c_str(),"RECREATE"); - _nt = new TNtupleD("nt","nt","nev:nj:e:px:py:pz:btag:ctag:otag:bbtag:bctag:cctag"); - _ntev = new TNtupleD("ntev","ntev","nev:btag1:btag2:btag3:btag4:btag5:btag6:ctag1:ctag2:ctag3:ctag4:ctag5:ctag6"); + if(_is6cat){ + _nt = new TNtupleD("nt","nt","nev:nj:e:px:py:pz:btag:ctag:otag:bbtag:bctag:cctag"); + }else{ + _nt = new TNtupleD("nt","nt","nev:nj:e:px:py:pz:btag:ctag:otag"); + } + + _ntev = new TNtupleD("ntev","ntev","nev:nbz:nbh:btag1:btag2:btag3:btag4:btag5:btag6:ctag1:ctag2:ctag3:ctag4:ctag5:ctag6"); _jets = 0; _nev = 0; @@ -1017,19 +1024,59 @@ void FlavtagReader::process() { for (unsigned int nj = 0; nj < _jets->size(); nj++) { const Jet* j = (*_jets)[nj]; - const Parameters* para = j->getParam("lcfiplus"); - _nt->Fill(_nev, nj, j->E(), j->Px(), j->Py(), j->Pz(), + double btag = 0.; + double ctag = 0.; + + const Parameters* para = j->getParam(_isWeaver ? "weaver" : "lcfiplus"); + + btag = para->get(_isWeaver ? "mc_b" : "BTag"); + ctag = para->get(_isWeaver ? "mc_c" : "CTag"); + + if(_isWeaver){ + _nt->Fill(_nev, nj, j->E(), j->Px(), j->Py(), j->Pz(), + para->get("mc_b"), para->get("mc_c"), para->get("mc_d")); + } + else if(_is6cat){ + _nt->Fill(_nev, nj, j->E(), j->Px(), j->Py(), j->Pz(), para->get("BTag"), para->get("CTag"), para->get("OTag"), para->get("BBTag"), para->get("CCTag"), para->get("BCTag")); - btags.push_back(para->get("BTag")); - ctags.push_back(para->get("CTag")); + }else{ + _nt->Fill(_nev, nj, j->E(), j->Px(), j->Py(), j->Pz(), + para->get("BTag"), para->get("CTag"), para->get("OTag")); + } + btags.push_back(btag); + ctags.push_back(ctag); - cout << "nvtx = " << para->get("nvtx") << ", nvtxall = " << para->get("nvtxall") << endl; + // if(!_isWeaver){ + // cout << "nvtx = " << para->get("nvtx") << ", nvtxall = " << para->get("nvtxall") << endl; + //} } + // counting MC b for ZHH/ZZH analysis + + int nbz = 0; + int nbh = 0; + Event *event = Event::Instance(); + if(event->IsExist(event->getDefaultMCParticles()) && event->getMCParticles().size() >= 12){ + for(unsigned int i=8;i<12;i++){ + //cout << i << " " << event->getMCParticles()[i]->getPDG() << endl; + if(fabs(event->getMCParticles()[i]->getPDG()) == 5)nbz ++; + } + for(unsigned int i=0;igetMCParticles().size();i++){ + //cout << "quark " << i << endl; + if(event->getMCParticles()[i]->getPDG() == 25){//Higgs + const MCParticle *mcp = event->getMCParticle(i); + for(unsigned int j=0;jgetDaughters().size();j++){ + //cout << "higgs " << j << endl; + if(fabs(mcp->getDaughters()[j]->getPDG()) == 5)nbh ++; + } + } + } + } + std::sort(btags.begin(), btags.end()); std::sort(ctags.begin(), ctags.end()); if (_jets->size() >= 6) - _ntev->Fill(_nev, btags[0], btags[1], btags[2], btags[3], btags[4], btags[5], ctags[0], ctags[1], ctags[2], ctags[3], ctags[4], ctags[5]); + _ntev->Fill(_nev, nbz, nbh, btags[0], btags[1], btags[2], btags[3], btags[4], btags[5], ctags[0], ctags[1], ctags[2], ctags[3], ctags[4], ctags[5]); _nev ++; } @@ -1039,6 +1086,44 @@ void FlavtagReader::end() { _file->Close(); } +void WeaverReader::init(Parameters* param) { + Algorithm::init(param); + + string filename = param->get("WeaverReader.FileName",string("test.root")); + _jetname = param->get("WeaverReader.JetCollectionName",string("RefinedJets")); + string primvtxcolname = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + Event::Instance()->setDefaultPrimaryVertex(primvtxcolname.c_str()); + + _file = new TFile(filename.c_str(),"RECREATE"); + _nt = new TNtupleD("nt","nt","nev:nj:e:px:py:pz:mc_b:mc_c:mc_s:mc_u:mc_d:mc_g"); + //_ntev = new TNtupleD("ntev","ntev","nev:btag1:btag2:btag3:btag4:btag5:btag6:ctag1:ctag2:ctag3:ctag4:ctag5:ctag6"); + + _jets = 0; + _nev = 0; +} + +void WeaverReader::process() { + if (!_jets) { + Event::Instance()->Get(_jetname.c_str(), _jets); + } + //const Vertex * privtx = Event::Instance()->getPrimaryVertex(); + + for (unsigned int nj = 0; nj < _jets->size(); nj++) { + const Jet* j = (*_jets)[nj]; + + const Parameters* para = j->getParam("weaver"); + _nt->Fill(_nev, nj, j->E(), j->Px(), j->Py(), j->Pz(), + para->get("mc_b"), para->get("mc_c"), para->get("mc_s"), para->get("mc_u"), para->get("mc_d"), para->get("mc_g")); + } + + _nev ++; +} + +void WeaverReader::end() { + _file->Write(); + _file->Close(); +} + #if 0 void TestAlgo::init(Parameters* param) { @@ -1411,14 +1496,153 @@ void TestAlgo::end() { } #endif +void TrackPairTree::init(Parameters* param) { + Algorithm::init(param); + _privtxname = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + + _jetname = param->get("TrackPairTree.JetCollectionName",string("RefinedJets")); + string filename = param->get("TrackPairTree.FileName",string("")); + + if(filename != "") + _file = new TFile(filename.c_str(),"RECREATE"); + else + _file = 0; + + _trackTree = new TTree("TrackTree","TrackTree"); + _trackTree->Branch("nev",&_dataTrack.nev,"nev/I"); + _trackTree->Branch("njet",&_dataTrack.njet,"njet/I"); + _trackTree->Branch("idx",&_dataTrack.idx,"idx/I"); + _trackTree->Branch("d0",&_dataTrack.d0,"d0/F"); + _trackTree->Branch("phi",&_dataTrack.phi,"phi/F"); + _trackTree->Branch("omega",&_dataTrack.omega,"omega/F"); + _trackTree->Branch("z0",&_dataTrack.z0,"z0/F"); + _trackTree->Branch("tanLambda",&_dataTrack.tanLambda,"tanLambda/F"); + _trackTree->Branch("covMatrix",_dataTrack.covMatrix,"covMatrix[15]/F"); + _trackTree->Branch("chi2",&_dataTrack.chi2,"chi2/F"); + _trackTree->Branch("ndf",&_dataTrack.ndf,"ndf/I"); + + _vtxTree = new TTree("VtxTree","VtxTree"); + _vtxTree->Branch("nev",&_dataVtx.nev,"nev/I"); + _vtxTree->Branch("njet",&_dataVtx.njet,"njet/I"); + _vtxTree->Branch("idx",&_dataVtx.idx,"idx/I"); + _vtxTree->Branch("tracks",_dataVtx.tracks,"tracks[2]/I"); + _vtxTree->Branch("pos",_dataVtx.pos,"pos[3]/F"); + _vtxTree->Branch("covMatrix",_dataVtx.covMatrix,"covMatrix[6]/F"); + _vtxTree->Branch("chi2",&_dataVtx.chi2,"chi2/F"); + _vtxTree->Branch("prob",&_dataVtx.prob,"prob/F"); + _vtxTree->Branch("trueVtx",&_dataVtx.trueVtx, "trueVtx/I"); + + _nev = 0; +} + +void TrackPairTree::process(){ + //TrackVec& tracks = Event::Instance()->getTracks(); + + JetVec *pjet = 0; + if (_jetname != "") + Event::Instance()->Get(_jetname.c_str(), pjet); + + cout << "jetname " << _jetname << endl; + + if(pjet == 0)return; + + cout << "jet size " << pjet->size() << endl; + + for(unsigned int nj=0;njsize();nj++){ + const Jet *jet = (*pjet)[nj]; + TrackVec &tracks = jet->getAllTracks(); + + for(unsigned int i=0;igetD0(); + _dataTrack.phi = tr->getPhi(); + _dataTrack.omega = tr->getOmega(); + _dataTrack.z0 = tr->getZ0(); + _dataTrack.tanLambda = tr->getTanLambda(); + for(int j=0;j<15;j++) + _dataTrack.covMatrix[j] = tr->getCovMatrix()[j]; + _dataTrack.chi2 = tr->getChi2(); + _dataTrack.ndf = tr->getNdf(); + + _trackTree->Fill(); + } + + if(tracks.size()>1){ + int n = 0; + for(unsigned int i=0;i trvec; + trvec.push_back(tracks[i]); + trvec.push_back(tracks[j]); + Vertex *vtx = VertexFitterSimple_V() (trvec.begin(), trvec.end(),0); + + _dataVtx.nev = _nev; + _dataVtx.njet = nj; + _dataVtx.idx = n++; + _dataVtx.tracks[0] = i; + _dataVtx.tracks[1] = j; + _dataVtx.pos[0] = vtx->getX(); + _dataVtx.pos[1] = vtx->getY(); + _dataVtx.pos[2] = vtx->getZ(); + for(int k=0;k<6;k++) + _dataVtx.covMatrix[k] = vtx->getCov()[k]; + _dataVtx.chi2 = vtx->getChi2(); + _dataVtx.prob = vtx->getProb(); + _dataVtx.trueVtx = 0; // no assignment + const MCParticle *mcp[2]; + const MCParticle *mcpp[2]; + if((mcp[0] = tracks[i]->getMcp()) && (mcp[1] = tracks[j]->getMcp())){ + if((mcp[0]->getSemiStableParent()) == 0 && (mcp[1]->getSemiStableParent()) == 0) + _dataVtx.trueVtx = 1; // primary vertex + else if((mcpp[0] = mcp[0]->getSemiStableParent()) == (mcpp[1] = mcp[1]->getSemiStableParent())){ // parent is same + if(mcpp[0]->isSemiStableB()) _dataVtx.trueVtx = 2; // b, same + else if(mcpp[0]->isSemiStableC()) _dataVtx.trueVtx = 3; // c, same + else if(mcpp[0]->getId() <= 4) _dataVtx.trueVtx = 1; // primary, same (temporary solution) + else _dataVtx.trueVtx = 4; // other, same + + cout << mcp[0]->getPDG() << " " << mcp[1]->getPDG() << " " << mcpp[0]->getPDG() << " " << mcpp[0]->E() << " " << _dataVtx.trueVtx << endl; + } + else if(mcp[0]->getSemiStableCParent() && mcp[0]->getSemiStableCParent() == mcp[1]->getSemiStableCParent() && mcp[0]->getSemiStableCParent()->isSemiStableC()) + _dataVtx.trueVtx = 13; // c, different + else if(mcp[0]->getSemiStableBParent() && mcp[0]->getSemiStableBParent() == mcp[1]->getSemiStableBParent() && mcp[0]->getSemiStableBParent()->isSemiStableB()) + _dataVtx.trueVtx = 12; // b, different + else _dataVtx.trueVtx = 0; // completely different + } + + _vtxTree->Fill(); + + delete vtx; + } + } + } + } + + _nev ++; +} + +void TrackPairTree::end(){ + if(_file){ + _file->Write(); + _file->Close(); + }else{ + _vtxTree->Write(); + _trackTree->Write(); + } +} + + void VertexAnalysis::init(Parameters* param) { Algorithm::init(param); _privtxname = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); string filename = param->get("VertexAnalysis.FileName",string("VertexAnalysis.root")); - _secvtxname = param->get("VertexAnalysis.SecondaryVertexCollectionName",string("BuildUpVertex")); + _secvtxname = param->get("VertexAnalysis.SecondaryVertexCollectionName",string("RefinedVertex")); + _jetname = param->get("VertexAnalysis.JetCollectionName",string("RefinedJets")); _file = new TFile(filename.c_str(),"RECREATE"); - _nt = new TNtupleD("vtxtree","vtxtree","track:invtx:d0:d0err:z0:z0err:e:pt:chi2:ndf:vtxftdhits"); + _nt = new TNtupleD("vtxtree","vtxtree","track:invtx:d0:d0err:z0:z0err:e:pt:chi2:ndf:vtxftdhits:parent:bhadron:bvtxdist:blepton:chadron:cvtxdist:clepton:recvtxdist:recvtxdist2"); } void VertexAnalysis::process() { @@ -1429,45 +1653,145 @@ void VertexAnalysis::process() { if (_secvtxname != "") Event::Instance()->Get(_secvtxname.c_str(), psecvtx); + JetVec *pjet = 0; + if (_jetname != "") + Event::Instance()->Get(_jetname.c_str(), pjet); + if (mcps.size() == 0) { cout << "MCParticle collection not specified. We need the MCParticle collection for vertex analysis." << endl; return; } + + if(pjet){ + for(unsigned int j=0;jsize();j++){ + const Jet *jet = (*pjet)[j]; + VertexVec &v = jet->getVertices(); + cout << "Jet " << j << " has " << v.size() << " vertices. Direction (" << jet->Px() << ", " << jet->Py() << ", " << jet->Pz() << ")" << endl; + + for(unsigned int k=0;kgetTracks(); + cout << "Vertex " << k << " has " << vtrs.size() << " tracks. position (" << v[k]->getX() << ", " << v[k]->getY() << ", " << v[k]->getZ() << ")" << endl; + + for(unsigned int l=0;lgetMcp(); + if(!vmcp){ + cout << "Track " << l << " cannot be associated with MCP. Energy = " << vtr->E() << endl; + }else{ + const MCParticle *parb, *parc, *par; + parb = vmcp->getSemiStableBParent(); + parc = vmcp->getSemiStableCParent(); + par = vmcp->getSemiStableParent(); + cout << "Track " << l << ": PDG " << vmcp->getPDG() << ", B parent " << (parb ? parb->getPDG() : 0) + << ", C parent " << (parc ? parc->getPDG() : 0) << ", parent " << (par ? par->getPDG() : 0) << endl; + cout << "Start point (" << vmcp->getVertex().x() << ", " << vmcp->getVertex().y() << ", " << vmcp->getVertex().z() << ")" << endl; + } + } + } + } + } + + for (unsigned int i=0; i < tracks.size(); ++i) { const Track* tr = tracks[i]; const MCParticle* mcp = tr->getMcp(); - double trackseed, invtx, d0, d0err, z0, z0err, e, pt, chi2, ndf, vtxftdhits; + struct { + double trackseed; + double invtx; + double d0; + double d0err; + double z0; + double z0err; + double e; + double pt; + double chi2; + double ndf; + double vtxftdhits; + double parent; + double bhadron; + double bvtxdist; + double blepton; + double chadron; + double cvtxdist; + double clepton; + double recvtxdist; + double recvtxdist2; + } d; + + memset(&d,0,sizeof(d)); + + if(pjet){ + bool found = false; + for(unsigned int j=0;jsize();j++){ + TrackVec& vtr = (*pjet)[j]->getAllTracks(); + if (find(vtr.begin(), vtr.end(), tr) != vtr.end()){ + const Jet *jet = (*pjet)[j]; + TVector3 pospri = privtx->getPos(); + if(jet->getVertices().size() > 0){ + TVector3 posvtx = jet->getVertices()[0]->getPos(); + d.recvtxdist = (posvtx - pospri).Mag(); + //cout << "privtx " << pospri.Mag() << " sec " << posvtx.Mag() << " dist " << d.recvtxdist << endl; + } + if(jet->getVertices().size() > 1){ + TVector3 posvtx = jet->getVertices()[1]->getPos(); + d.recvtxdist2 = (posvtx - pospri).Mag(); + } + found = true; + break; + } + } + if(!found)continue; + } - if (mcp->getSemiStableParent() == 0) trackseed = 0.; - else if (mcp->getSemiStableBParent() != 0) trackseed = 1.; - else if (mcp->getSemiStableCParent() != 0) trackseed = 2.; - else trackseed = 3.; - invtx = (find(privtx->getTracks().begin(), privtx->getTracks().end(), tr) != privtx->getTracks().end()); - if (invtx == 0. && psecvtx) { + const MCParticle *mcd; + // if (mcp->getSemiStableParent() == 0) d.trackseed = 0.; + if (mcp->getSemiStableParent() == 0 || (mcp->getVertex().x() == 0 && mcp->getVertex().y() == 0)) d.trackseed = 0.; + else{ + d.parent = mcp->getSemiStableParent()->getPDG(); + d.trackseed = 3.; + if ((mcd = mcp->getSemiStableBParent()) != 0){ + d.trackseed = 1.; + d.bhadron = mcd->getPDG(); + d.bvtxdist = mcd->decayDistance(); + const MCParticle *mcdl = mcd->semileptonicDecay(); + d.blepton = (mcdl ? mcdl->getPDG() : 0); + } + if ((mcd = mcp->getSemiStableCParent()) != 0){ + d.trackseed = 2.; + d.chadron = mcd->getPDG(); + d.cvtxdist = mcd->decayDistance(); + const MCParticle *mcdl = mcd->semileptonicDecay(); + d.clepton = (mcdl ? mcdl->getPDG() : 0); + } + } + + d.invtx = (find(privtx->getTracks().begin(), privtx->getTracks().end(), tr) != privtx->getTracks().end()); + if (d.invtx == 0. && psecvtx) { // looking for secondary vertices for (unsigned int j=0; jsize(); j++) { TrackVec& vtr = (*psecvtx)[j]->getTracks(); if (find(vtr.begin(), vtr.end(), tr) != vtr.end()) - invtx = 2.; + d.invtx = 2.; } } - d0 = tr->getD0(); - d0err = sqrt(tr->getCovMatrix()[tpar::d0d0]); - z0 = tr->getZ0(); - z0err = sqrt(tr->getCovMatrix()[tpar::z0z0]); + d.d0 = tr->getD0(); + d.d0err = sqrt(tr->getCovMatrix()[tpar::d0d0]); + d.z0 = tr->getZ0(); + d.z0err = sqrt(tr->getCovMatrix()[tpar::z0z0]); + + d.e = tr->E(); + d.pt = tr->Pt(); - e = tr->E(); - pt = tr->Pt(); + d.chi2 = tr->getChi2(); + d.ndf = tr->getNdf(); + d.vtxftdhits = tr->getVtxHits() + tr->getFtdHits(); - chi2 = tr->getChi2(); - ndf = tr->getNdf(); - vtxftdhits = tr->getVtxHits() + tr->getFtdHits(); + _nt->Fill((double *)&d); - _nt->Fill(trackseed, invtx, d0, d0err, z0, z0err, e, pt, chi2, ndf, vtxftdhits); } } diff --git a/src/vertexfinderdnn.cc b/src/vertexfinderdnn.cc new file mode 100644 index 0000000..b3302f8 --- /dev/null +++ b/src/vertexfinderdnn.cc @@ -0,0 +1,485 @@ +#include + +#include "TFile.h" +#include "TROOT.h" +#include "TNtuple.h" +#include "TTree.h" +#include "TNtupleD.h" +#include "TSystem.h" +#include "TPad.h" +#include "TStyle.h" +#include "TVector3.h" +#include "TLorentzVector.h" +#include "TMatrixFSym.h" +#include +#include "TMatrixDSym.h" +#include "HelixClass.h" +#include +#include "TVectorD.h" +#include "TMethodCall.h" + +#include "lcfiplus.h" +#include "process.h" +#include "geometry.h" +#include "vertexfinderdnn.h" +#include "VertexSelector.h" +#include "algoEtc.h" +#include "algoSigProb.h" +#include "VertexFinderSuehara.h" +#include "TrackSelector.h" +#include "VertexFitterSimple.h" +#include "TrackNtuple.h" +#include "EventStore.h" +#include "LcfiplusProcessor.h" +#include + +#include "math.h" + + +using namespace lcfiplus; +using namespace lcfiplus::algoSigProb; +using namespace lcfiplus::algoEtc; + +using namespace lcfiplus; + +namespace lcfiplus { + +// ================================================================================================================ // +// ================================================================================================================ // +void VertexFinderDNNccbar::init(Parameters* param) { + Algorithm::init(param); + + string filename = param->get("VertexFinderDNNccbar.RootFileName", string("my_track1.root")); + string primvtxcolname = param->get("PrimaryVertexCollectionName",string("PrimaryVertex")); + Event::Instance()->setDefaultPrimaryVertex(primvtxcolname.c_str()); + + _file = new TFile(filename.c_str(), "RECREATE"); + + _ntp = new TTree("track0", "track0"); + + nEvt = 0; + ntr1Trk = 0; + ntr2Trk = 0; + + TracksData& d = _data; + _ntp->Branch("nevent", &nEvt, "nevent/I"); // Event number + _ntp->Branch("ntr1track", &ntr1Trk, "ntr1track/I"); // Event number + _ntp->Branch("ntr2track", &ntr2Trk, "ntr2track/I"); // Event number + // Track1 low data + _ntp->Branch("tr1d0", &d.tr1d0, "tr1d0/D"); // Impact paramter of the track in (r-phi) + _ntp->Branch("tr1z0", &d.tr1z0, "tr1z0/D"); // mpact paramter of the track in (r-z) + _ntp->Branch("tr1phi", &d.tr1phi, "tr1phi/D"); // Phi of the track at the reference point + _ntp->Branch("tr1omega", &d.tr1omega, "tr1omega/D"); // Omega is the signed curvature of the track in [1/mm] + _ntp->Branch("tr1tanlam", &d.tr1tanlam, "tr1tanlam/D"); // Lambda is the dip angle of the track in r-z at the reference point + _ntp->Branch("tr1x", &d.tr1x, "tr1x/D"); + _ntp->Branch("tr1y", &d.tr1y, "tr1y/D"); + _ntp->Branch("tr1z", &d.tr1z, "tr1z/D"); + _ntp->Branch("tr1charge", &d.tr1charge, "tr1charge/I"); + _ntp->Branch("tr1covmatrixd0d0", &d.tr1covmatrixd0d0, "tr1covmatrixd0d0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixd0z0", &d.tr1covmatrixd0z0, "tr1covmatrixd0z0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixd0ph", &d.tr1covmatrixd0ph, "tr1covmatrixd0ph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixd0om", &d.tr1covmatrixd0om, "tr1covmatrixd0om/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixd0tl", &d.tr1covmatrixd0tl, "tr1covmatrixd0tl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixz0z0", &d.tr1covmatrixz0z0, "tr1covmatrixz0z0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixz0ph", &d.tr1covmatrixz0ph, "tr1covmatrixz0ph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixz0om", &d.tr1covmatrixz0om, "tr1covmatrixz0om/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixz0tl", &d.tr1covmatrixz0tl, "tr1covmatrixz0tl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixphph", &d.tr1covmatrixphph, "tr1covmatrixphph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixphom", &d.tr1covmatrixphom, "tr1covmatrixphom/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixphtl", &d.tr1covmatrixphtl, "tr1covmatrixphtl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixomom", &d.tr1covmatrixomom, "tr1covmatrixomom/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixomtl", &d.tr1covmatrixomtl, "tr1covmatrixomtl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr1covmatrixtltl", &d.tr1covmatrixtltl, "tr1covmatrixtltl/D"); // Covariance matrix of the track parameters + // ======================================================= // + // Track2 low data + _ntp->Branch("tr2d0", &d.tr2d0, "tr2d0/D"); // Impact paramter of the track in (r-phi) + _ntp->Branch("tr2z0", &d.tr2z0, "tr2z0/D"); // mpact paramter of the track in (r-z) + _ntp->Branch("tr2phi", &d.tr2phi, "tr2phi/D"); // Phi of the track at the reference point + _ntp->Branch("tr2omega", &d.tr2omega, "tr2omega/D"); // Omega is the signed curvature of the track in [1/mm] + _ntp->Branch("tr2tanlam", &d.tr2tanlam, "tr2tanlam/D"); // Lambda is the dip angle of the track in r-z at the reference point + _ntp->Branch("tr2x", &d.tr2x, "tr2x/D"); + _ntp->Branch("tr2y", &d.tr2y, "tr2y/D"); + _ntp->Branch("tr2z", &d.tr2z, "tr2z/D"); + _ntp->Branch("tr2charge", &d.tr2charge, "tr2charge/I"); + _ntp->Branch("tr2covmatrixd0d0", &d.tr2covmatrixd0d0, "tr2covmatrixd0d0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixd0z0", &d.tr2covmatrixd0z0, "tr2covmatrixd0z0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixd0ph", &d.tr2covmatrixd0ph, "tr2covmatrixd0ph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixd0om", &d.tr2covmatrixd0om, "tr2covmatrixd0om/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixd0tl", &d.tr2covmatrixd0tl, "tr2covmatrixd0tl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixz0z0", &d.tr2covmatrixz0z0, "tr2covmatrixz0z0/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixz0ph", &d.tr2covmatrixz0ph, "tr2covmatrixz0ph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixz0om", &d.tr2covmatrixz0om, "tr2covmatrixz0om/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixz0tl", &d.tr2covmatrixz0tl, "tr2covmatrixz0tl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixphph", &d.tr2covmatrixphph, "tr2covmatrixphph/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixphom", &d.tr2covmatrixphom, "tr2covmatrixphom/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixphtl", &d.tr2covmatrixphtl, "tr2covmatrixphtl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixomom", &d.tr2covmatrixomom, "tr2covmatrixomom/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixomtl", &d.tr2covmatrixomtl, "tr2covmatrixomtl/D"); // Covariance matrix of the track parameters + _ntp->Branch("tr2covmatrixtltl", &d.tr2covmatrixtltl, "tr2covmatrixtltl/D"); // Covariance matrix of the track parameters + // ======================================================= // + // Track1 MC particle data + _ntp->Branch("tr1mcx", &d.tr1mcx, "tr1mcx/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr1mcy", &d.tr1mcy, "tr1mcy/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr1mcz", &d.tr1mcz, "tr1mcz/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr1id", &d.tr1id, "tr1id/I"); // the number of parents of this particle - 0 if mother + _ntp->Branch("tr1pdg", &d.tr1pdg, "tr1pdg/I"); // the number of parents of this particle - 0 if mother + _ntp->Branch("tr1sscid", &tr1sscid, "tr1sscid/I"); // other : 0, semistable c : xxx, primary : 2 + _ntp->Branch("tr1sscpdg", &tr1sscpdg, "tr1sscpdg/I"); // other : 0, semistable c : 1, primary : 2 + _ntp->Branch("tr1ssc", &tr1ssc, "tr1ssc/I"); // other : 0, semistable c : 1, primary : 2 + // ======================================================= // + // Track2 MC particle data + _ntp->Branch("tr2mcx", &d.tr2mcx, "tr2mcx/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr2mcy", &d.tr2mcy, "tr2mcy/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr2mcz", &d.tr2mcz, "tr2mcz/D"); // the production vertex of the particle in [mm]. + _ntp->Branch("tr2id", &d.tr2id, "tr2id/I"); // the number of parents of this particle - 0 if mother + _ntp->Branch("tr2pdg", &d.tr2pdg, "tr2pdg/I"); // the number of parents of this particle - 0 if mother + _ntp->Branch("tr2sscid", &tr2sscid, "tr2sscid/I"); // semistable exist : 1 other : 0 + _ntp->Branch("tr2sscpdg", &tr2sscpdg, "tr2sscpdg/I"); // semistable exist : 1 other : 0 + _ntp->Branch("tr2ssc", &tr2ssc, "tr2ssc/I"); // semistable exist : 1 other : 0 + // ======================================================= // + _ntp->Branch("coSemiStableC", &d.cosemistablec, "coSemiStableC/I"); // cosemistableC is exist : 1 other : 0 + _ntp->Branch("connect", &d.connect, "connect/I"); // 0 : not connect, 1 : semistablecparent, 2 : primary vertex + _ntp->Branch("lcfiplustag", &d.lcfiplustag, "lcfiplustag/I"); // 0 : not connect, 1 : semistablecparent, 2 : primary vertex + // ======================================================= // + // Parameters From Vertex Finder Suehara + /* + _ntp->Branch("tr1tlvx", &d.tr1tlvx, "tr1tlvx/I"); // TLotentzVector track 1 x + _ntp->Branch("tr1tlvy", &d.tr1tlvy, "tr1tlvy/I"); // TLotentzVector track 1 y + _ntp->Branch("tr1tlvz", &d.tr1tlvz, "tr1tlvz/I"); // TLotentzVector track 1 z + _ntp->Branch("tr2tlvx", &d.tr2tlvx, "tr2tlvx/I"); // TLotentzVector track 2 x + _ntp->Branch("tr2tlvy", &d.tr2tlvy, "tr2tlvy/I"); // TLotentzVector track 2 y + _ntp->Branch("tr2tlvz", &d.tr2tlvz, "tr2tlvz/I"); // TLotentzVector track 2 z + _ntp->Branch("l0mass", &d.l0mass, "l0mass/I"); // l0 mass + _ntp->Branch("minE", &d.minE, "minE/I"); // min E + */ + _ntp->Branch("chi2", &d.chi2, "chi2/D"); // chi2 + _ntp->Branch("vchi2", &d.vchi2, "vchi2/D"); // vertex chi2 + _ntp->Branch("vprob", &d.vprob, "vprob/D"); // vertex probability + _ntp->Branch("vposx", &d.vposx, "vposx/D"); // vertex position x + _ntp->Branch("vposy", &d.vposy, "vposy/D"); // vertex position y + _ntp->Branch("vposz", &d.vposz, "vposz/D"); // vertex position z + _ntp->Branch("mass", &d.mass, "mass/D"); // mass + + // track parmeter notation + for(int i=0;i<10;i++){ + for(int j=0;j<3;j++){ + _ntp->Branch(TString::Format("tr1_%d_%d",i,j), &d.tr1xyz[i][j], TString::Format("tr1_%d_%d/D",i,j)); + _ntp->Branch(TString::Format("tr2_%d_%d",i,j), &d.tr2xyz[i][j], TString::Format("tr2_%d_%d/D",i,j)); + } + } +} + +void VertexFinderDNNccbar::process() { + + Event* event = Event::Instance(); + const TrackVec& track_list = event->getTracks(); + + for (unsigned int i=0; i < track_list.size(); ++i) { + ntr1Trk = i; + const Track* track1 = track_list[i]; + if (!track1) continue; + const MCParticle* mcpc1 = event->getMCParticle(track1); + _hel1 = new lcfiplus::Helix(track1, PointBase::NOTUSED); + + _data = TracksData(); + + _data.tr1d0 = track1->getD0(); + _data.tr1z0 = track1->getZ0(); + _data.tr1phi = track1->getPhi(); + _data.tr1omega = track1->getOmega(); + _data.tr1tanlam = track1->getTanLambda(); + + _data.tr1charge = (int)track1->getCharge(); + + _data.tr1x = _hel1->GetPos(0.).X(); + _data.tr1y = _hel1->GetPos(0.).Y(); + _data.tr1z = _hel1->GetPos(0.).Z(); + + /* + TVector3 mom1 = track1->Vect(); + TLorentzVector tr1tlv; + tr1tlv.SetVectM(mom1, 0.1396); + + _data.tr1tlvx = tr1tlv.X(); + _data.tr1tlvy = tr1tlv.Y(); + _data.tr1tlvz = tr1tlv.Z(); + */ + + const double* cov1 = track1->getCovMatrix(); + + _data.tr1covmatrixd0d0 = cov1[0]; // d0d0 + _data.tr1covmatrixd0z0 = cov1[1]; // d0z0 + _data.tr1covmatrixd0ph = cov1[2]; // d0ph + _data.tr1covmatrixd0om = cov1[3]; // d0om + _data.tr1covmatrixd0tl = cov1[4]; // d0tl + _data.tr1covmatrixz0z0 = cov1[5]; // z0z0 + _data.tr1covmatrixz0ph = cov1[6]; // z0ph + _data.tr1covmatrixz0om = cov1[7]; // z0om + _data.tr1covmatrixz0tl = cov1[8]; // z0tl + _data.tr1covmatrixphph = cov1[9]; // phph + _data.tr1covmatrixphom = cov1[10]; // phom + _data.tr1covmatrixphtl = cov1[11]; // phtl + _data.tr1covmatrixomom = cov1[12]; // omom + _data.tr1covmatrixomtl = cov1[13]; // omtl + _data.tr1covmatrixtltl = cov1[14]; // tltl + + for(int i=0;i<10;i++){ + track1->setFlightLength(i-5); + _data.tr1xyz[i][0] = track1->getX(); + _data.tr1xyz[i][1] = track1->getY(); + _data.tr1xyz[i][2] = track1->getZ(); + //std::cout << "t = " << i << ", x = " << _data.tr1xyz[i][0] << ", y = " << _data.tr1xyz[i][1] << ", z = " << _data.tr1xyz[i][2] << std::endl; + } + + if (mcpc1){ + _data.tr1mcx = mcpc1->getVertex().X(); + _data.tr1mcy = mcpc1->getVertex().Y(); + _data.tr1mcz = mcpc1->getVertex().Z(); + _data.tr1id = mcpc1->getId(); + _data.tr1pdg = mcpc1->getPDG(); + + const MCParticle* ssmcpc1 = mcpc1->getSemiStableParent(); + + if (ssmcpc1==0){ + tr1ssc = 2; + tr1sscid = 0; + tr1sscid = 0; + } + if (ssmcpc1){ + tr1sscid = ssmcpc1->getId(); + tr1sscpdg = ssmcpc1->getPDG(); + if (ssmcpc1->getParent()->getId()==0){ + tr1ssc = 2; + } + else if (ssmcpc1->isSemiStableC()){ + tr1ssc = 1; + } + else { + tr1ssc = 0; + } + cout << "SSC : " << tr1ssc << ", ID : " << ssmcpc1->getId() << ", PDG : " << ssmcpc1->getPDG() << + ", Parent : " << ssmcpc1->getParent()->getId() << endl; + }else{ + _data.tr1mcx = 0.; + _data.tr1mcy = 0.; + _data.tr1mcz = 0.; + _data.tr1id = 0.; + _data.tr1pdg = 0.; + tr1ssc = -1; + tr1sscid = 0; + tr1sscpdg = 0; + } + }else{ + _data.tr1mcx = 0.; + _data.tr1mcy = 0.; + _data.tr1mcz = 0.; + _data.tr1id = 0.; + _data.tr1pdg = 0.; + tr1ssc = -1; + tr1sscid = 0; + tr1sscpdg = 0; + } + + for (unsigned int j=0; j < i; ++j) { + ntr2Trk = j; + const Track* track2 = track_list[j]; + if (!track2) continue; + + const MCParticle* mcpc2 = event->getMCParticle(track2); + _hel2 = new lcfiplus::Helix(track2, PointBase::NOTUSED); + + _data.tr2d0 = track2->getD0(); + _data.tr2z0 = track2->getZ0(); + _data.tr2phi = track2->getPhi(); + _data.tr2omega = track2->getOmega(); + _data.tr2tanlam = track2->getTanLambda(); + + _data.tr2charge = (int)track2->getCharge(); + + _data.tr2x = _hel2->GetPos(0.).X(); + _data.tr2y = _hel2->GetPos(0.).Y(); + _data.tr2z = _hel2->GetPos(0.).Z(); + + /* + TVector3 mom2 = track2->Vect(); + TLorentzVector tr2tlv; + tr2tlv.SetVectM(mom2, 0.1396); + + _data.tr2tlvx = tr2tlv.X(); + _data.tr2tlvy = tr2tlv.Y(); + _data.tr2tlvz = tr2tlv.Z(); + */ + + const double* cov2 = track2->getCovMatrix(); + + _data.tr2covmatrixd0d0 = cov2[0]; // d0d0 + _data.tr2covmatrixd0z0 = cov2[1]; // d0z0 + _data.tr2covmatrixd0ph = cov2[2]; // d0ph + _data.tr2covmatrixd0om = cov2[3]; // d0om + _data.tr2covmatrixd0tl = cov2[4]; // d0tl + _data.tr2covmatrixz0z0 = cov2[5]; // z0z0 + _data.tr2covmatrixz0ph = cov2[6]; // z0ph + _data.tr2covmatrixz0om = cov2[7]; // z0om + _data.tr2covmatrixz0tl = cov2[8]; // z0tl + _data.tr2covmatrixphph = cov2[9]; // phph + _data.tr2covmatrixphom = cov2[10]; // phom + _data.tr2covmatrixphtl = cov2[11]; // phtl + _data.tr2covmatrixomom = cov2[12]; // omom + _data.tr2covmatrixomtl = cov2[13]; // omtl + _data.tr2covmatrixtltl = cov2[14]; // tltl + + for(int i=0;i<10;i++){ + track2->setFlightLength(i-5); + _data.tr2xyz[i][0] = track2->getX(); + _data.tr2xyz[i][1] = track2->getY(); + _data.tr2xyz[i][2] = track2->getZ(); + //std::cout << "t = " << i << ", x = " << _data.tr2xyz[i][0] << ", y = " << _data.tr2xyz[i][1] << ", z = " << _data.tr2xyz[i][2] << std::endl; + } + + if (mcpc2){ + _data.tr2mcx = mcpc2->getVertex().X(); + _data.tr2mcy = mcpc2->getVertex().Y(); + _data.tr2mcz = mcpc2->getVertex().Z(); + _data.tr2id = mcpc2->getId(); + _data.tr2pdg = mcpc2->getPDG(); + + const MCParticle* ssmcpc2 = mcpc2->getSemiStableParent(); + + if (ssmcpc2==0){ + tr2ssc = 2; + tr2sscid = 0; + tr2sscid = 0; + } + if (ssmcpc2){ + tr2sscid = ssmcpc2->getId(); + tr2sscpdg = ssmcpc2->getPDG(); + if (ssmcpc2->getParent()->getId()==0){ + tr2ssc = 2; + } + else if (ssmcpc2->isSemiStableC()){ + tr2ssc = 1; + } + else { + tr2ssc = 0; + } + }else{ + _data.tr2mcx = 0.; + _data.tr2mcy = 0.; + _data.tr2mcz = 0.; + _data.tr2id = 0.; + _data.tr2pdg = 0.; + tr2ssc = -1; + tr2sscid = 0; + tr2sscpdg = 0; + } + }else{ + _data.tr2mcx = 0.; + _data.tr2mcy = 0.; + _data.tr2mcz = 0.; + _data.tr2id = 0.; + _data.tr2pdg = 0.; + tr2ssc = -1; + tr2sscid = 0; + tr2sscpdg = 0; + } + + if ((tr1ssc == -1) || (tr2ssc == -1))_data.connect = -1; + else if (((tr1ssc==1)&&(tr2ssc==1))&&(tr1sscid==tr2sscid)) _data.connect = 1; + else if ((tr1ssc==2)&&(tr2ssc==2)) _data.connect = 2; + else _data.connect = 0; + + /* + + TLorentzVector protonForLambda, pionForLambda; + if(mom1.Mag() > mom2.Mag()){ + protonForLambda.SetVectM( mom1, 0.9383 ); + pionForLambda.SetVectM( mom2, 0.1396 ); + } else { + protonForLambda.SetVectM( mom2, 0.9383 ); + pionForLambda.SetVectM( mom1, 0.1396 ); + } + _data.l0mass = (protonForLambda+pionForLambda).M(); + _data.minE = min(tr1tlv.E(), tr2tlv.E()); + */ + + vector vttmp; + vttmp.push_back(track1); + vttmp.push_back(track2); + + Vertex* vtx = VertexFitterSimple_V() (vttmp.begin(), vttmp.end(), 0); + + _data.chi2 = max(vtx->getChi2Track(track1), vtx->getChi2Track(track2)); + _data.vchi2 = vtx->getChi2(); + _data.vprob = vtx->getProb(); + /* + cout << "max(vtx->getChi2Track()) " << _data.chi2 << " : " << max(vtx->getChi2Track(track1), vtx->getChi2Track(track2)) << + " : vtx->getChi2 " << _data.vchi2 << " : " << vtx->getChi2() << endl; + */ + _data.vposx = vtx->getPos().X(); + _data.vposy = vtx->getPos().Y(); + _data.vposz = vtx->getPos().Z(); + + TLorentzVector tr1tlv = *track1; + TLorentzVector tr2tlv = *track2; + + _data.mass = (tr1tlv+tr2tlv).M(); + + + // suehara 200319 + // TrackSelector and VertexSelector config + TrackSelectorConfig cfgtrack; + cfgtrack.maxD0 = 10.; + cfgtrack.maxZ0 = 20.; + cfgtrack.minPt = 0.1; + cfgtrack.maxInnermostHitRadius = 1e10; + + cfgtrack.maxD0Err = 0.1; + cfgtrack.maxZ0Err = 0.1; + + cfgtrack.minTpcHits = 20; + cfgtrack.minFtdHits = 3; + cfgtrack.minVtxHits = 3; + cfgtrack.minVtxPlusFtdHits = 1; + + VertexFinderSuehara::VertexFinderSueharaConfig cfgvtx; + // chi2 threshold = 9. + // mass threshold = 10. + // minimum distance = 0.3; + + TrackSelector sel; + + _data.lcfiplustag = 0; // 0: NC, 1: secondary, 2: primary + const Vertex *ip = event->getPrimaryVertex(); + TrackVec &iptracks = ip->getTracks(); + if(std::find(iptracks.begin(), iptracks.end(), track1) != iptracks.end() + && std::find(iptracks.begin(), iptracks.end(), track2) != iptracks.end()) + _data.lcfiplustag = 2; // primary tracks + else{ + // track selection + if(sel.passesCut(track1, cfgtrack) && sel.passesCut(track2, cfgtrack)){ + // mass, chi2, vpos, v0 selection + TLorentzVector v1 = *track1; + TLorentzVector v2 = *track2; + double mass = (v1+v2).M(); + double chi2 = max(vtx->getChi2Track(track1), vtx->getChi2Track(track2)); + + if(mass < min(v1.E(), v2.E()) && mass < 10. && chi2 < 9. && vtx->getPos().Mag() > 0.3 + && vtx->getPos().Dot((v1+v2).Vect()) > 0 && !VertexSelector().passesCut(vtx, cfgvtx.v0selVertex, ip)){ + _data.lcfiplustag = 1; // secondary tracks + } + } + } + + _ntp->Fill(); + } + + } + nEvt++; +} + +void VertexFinderDNNccbar::end() { + _file->Write(); + _file->Close(); +} + +}