diff --git a/.gitignore b/.gitignore index d1458d2d..54c57605 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,5 @@ benchmark/* /profile/* /statprof/* /debug/* +ext/JetFlavourTagging/PROGRESS.md +examples/flavour-tagging/data/events_080263084.root diff --git a/Project.toml b/Project.toml index 77c12d07..95a486a9 100644 --- a/Project.toml +++ b/Project.toml @@ -18,10 +18,14 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" [weakdeps] EDM4hep = "eb32b910-dde9-4347-8fce-cd6be3498f0c" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" +ONNXRunTime = "e034b28e-924e-41b2-b98f-d2bbeb830c6a" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab" [extensions] EDM4hepJets = "EDM4hep" JetVisualisation = "Makie" +JetFlavourTagging = ["EDM4hep", "ONNXRunTime", "JSON", "PhysicalConstants"] [compat] Accessors = "0.1.36" @@ -38,9 +42,11 @@ Makie = "0.20, 0.21, 0.22, 0.23, 0.24" MuladdMacro = "0.2.4" Random = "1.9" Statistics = "1.9" +ONNXRunTime = "1.3.1" StructArrays = "0.6.18, 0.7" Test = "1.9" julia = "1.9" +PhysicalConstants = "0.2.4" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" diff --git a/examples/flavour-tagging/Project.toml b/examples/flavour-tagging/Project.toml new file mode 100644 index 00000000..202edbe4 --- /dev/null +++ b/examples/flavour-tagging/Project.toml @@ -0,0 +1,20 @@ +[deps] +Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" +CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193" +CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2" +EDM4hep = "eb32b910-dde9-4347-8fce-cd6be3498f0c" +EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" +IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +JetReconstruction = "44e8cb2c-dfab-4825-9c70-d4808a591196" +Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +LorentzVectorHEP = "f612022c-142a-473f-8cfd-a09cf3793c6c" +LorentzVectors = "3f54b04b-17fc-5cd4-9758-90c048d965e3" +MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" +ONNXRunTime = "e034b28e-924e-41b2-b98f-d2bbeb830c6a" +PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab" +StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" + +[sources] +JetReconstruction = {path = "/Users/harrywanghc/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl"} diff --git a/examples/flavour-tagging/README.md b/examples/flavour-tagging/README.md new file mode 100644 index 00000000..b2415c3f --- /dev/null +++ b/examples/flavour-tagging/README.md @@ -0,0 +1,13 @@ +# Jet Reconstruction Flavour Tagging Module Examples + +## `simple-flavour-tagging.jl` + +```bash +julia --project simple-flavour-tagging.jl +``` +This script will perform a simple flavour tagging with one events only. +It will use the `FlavourTagging` module to tag the jets in the event and print the results. + +## `simple-flavour-tagging.ipynb` + +The same as the above example, but using a Jupyter notebook. diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_fl.swp b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_fl.swp new file mode 100644 index 00000000..3432f135 Binary files /dev/null and b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_fl.swp differ diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.json b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.json new file mode 100644 index 00000000..a3e8c6bc --- /dev/null +++ b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.json @@ -0,0 +1,407 @@ +{ + "output_names": [ + "recojet_isG", + "recojet_isU", + "recojet_isS", + "recojet_isC", + "recojet_isB", + "recojet_isTAU", + "recojet_isD" + ], + "input_names": [ + "pf_points", + "pf_features", + "pf_vectors", + "pf_mask" + ], + "pf_points": { + "var_names": [ + "pfcand_thetarel", + "pfcand_phirel" + ], + "var_infos": { + "pfcand_thetarel": { + "median": 0.34896841645240784, + "norm_factor": 1.623847928023531, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phirel": { + "median": 0.00031830096850171685, + "norm_factor": 0.4692355255199169, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_features": { + "var_names": [ + "pfcand_erel_log", + "pfcand_thetarel", + "pfcand_phirel", + "pfcand_dptdpt", + "pfcand_detadeta", + "pfcand_dphidphi", + "pfcand_dxydxy", + "pfcand_dzdz", + "pfcand_dxydz", + "pfcand_dphidxy", + "pfcand_dlambdadz", + "pfcand_dxyc", + "pfcand_dxyctgtheta", + "pfcand_phic", + "pfcand_phidz", + "pfcand_phictgtheta", + "pfcand_cdz", + "pfcand_cctgtheta", + "pfcand_mtof", + "pfcand_dndx", + "pfcand_charge", + "pfcand_isMu", + "pfcand_isEl", + "pfcand_isChargedHad", + "pfcand_isGamma", + "pfcand_isNeutralHad", + "pfcand_dxy", + "pfcand_dz", + "pfcand_btagSip2dVal", + "pfcand_btagSip2dSig", + "pfcand_btagSip3dVal", + "pfcand_btagSip3dSig", + "pfcand_btagJetDistVal", + "pfcand_btagJetDistSig", + "pfcand_type" + ], + "var_infos": { + "pfcand_erel_log": { + "median": -1.8002910614013672, + "norm_factor": 1.5911575382168794, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_thetarel": { + "median": 0.34896841645240784, + "norm_factor": 1.623847928023531, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phirel": { + "median": 0.00031830096850171685, + "norm_factor": 0.4692355255199169, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dptdpt": { + "median": -9.0, + "norm_factor": 0.11111111111111001, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_detadeta": { + "median": -9.0, + "norm_factor": 0.11111111039722958, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dphidphi": { + "median": -9.0, + "norm_factor": 0.11111107873678364, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxydxy": { + "median": -9.0, + "norm_factor": 0.11110580581124187, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dzdz": { + "median": -9.0, + "norm_factor": 0.11111063977819315, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxydz": { + "median": -9.0, + "norm_factor": 0.11111103621962409, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dphidxy": { + "median": -9.0, + "norm_factor": 0.11111118663623733, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dlambdadz": { + "median": -9.0, + "norm_factor": 0.11111111196002012, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxyc": { + "median": -9.0, + "norm_factor": 0.1111111110812715, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxyctgtheta": { + "median": -9.0, + "norm_factor": 0.11111111114037472, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phic": { + "median": -9.0, + "norm_factor": 0.11111111111115293, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phidz": { + "median": -9.0, + "norm_factor": 0.11111111110696342, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phictgtheta": { + "median": -9.0, + "norm_factor": 0.11111111049504692, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_cdz": { + "median": -9.0, + "norm_factor": 0.11111111111018916, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_cctgtheta": { + "median": -9.0, + "norm_factor": 0.11111111111112591, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_mtof": { + "median": 0.07032066211104393, + "norm_factor": 5.0446882868657825, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dndx": { + "median": 0.0, + "norm_factor": 0.7165380157879775, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_charge": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isMu": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isEl": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isChargedHad": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isGamma": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isNeutralHad": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxy": { + "median": -9.0, + "norm_factor": 0.11104036349544703, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dz": { + "median": -9.0, + "norm_factor": 0.1110838280661911, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip2dVal": { + "median": -9.0, + "norm_factor": 0.11100882796219, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip2dSig": { + "median": -9.0, + "norm_factor": 0.10339552523584994, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip3dVal": { + "median": -9.0, + "norm_factor": 0.11100379369461871, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip3dSig": { + "median": -9.0, + "norm_factor": 0.10327416627313275, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagJetDistVal": { + "median": -9.0, + "norm_factor": 0.11106689226469424, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagJetDistSig": { + "median": -9.0, + "norm_factor": 0.10776558597370688, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_type": { + "median": 22.0, + "norm_factor": 0.045454545454545456, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_vectors": { + "var_names": [ + "pfcand_e", + "pfcand_p", + "pfcand_e", + "pfcand_e" + ], + "var_infos": { + "pfcand_e": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + }, + "pfcand_p": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_mask": { + "var_names": [ + "pfcand_mask" + ], + "var_infos": { + "pfcand_mask": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + } + }, + "var_length": 75 + } +} \ No newline at end of file diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.onnx b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.onnx new file mode 100644 index 00000000..0fec730e Binary files /dev/null and b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.onnx differ diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.yaml b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.yaml new file mode 100644 index 00000000..bc9fcfa2 --- /dev/null +++ b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc.yaml @@ -0,0 +1,430 @@ +treename: null +selection: null +test_time_selection: null +preprocess: + method: auto + data_fraction: 0.1 + params: + pfcand_thetarel: + length: 75 + pad_mode: wrap + center: 0.34896841645240784 + scale: 1.623847928023531 + min: -5 + max: 5 + pad_value: 0 + pfcand_phirel: + length: 75 + pad_mode: wrap + center: 0.00031830096850171685 + scale: 0.4692355255199169 + min: -5 + max: 5 + pad_value: 0 + pfcand_erel_log: + length: 75 + pad_mode: wrap + center: -1.8002910614013672 + scale: 1.5911575382168794 + min: -5 + max: 5 + pad_value: 0 + pfcand_dptdpt: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111111111001 + min: -5 + max: 5 + pad_value: 0 + pfcand_detadeta: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111039722958 + min: -5 + max: 5 + pad_value: 0 + pfcand_dphidphi: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111107873678364 + min: -5 + max: 5 + pad_value: 0 + pfcand_dxydxy: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11110580581124187 + min: -5 + max: 5 + pad_value: 0 + pfcand_dzdz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111063977819315 + min: -5 + max: 5 + pad_value: 0 + pfcand_dxydz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111103621962409 + min: -5 + max: 5 + pad_value: 0 + pfcand_dphidxy: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111118663623733 + min: -5 + max: 5 + pad_value: 0 + pfcand_dlambdadz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111196002012 + min: -5 + max: 5 + pad_value: 0 + pfcand_dxyc: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.1111111110812715 + min: -5 + max: 5 + pad_value: 0 + pfcand_dxyctgtheta: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111114037472 + min: -5 + max: 5 + pad_value: 0 + pfcand_phic: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111111115293 + min: -5 + max: 5 + pad_value: 0 + pfcand_phidz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111110696342 + min: -5 + max: 5 + pad_value: 0 + pfcand_phictgtheta: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111049504692 + min: -5 + max: 5 + pad_value: 0 + pfcand_cdz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111111018916 + min: -5 + max: 5 + pad_value: 0 + pfcand_cctgtheta: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11111111111112591 + min: -5 + max: 5 + pad_value: 0 + pfcand_mtof: + length: 75 + pad_mode: wrap + center: 0.07032066211104393 + scale: 5.0446882868657825 + min: -5 + max: 5 + pad_value: 0 + pfcand_dndx: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 0.7165380157879775 + min: -5 + max: 5 + pad_value: 0 + pfcand_charge: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_isMu: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_isEl: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_isChargedHad: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_isGamma: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_isNeutralHad: + length: 75 + pad_mode: wrap + center: 0.0 + scale: 1.0 + min: -5 + max: 5 + pad_value: 0 + pfcand_dxy: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11104036349544703 + min: -5 + max: 5 + pad_value: 0 + pfcand_dz: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.1110838280661911 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagSip2dVal: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11100882796219 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagSip2dSig: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.10339552523584994 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagSip3dVal: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11100379369461871 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagSip3dSig: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.10327416627313275 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagJetDistVal: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.11106689226469424 + min: -5 + max: 5 + pad_value: 0 + pfcand_btagJetDistSig: + length: 75 + pad_mode: wrap + center: -9.0 + scale: 0.10776558597370688 + min: -5 + max: 5 + pad_value: 0 + pfcand_type: + length: 75 + pad_mode: wrap + center: 22.0 + scale: 0.045454545454545456 + min: -5 + max: 5 + pad_value: 0 + pfcand_e: + length: 75 + pad_mode: wrap + center: null + scale: 1 + min: -5 + max: 5 + pad_value: 0 + pfcand_p: + length: 75 + pad_mode: wrap + center: null + scale: 1 + min: -5 + max: 5 + pad_value: 0 + pfcand_mask: + length: 75 + pad_mode: constant + center: null + scale: 1 + min: -5 + max: 5 + pad_value: 0 +new_variables: + pfcand_mask: ak.ones_like(pfcand_e) +inputs: + pf_points: + pad_mode: wrap + length: 75 + vars: + - pfcand_thetarel + - pfcand_phirel + pf_features: + pad_mode: wrap + length: 75 + vars: + - pfcand_erel_log + - pfcand_thetarel + - pfcand_phirel + - pfcand_dptdpt + - pfcand_detadeta + - pfcand_dphidphi + - pfcand_dxydxy + - pfcand_dzdz + - pfcand_dxydz + - pfcand_dphidxy + - pfcand_dlambdadz + - pfcand_dxyc + - pfcand_dxyctgtheta + - pfcand_phic + - pfcand_phidz + - pfcand_phictgtheta + - pfcand_cdz + - pfcand_cctgtheta + - pfcand_mtof + - pfcand_dndx + - pfcand_charge + - pfcand_isMu + - pfcand_isEl + - pfcand_isChargedHad + - pfcand_isGamma + - pfcand_isNeutralHad + - pfcand_dxy + - pfcand_dz + - pfcand_btagSip2dVal + - pfcand_btagSip2dSig + - pfcand_btagSip3dVal + - pfcand_btagSip3dSig + - pfcand_btagJetDistVal + - pfcand_btagJetDistSig + - pfcand_type + pf_vectors: + length: 75 + pad_mode: wrap + vars: + - - pfcand_e + - null + - - pfcand_p + - null + - - pfcand_e + - null + - - pfcand_e + - null + pf_mask: + length: 75 + pad_mode: constant + vars: + - - pfcand_mask + - null +labels: + type: simple + value: + - recojet_isG + - recojet_isU + - recojet_isS + - recojet_isC + - recojet_isB + - recojet_isTAU + - recojet_isD +observers: [] +monitor_variables: [] +weights: + use_precomputed_weights: false + reweight_method: flat + reweight_vars: + jet_phi: + - -10.0 + - 10.0 + jet_theta: + - -10.0 + - 10.0 + reweight_classes: + - recojet_isG + - recojet_isU + - recojet_isS + - recojet_isC + - recojet_isB + - recojet_isTAU + - recojet_isD + class_weights: + - 1 + - 1 + - 1 + - 1 + - 1 + - 1 + - 1 + reweight_hists: + recojet_isG: + - - 0.803920567035675 + recojet_isU: + - - 0.7541366219520569 + recojet_isS: + - - 0.802278459072113 + recojet_isC: + - - 0.7855589985847473 + recojet_isB: + - - 0.7855589985847473 + recojet_isTAU: + - - 0.8999999761581421 + recojet_isD: + - - 0.7541366219520569 diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.json b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.json new file mode 100644 index 00000000..83b4da0a --- /dev/null +++ b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.json @@ -0,0 +1,396 @@ +{ + "output_names": [ + "recojet_isG", + "recojet_isQ", + "recojet_isS", + "recojet_isC", + "recojet_isB" + ], + "input_names": [ + "pf_points", + "pf_features", + "pf_vectors", + "pf_mask" + ], + "pf_points": { + "var_names": [ + "pfcand_thetarel", + "pfcand_phirel" + ], + "var_infos": { + "pfcand_thetarel": { + "median": 0.3606320321559906, + "norm_factor": 1.6306203960800496, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phirel": { + "median": 0.0002556022081989795, + "norm_factor": 0.46948837373539853, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_features": { + "var_names": [ + "pfcand_erel_log", + "pfcand_thetarel", + "pfcand_phirel", + "pfcand_dptdpt", + "pfcand_detadeta", + "pfcand_dphidphi", + "pfcand_dxydxy", + "pfcand_dzdz", + "pfcand_dxydz", + "pfcand_dphidxy", + "pfcand_dlambdadz", + "pfcand_dxyc", + "pfcand_dxyctgtheta", + "pfcand_phic", + "pfcand_phidz", + "pfcand_phictgtheta", + "pfcand_cdz", + "pfcand_cctgtheta", + "pfcand_mtof", + "pfcand_dndx", + "pfcand_charge", + "pfcand_isMu", + "pfcand_isEl", + "pfcand_isChargedHad", + "pfcand_isGamma", + "pfcand_isNeutralHad", + "pfcand_dxy", + "pfcand_dz", + "pfcand_btagSip2dVal", + "pfcand_btagSip2dSig", + "pfcand_btagSip3dVal", + "pfcand_btagSip3dSig", + "pfcand_btagJetDistVal", + "pfcand_btagJetDistSig" + ], + "var_infos": { + "pfcand_erel_log": { + "median": -1.8151949644088745, + "norm_factor": 1.6230400937923801, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_thetarel": { + "median": 0.3606320321559906, + "norm_factor": 1.6306203960800496, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phirel": { + "median": 0.0002556022081989795, + "norm_factor": 0.46948837373539853, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dptdpt": { + "median": -9.0, + "norm_factor": 0.11111111111110991, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_detadeta": { + "median": -9.0, + "norm_factor": 0.11111111034557275, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dphidphi": { + "median": -9.0, + "norm_factor": 0.11111107770056672, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxydxy": { + "median": -9.0, + "norm_factor": 0.1111054124522878, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dzdz": { + "median": -9.0, + "norm_factor": 0.1111105908437637, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxydz": { + "median": -9.0, + "norm_factor": 0.11111102914593461, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dphidxy": { + "median": -9.0, + "norm_factor": 0.11111119697916892, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dlambdadz": { + "median": -9.0, + "norm_factor": 0.11111111196585105, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxyc": { + "median": -9.0, + "norm_factor": 0.11111111107993017, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxyctgtheta": { + "median": -9.0, + "norm_factor": 0.11111111115695796, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phic": { + "median": -9.0, + "norm_factor": 0.11111111111116635, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phidz": { + "median": -9.0, + "norm_factor": 0.11111111111495718, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_phictgtheta": { + "median": -9.0, + "norm_factor": 0.11111111043836697, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_cdz": { + "median": -9.0, + "norm_factor": 0.11111111111012681, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_cctgtheta": { + "median": -9.0, + "norm_factor": 0.1111111111111261, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_mtof": { + "median": 0.07332593947649002, + "norm_factor": 5.140589694831398, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dndx": { + "median": 0.0, + "norm_factor": 0.72003525410647, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_charge": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isMu": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isEl": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isChargedHad": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isGamma": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_isNeutralHad": { + "median": 0.0, + "norm_factor": 1.0, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dxy": { + "median": -9.0, + "norm_factor": 0.11103669093395387, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_dz": { + "median": -9.0, + "norm_factor": 0.11108313245378311, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip2dVal": { + "median": -9.0, + "norm_factor": 0.11100134875917528, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip2dSig": { + "median": -9.0, + "norm_factor": 0.10324104907259508, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip3dVal": { + "median": -9.0, + "norm_factor": 0.11099737353525808, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagSip3dSig": { + "median": -9.0, + "norm_factor": 0.10317010115501339, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagJetDistVal": { + "median": -9.0, + "norm_factor": 0.11106430427459978, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + }, + "pfcand_btagJetDistSig": { + "median": -9.0, + "norm_factor": 0.10770140600206164, + "replace_inf_value": 0, + "lower_bound": -5, + "upper_bound": 5, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_vectors": { + "var_names": [ + "pfcand_e", + "pfcand_p", + "pfcand_e", + "pfcand_e" + ], + "var_infos": { + "pfcand_e": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + }, + "pfcand_p": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + } + }, + "var_length": 75 + }, + "pf_mask": { + "var_names": [ + "pfcand_mask" + ], + "var_infos": { + "pfcand_mask": { + "median": 0, + "norm_factor": 1, + "replace_inf_value": 0, + "lower_bound": -1e+32, + "upper_bound": 1e+32, + "pad": 0 + } + }, + "var_length": 75 + } +} \ No newline at end of file diff --git a/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.onnx b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.onnx new file mode 100644 index 00000000..cf439a4e Binary files /dev/null and b/examples/flavour-tagging/data/wc_pt_7classes_12_04_2023/fccee_flavtagging_edm4hep_wc_v1.onnx differ diff --git a/examples/flavour-tagging/simple-flavour-tagging.ipynb b/examples/flavour-tagging/simple-flavour-tagging.ipynb new file mode 100644 index 00000000..08b61afd --- /dev/null +++ b/examples/flavour-tagging/simple-flavour-tagging.ipynb @@ -0,0 +1,299 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "21a84728", + "metadata": {}, + "outputs": [], + "source": [ + "using EDM4hep\n", + "using EDM4hep.RootIO\n", + "using StaticArrays\n", + "using LorentzVectorHEP\n", + "using JSON\n", + "using ONNXRunTime\n", + "using StructArrays" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "3b92ca23", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\u001b[32m\u001b[1m No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Project.toml`\n", + "\u001b[32m\u001b[1m No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Manifest.toml`\n", + "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n", + "\u001b[32m\u001b[1m No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Project.toml`\n", + "\u001b[32m\u001b[1m No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Manifest.toml`\n", + "WARNING: Method definition get_erel_log_cluster(Array{JetReconstruction.EEJet, 1}, Array{var\"#s102\", 1} where var\"#s102\"<:(StructArrays.StructArray{EDM4hep.ReconstructedParticle, 1, var\"#s1\", I} where I where var\"#s1\")) in module JetFlavourHelper at /Users/harrywanghc/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl/ext/JetFlavourTagging/JetConstituentUtils.jl:1426 overwritten at /Users/harrywanghc/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl/ext/JetFlavourTagging/JetConstituentUtils.jl:1468.\n", + "ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[32m\u001b[1mStatus\u001b[22m\u001b[39m `~/.julia/environments/v1.11/Project.toml`\n", + " \u001b[90m[c7e460c6] \u001b[39mArgParse v1.2.0\n", + " \u001b[90m[6e4b80f9] \u001b[39mBenchmarkTools v1.6.0\n", + " \u001b[90m[150eb455] \u001b[39mCoordinateTransformations v0.6.4\n", + " \u001b[90m[a93c6f00] \u001b[39mDataFrames v1.7.0\n", + " \u001b[90m[eb32b910] \u001b[39mEDM4hep v0.4.3\n", + "\u001b[32m⌃\u001b[39m \u001b[90m[5c1252a2] \u001b[39mGeometryBasics v0.5.7\n", + "\u001b[32m⌃\u001b[39m \u001b[90m[7073ff75] \u001b[39mIJulia v1.27.0\n", + " \u001b[90m[682c06a0] \u001b[39mJSON v0.21.4\n", + " \u001b[90m[44e8cb2c] \u001b[39mJetReconstruction v0.5.0-dev `~/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl`\n", + " \u001b[90m[98e50ef6] \u001b[39mJuliaFormatter v2.1.2\n", + " \u001b[90m[bdcacae8] \u001b[39mLoopVectorization v0.12.172\n", + "\u001b[32m⌃\u001b[39m \u001b[90m[f612022c] \u001b[39mLorentzVectorHEP v0.1.6\n", + " \u001b[90m[e034b28e] \u001b[39mONNXRunTime v1.3.1\n", + " \u001b[90m[5ad8b20f] \u001b[39mPhysicalConstants v0.2.4\n", + "\u001b[32m⌃\u001b[39m \u001b[90m[91a5bcdd] \u001b[39mPlots v1.40.13\n", + " \u001b[90m[6038ab10] \u001b[39mRotations v1.7.1\n", + " \u001b[90m[90137ffa] \u001b[39mStaticArrays v1.9.13\n", + " \u001b[90m[10745b16] \u001b[39mStatistics v1.11.1\n", + " \u001b[90m[2913bbd2] \u001b[39mStatsBase v0.34.5\n", + "\u001b[33m⌅\u001b[39m \u001b[90m[09ab397b] \u001b[39mStructArrays v0.6.21\n", + "\u001b[32m⌃\u001b[39m \u001b[90m[1e6cf692] \u001b[39mTestEnv v1.102.0\n", + " \u001b[90m[3cd96dde] \u001b[39mUnROOT v0.10.36\n", + " \u001b[90m[37e2e46d] \u001b[39mLinearAlgebra v1.11.0\n", + "\u001b[36m\u001b[1mInfo\u001b[22m\u001b[39m Packages marked with \u001b[32m⌃\u001b[39m and \u001b[33m⌅\u001b[39m have new versions available. Those with \u001b[32m⌃\u001b[39m may be upgradable, but those with \u001b[33m⌅\u001b[39m are restricted by compatibility constraints from upgrading. To see why use `status --outdated`\n" + ] + } + ], + "source": [ + "using Pkg\n", + "Pkg.resolve()\n", + "Pkg.develop(Pkg.PackageSpec(path = \"/Users/harrywanghc/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl\"))\n", + "using JetReconstruction\n", + "Pkg.status()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "379f4304", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The model predicts these flavor classes:\n", + " - recojet_isG\n", + " - recojet_isQ\n", + " - recojet_isS\n", + " - recojet_isC\n", + " - recojet_isB\n" + ] + } + ], + "source": [ + "# Paths to model files\n", + "model_dir = \"data/wc_pt_7classes_12_04_2023\"\n", + "onnx_path = joinpath(model_dir, \"fccee_flavtagging_edm4hep_wc_v1.onnx\")\n", + "json_path = joinpath(model_dir, \"fccee_flavtagging_edm4hep_wc_v1.json\")\n", + "\n", + "# Load the configuration and model\n", + "config = JSON.parsefile(json_path)\n", + "model = ONNXRunTime.load_inference(onnx_path)\n", + "\n", + "# Display the output classes we'll predict\n", + "println(\"The model predicts these flavor classes:\")\n", + "for class_name in config[\"output_names\"]\n", + " println(\" - \", class_name)\n", + "end" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "51c71156", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loaded 100000 events\n", + "Processing event #15\n", + "Loaded 65 reconstructed particles\n", + "Loaded 147 Monte Carlo particles\n", + "Loaded 29 tracks\n" + ] + } + ], + "source": [ + "# Path to ROOT file with EDM4hep data\n", + "edm4hep_path = \"data/events_080263084.root\"\n", + "reader = RootIO.Reader(edm4hep_path)\n", + "\n", + "# Get event information\n", + "events = RootIO.get(reader, \"events\")\n", + "println(\"Loaded $(length(events)) events\")\n", + "\n", + "# Choose a specific event to analyze (event #15)\n", + "event_id = 15\n", + "evt = events[event_id]\n", + "println(\"Processing event #$event_id\")\n", + "\n", + "# Get reconstructed particles and tracks\n", + "recps = RootIO.get(reader, evt, \"ReconstructedParticles\")\n", + "mcps = RootIO.get(reader, evt, \"Particle\")\n", + "MCRecoLinks = RootIO.get(reader, evt, \"MCRecoAssociations\")\n", + "tracks = RootIO.get(reader, evt, \"EFlowTrack_1\")\n", + "\n", + "# Get needed collections for feature extraction\n", + "bz = RootIO.get(reader, evt, \"magFieldBz\", register = false)[1]\n", + "trackdata = RootIO.get(reader, evt, \"EFlowTrack\")\n", + "trackerhits = RootIO.get(reader, evt, \"TrackerHits\")\n", + "gammadata = RootIO.get(reader, evt, \"EFlowPhoton\")\n", + "nhdata = RootIO.get(reader, evt, \"EFlowNeutralHadron\")\n", + "calohits = RootIO.get(reader, evt, \"CalorimeterHits\")\n", + "dNdx = RootIO.get(reader, evt, \"EFlowTrack_2\")\n", + "track_L = RootIO.get(reader, evt, \"EFlowTrack_L\", register = false)\n", + "\n", + "mc_vertices = Vector{LorentzVector{Float32}}(undef, length(recps))\n", + "reco_to_mc = Dict(link.rec_idx.index => link.sim_idx.index for link in MCRecoLinks)\n", + "for (rec_idx, mc_idx) in reco_to_mc\n", + " mc_vertices[rec_idx+1] = LorentzVector(mcps[mc_idx+1].vertex.x, mcps[mc_idx+1].vertex.y, mcps[mc_idx+1].vertex.z, mcps[mc_idx+1].time)\n", + "end\n", + "\n", + "println(\"Loaded $(length(recps)) reconstructed particles\")\n", + "println(\"Loaded $(length(mcps)) Monte Carlo particles\")\n", + "println(\"Loaded $(length(tracks)) tracks\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "f96fe81b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "2-element Vector{StructVector{ReconstructedParticle, @NamedTuple{index::StructVector{ObjectID{ReconstructedParticle}, @NamedTuple{index::Vector{Int64}, collectionID::Vector{UInt32}}, Int64}, type::Vector{Int32}, energy::Vector{Float32}, momentum::StructVector{Vector3f, @NamedTuple{x::Vector{Float32}, y::Vector{Float32}, z::Vector{Float32}}, Int64}, referencePoint::StructVector{Vector3f, @NamedTuple{x::Vector{Float32}, y::Vector{Float32}, z::Vector{Float32}}, Int64}, charge::Vector{Float32}, mass::Vector{Float32}, goodnessOfPID::Vector{Float32}, covMatrix::StructVector{SVector{10, Float32}, NTuple{10, Vector{Float32}}, Int64}, clusters::StructVector{Relation{ReconstructedParticle, Cluster, 1}, @NamedTuple{first::Vector{UInt32}, last::Vector{UInt32}, collid::Vector{UInt32}}, Int64}, tracks::StructVector{Relation{ReconstructedParticle, Track, 2}, @NamedTuple{first::Vector{UInt32}, last::Vector{UInt32}, collid::Vector{UInt32}}, Int64}, particles::StructVector{Relation{ReconstructedParticle, ReconstructedParticle, 3}, @NamedTuple{first::Vector{UInt32}, last::Vector{UInt32}, collid::Vector{UInt32}}, Int64}, particleIDs::StructVector{Relation{ReconstructedParticle, ParticleID, 4}, @NamedTuple{first::Vector{UInt32}, last::Vector{UInt32}, collid::Vector{UInt32}}, Int64}, startVertex_idx::StructVector{ObjectID{Vertex}, @NamedTuple{index::Vector{Int32}, collectionID::Vector{Int32}}, Int64}, particleIDUsed_idx::StructVector{ObjectID{ParticleID}, @NamedTuple{index::Vector{Int32}, collectionID::Vector{Int32}}, Int64}}, Int64}}:\n", + " [ReconstructedParticle(#4, 0, 20.46865f0, (-19.487997, -1.9470071, 5.9491873), (0.0, 0.0, 0.0), -1.0f0, 0.00051099894f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[4], ReconstructedParticle#[], ParticleID#[4], #0, #0), ReconstructedParticle(#1, 0, 41.66906f0, (-38.420082, -8.888379, -13.460811), (0.0, 0.0, 0.0), 1.0f0, 0.105658375f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[1], ReconstructedParticle#[], ParticleID#[1], #0, #0), ReconstructedParticle(#30, 22, 5.97301f0, (-5.5111613, -1.3775525, -1.8456159), (0.0, 0.0, 0.0), 0.0f0, 0.00072206435f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[1], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0)]\n", + " [ReconstructedParticle(#49, 22, 0.038685087f0, (0.0016564546, 0.030539023, 0.023688823), (0.0, 0.0, 0.0), 0.0f0, -7.2657667f-6, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[20], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#50, 22, 0.07068448f0, (0.009334846, 0.04360084, 0.054846358), (0.0, 0.0, 0.0), 0.0f0, -1.211525f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[21], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#20, 0, 2.3770719f0, (0.18115185, 0.92913306, 2.1759794), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[20], ReconstructedParticle#[], ParticleID#[20], #0, #0), ReconstructedParticle(#52, 22, 0.49115527f0, (0.044869967, 0.21402507, 0.43978798), (0.0, 0.0, 0.0), 0.0f0, -6.145881f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[23], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#58, 22, 2.2526124f0, (-0.008166414, 0.70710534, 2.1387374), (0.0, 0.0, 0.0), 0.0f0, -0.00067013147f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[29], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#59, 22, 1.1003908f0, (-0.00017319276, 0.32061982, 1.0526457), (0.0, 0.0, 0.0), 0.0f0, -0.00021608449f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[30], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#54, 22, 1.757357f0, (0.08186126, 0.5793458, 1.6570942), (0.0, 0.0, 0.0), 0.0f0, -0.0007735479f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[25], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#56, 22, 5.843931f0, (0.19479804, 1.8359567, 5.5446234), (0.0, 0.0, 0.0), 0.0f0, -0.0012316033f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[27], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#25, 0, 2.5155842f0, (0.28472182, 0.63002336, 2.414682), (0.0, 0.0, 0.0), 1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[25], ReconstructedParticle#[], ParticleID#[25], #0, #0), ReconstructedParticle(#53, 22, 0.10352973f0, (0.032321043, 0.012602998, 0.09754445), (0.0, 0.0, 0.0), 0.0f0, 2.5103373f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[24], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0) … ReconstructedParticle(#36, 22, 1.2084657f0, (1.1462502, 0.09732889, 0.37017113), (0.0, 0.0, 0.0), 0.0f0, 9.1004724f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[7], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#7, 0, 3.8613276f0, (3.6221685, -0.1164991, 1.3254038), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[7], ReconstructedParticle#[], ParticleID#[7], #0, #0), ReconstructedParticle(#11, 0, 4.4815288f0, (4.0670757, -0.13165577, 1.8724797), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[11], ReconstructedParticle#[], ParticleID#[11], #0, #0), ReconstructedParticle(#28, 0, 0.5223787f0, (0.4566623, 0.0062101614, 0.21170896), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[28], ReconstructedParticle#[], ParticleID#[28], #0, #0), ReconstructedParticle(#41, 22, 0.3654142f0, (0.33327267, 0.007723306, 0.14965703), (0.0, 0.0, 0.0), 0.0f0, 3.5387056f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[12], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#10, 0, 1.0771525f0, (0.83581126, -0.55972606, 0.35903203), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[10], ReconstructedParticle#[], ParticleID#[10], #0, #0), ReconstructedParticle(#37, 22, 0.055122558f0, (0.046858095, -0.02121588, 0.0198167), (0.0, 0.0, 0.0), 0.0f0, 1.1097479f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[8], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#17, 0, 1.157116f0, (0.7993353, -0.37195656, 0.73630774), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[17], ReconstructedParticle#[], ParticleID#[17], #0, #0), ReconstructedParticle(#46, 22, 0.7516057f0, (0.59814185, -0.19748434, 0.41004556), (0.0, 0.0, 0.0), 0.0f0, -6.0317438f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[17], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#48, 22, 0.29473728f0, (0.23144196, -0.054464586, 0.1741789), (0.0, 0.0, 0.0), 0.0f0, 7.500272f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[19], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0)]" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Cluster jets using the EEkt algorithm with R=2.0 and p=1.0\n", + "cs = jet_reconstruct(recps; p = 1.0, R = 2.0, algorithm = JetAlgorithm.EEKt)\n", + "\n", + "# Get 2 exclusive jets\n", + "jets = exclusive_jets(cs; njets=2, T=EEJet)\n", + "\n", + "# For each jet, get its constituent particles\n", + "constituent_indices = [constituent_indexes(jet, cs) for jet in jets]\n", + "jet_constituents = build_constituents_cluster(recps, constituent_indices)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "4f46c6c3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Extracting features for flavor tagging...\n", + "Step 1: Feature extraction completed.\n", + "Step 2: Weaver setup completed.\n", + "Step 3: Input tensor preparation completed.\n", + "Running flavor tagging inference...\n", + "Step 4: Weights retrieval completed.\n", + "Jet scores:\n", + " - recojet_isG: 0.20045522\n", + " - recojet_isB: 0.7949252\n", + " - recojet_isQ: 0.001174435\n", + " - recojet_isC: 0.0027448845\n", + " - recojet_isS: 0.0007002373\n" + ] + } + ], + "source": [ + "println(\"Extracting features for flavor tagging...\")\n", + "feature_data = extract_features(\n", + " jets, \n", + " jet_constituents, \n", + " tracks, \n", + " bz, \n", + " track_L, \n", + " trackdata, \n", + " trackerhits, \n", + " gammadata, \n", + " nhdata, \n", + " calohits, \n", + " dNdx\n", + ")\n", + "println(\"Step 1: Feature extraction completed.\")\n", + "\n", + "model, config = setup_weaver(\n", + " onnx_path,\n", + " json_path\n", + ")\n", + "\n", + "println(\"Step 2: Weaver setup completed.\")\n", + "\n", + "input_tensors = prepare_input_tensor(\n", + " jet_constituents,\n", + " jets,\n", + " config,\n", + " feature_data\n", + ")\n", + "\n", + "println(\"Step 3: Input tensor preparation completed.\")\n", + "\n", + "println(\"Running flavor tagging inference...\")\n", + "weights = get_weights(\n", + " 0, # Thread slot\n", + " feature_data,\n", + " jets,\n", + " jet_constituents,\n", + " config,\n", + " model\n", + ")\n", + "println(\"Step 4: Weights retrieval completed.\")\n", + "\n", + "jet_scores = Dict{String, Vector{Float32}}()\n", + "for (i, score_name) in enumerate(config[\"output_names\"])\n", + " jet_scores[score_name] = get_weight(weights, i-1)\n", + "end\n", + "\n", + "println(\"Jet scores:\")\n", + "for (name, scores) in jet_scores\n", + " println(\" - $name: $(scores[1])\")\n", + "end" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Julia 1.11.5", + "language": "julia", + "name": "julia-1.11" + }, + "language_info": { + "file_extension": ".jl", + "mimetype": "application/julia", + "name": "julia", + "version": "1.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/flavour-tagging/simple-flavour-tagging.jl b/examples/flavour-tagging/simple-flavour-tagging.jl new file mode 100644 index 00000000..9b6ff0b1 --- /dev/null +++ b/examples/flavour-tagging/simple-flavour-tagging.jl @@ -0,0 +1,257 @@ +#!/usr/bin/env julia + +""" +Simple Jet Flavour Tagging Example + +This script demonstrates how to: +1. Load EDM4hep event data +2. Reconstruct jets using JetReconstruction +3. Extract features for flavour tagging +4. Run ONNX neural network inference +5. Get flavour probabilities for each jet + +Run with: julia --project simple-flavour-tagging.jl +""" + +using EDM4hep +using EDM4hep.RootIO +using LorentzVectorHEP +using JSON +using ONNXRunTime +using PhysicalConstants +using StructArrays +using JetReconstruction + +function main() + # Paths to model files + model_dir = "data/wc_pt_7classes_12_04_2023" + onnx_path = joinpath(model_dir, "fccee_flavtagging_edm4hep_wc_v1.onnx") + json_path = joinpath(model_dir, "fccee_flavtagging_edm4hep_wc_v1.json") + + # Check if model files exist + if !isfile(onnx_path) + error("ONNX model not found at: $onnx_path") + end + if !isfile(json_path) + error("JSON config not found at: $json_path") + end + + println("Loading flavour tagging model...") + config = JSON.parsefile(json_path) + model = ONNXRunTime.load_inference(onnx_path) + + println("\nThe model predicts these flavour classes:") + for class_name in config["output_names"] + println(" - $class_name") + end + + # Path to ROOT file with EDM4hep data + edm4hep_path = "data/events_080263084.root" + # edm4hep_path = "/eos/experiment/fcc/ee/generation/DelphesEvents/winter2023/IDEA/wzp6_ee_nunuH_ecm240/events_080263084.root" + if !isfile(edm4hep_path) + error("EDM4hep data file not found at: $edm4hep_path") + end + + println("\nLoading EDM4hep data...") + reader = RootIO.Reader(edm4hep_path) + events = RootIO.get(reader, "events") + println("Loaded $(length(events)) events") + + # Process a specific event (event #12 as in the notebook) + event_id = 12 + println("\nProcessing event #$event_id") + evt = events[event_id] + + # Get reconstructed particles and tracks + recps = RootIO.get(reader, evt, "ReconstructedParticles") + tracks = RootIO.get(reader, evt, "EFlowTrack_1") + + # Get MC particles and links for vertex information + mcps = RootIO.get(reader, evt, "Particle") + MCRecoLinks = RootIO.get(reader, evt, "MCRecoAssociations") + + # Extract MC vertices for each reconstructed particle + mc_vertices = Vector{LorentzVector{Float32}}(undef, length(recps)) + reco_to_mc = Dict(link.rec_idx.index => link.sim_idx.index for link in MCRecoLinks) + for (rec_idx, mc_idx) in reco_to_mc + if rec_idx < length(recps) && mc_idx < length(mcps) + mc_vertices[rec_idx + 1] = LorentzVector(Float32(mcps[mc_idx + 1].vertex.x), + Float32(mcps[mc_idx + 1].vertex.y), + Float32(mcps[mc_idx + 1].vertex.z), + Float32(mcps[mc_idx + 1].time)) + end + end + + # Fill any missing vertices with (0,0,0,0) + for i in 1:length(recps) + if !isassigned(mc_vertices, i) + mc_vertices[i] = LorentzVector(0.0f0, 0.0f0, 0.0f0, 0.0f0) + end + end + + # Get needed collections for feature extraction + bz = RootIO.get(reader, evt, "magFieldBz", register = false)[1] + trackdata = RootIO.get(reader, evt, "EFlowTrack") + trackerhits = RootIO.get(reader, evt, "TrackerHits") + gammadata = RootIO.get(reader, evt, "EFlowPhoton") + nhdata = RootIO.get(reader, evt, "EFlowNeutralHadron") + calohits = RootIO.get(reader, evt, "CalorimeterHits") + dNdx = RootIO.get(reader, evt, "EFlowTrack_2") + track_L = RootIO.get(reader, evt, "EFlowTrack_L", register = false) + + println(" - $(length(recps)) reconstructed particles") + println(" - $(length(tracks)) tracks") + println(" - Magnetic field Bz = $bz T") + + # Print the primary vertex that will be used + primary_vertex = LorentzVector(0.0f0, 0.0f0, 0.0f0, 0.0f0) + for vertex in mc_vertices + if vertex.x != 0.0 || vertex.y != 0.0 || vertex.z != 0.0 + primary_vertex = vertex + break + end + end + println(" - Primary vertex: ($(round(primary_vertex.x, digits=3)), $(round(primary_vertex.y, digits=3)), $(round(primary_vertex.z, digits=3))) mm") + + # Reconstruct jets + println("\nReconstructing jets...") + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, algorithm = JetAlgorithm.EEKt) + + # Get 2 exclusive jets + jets = exclusive_jets(cs, EEJet; njets = 2) + println("Found $(length(jets)) jets") + + # Print jet properties + for (i, jet) in enumerate(jets) + println("\nJet $i:") + println(" - Energy: $(round(jet.E, digits=2)) GeV") + println(" - Pt: $(round(JetReconstruction.pt(jet), digits=2)) GeV") + println(" - Eta: $(round(JetReconstruction.eta(jet), digits=3))") + println(" - Phi: $(round(JetReconstruction.phi(jet), digits=3))") + println(" - Mass: $(round(JetReconstruction.mass(jet), digits=2)) GeV") + end + + # Get jet constituents + println("\nExtracting jet constituents...") + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + + # Access the extension module + ext_mod = Base.get_extension(JetReconstruction, :JetFlavourTagging) + if isnothing(ext_mod) + error("JetFlavourTagging extension not loaded") + end + + jet_constituents = ext_mod.build_constituents_cluster(recps, constituent_indices) + + for (i, constituents) in enumerate(jet_constituents) + println(" - Jet $i has $(length(constituents)) constituents") + end + + # Extract features for flavour tagging + println("\nExtracting features for flavour tagging...") + feature_data = JetReconstruction.extract_features(jets, + jet_constituents, + tracks, + bz, + track_L, + config, + trackdata, + trackerhits, + gammadata, + nhdata, + calohits, + dNdx, + mc_vertices) + + # Prepare input tensors + println("Preparing input tensors...") + input_tensors = JetReconstruction.prepare_input_tensor(jet_constituents, + jets, + config, + feature_data) + + # Run inference + println("Running neural network inference...") + weights = JetReconstruction.get_weights(0, # Thread slot + feature_data, + jets, + jet_constituents, + config, + model) + + # Extract and display results + println("\n" * "="^60) + println("FLAVOUR TAGGING RESULTS") + println("="^60) + + for (jet_idx, jet) in enumerate(jets) + println("\nJet $jet_idx (E=$(round(jet.E, digits=1)) GeV, pT=$(round(JetReconstruction.pt(jet), digits=1)) GeV):") + println("-"^40) + + # Collect scores for this jet + scores = Float32[] + labels = String[] + + for (i, score_name) in enumerate(config["output_names"]) + score = JetReconstruction.get_weight(weights, i - 1)[jet_idx] + push!(scores, score) + push!(labels, score_name) + end + + # Sort by probability (descending) + sorted_indices = sortperm(scores, rev = true) + + # Display scores + for idx in sorted_indices + label = labels[idx] + score = scores[idx] + + # Handle NaN or invalid scores + if isnan(score) || isinf(score) + flavor_map = Dict("recojet_isG" => "Gluon ", + "recojet_isQ" => "Light q ", + "recojet_isS" => "Strange ", + "recojet_isC" => "Charm ", + "recojet_isB" => "Bottom ") + formatted_label = get(flavor_map, label, label) + println(" $formatted_label: [Invalid score]") + continue + end + + bar_length = Int(round(score * 30)) + bar = "█"^bar_length + percentage = round(score * 100, digits = 1) + + # Format label + flavor_map = Dict("recojet_isG" => "Gluon ", + "recojet_isQ" => "Light q ", + "recojet_isS" => "Strange ", + "recojet_isC" => "Charm ", + "recojet_isB" => "Bottom ") + + formatted_label = get(flavor_map, label, label) + println(" $formatted_label: $bar $(percentage)%") + end + + # Identify most likely flavour + max_idx = argmax(scores) + max_label = labels[max_idx] + max_score = scores[max_idx] + + flavour_name = Dict("recojet_isG" => "gluon", + "recojet_isQ" => "light quark", + "recojet_isS" => "strange", + "recojet_isC" => "charm", + "recojet_isB" => "bottom")[max_label] + + println("\n → Most likely: $(flavour_name) ($(round(max_score * 100, digits=1))% confidence)") + end + + println("\n" * "="^60) + println("Processing complete!") +end + +# Run the main function +if abspath(PROGRAM_FILE) == @__FILE__ + main() +end diff --git a/ext/JetFlavourTagging/JetConstituentBuilder.jl b/ext/JetFlavourTagging/JetConstituentBuilder.jl new file mode 100644 index 00000000..60a16ae5 --- /dev/null +++ b/ext/JetFlavourTagging/JetConstituentBuilder.jl @@ -0,0 +1,74 @@ +module JetConstituentBuilder + +using EDM4hep +using StructArrays: StructVector + +# Define type aliases for clarity +const JetConstituents = StructVector{ReconstructedParticle, <:Any} +const JetConstituentsData = Vector{Float32} + +""" + build_constituents(jets::JetConstituents, + reco_particles::JetConstituents) -> Vector{JetConstituents} + +Build the collection of constituents (mapping jet -> reconstructed particles) for all jets in event. + +# Returns +A vector of JetConstituents, each containing the constituents for a specific jet. +""" +# TODO: Fix this function to be interpolate with Julia pipeline. Specificly, what would be the input jets? + +""" + build_constituents_cluster(reco_particles::JetConstituents, + indices::Vector{Vector{Int}}) -> Vector{JetConstituents} + +Build the collection of constituents using cluster indices. + +# Arguments +- reco_particles: a vector of `JetConstituents` representing reconstructed particles. +- indices: a vector of vectors, where each inner vector contains indices of particles for a specific cluster. + +# Returns +A vector of JetConstituents, each containing the constituents for a specific cluster. +""" +function build_constituents_cluster(reco_particles::JetConstituents, + indices::Vector{Vector{Int64}}) + return map(jet_indices -> reco_particles[jet_indices], indices) +end + +""" + get_jet_constituents(constituents_collection::Vector{JetConstituents}, jet_index::Int) -> JetConstituents + +Retrieve the constituents of an indexed jet in the event. +# Arguments +- constituents_collection: constituents collection, a vector of `JetConstituents`. +- jet_index: the index of the jet for which to retrieve constituents (1-based index) + +# Returns +The constituents of the specified jet, or an empty collection if the jet index is invalid. +""" +function get_jet_constituents(constituents_collection::Vector{JetConstituents}, + jet_index::Int) + return constituents_collection[jet_index] +end + +""" + get_constituents(constituents_collection::Vector{JetConstituents}, jet_indices::Vector{Int}) -> Vector{JetConstituents} + +Retrieve the constituents of a collection of indexed jets in the event. + +# Arguments +- constituents_collection: constituents collection, a vector of `JetConstituents`. +- jet_indices: a vector of jet indices (1-based index) for which to retrieve constituents. + +# Returns +A vector of `JetConstituents`, each containing the constituents for the specified jets. +""" +function get_constituents(constituents_collection::Vector{JetConstituents}, + jet_indices::Vector{Int}) + # Filter valid indices and map to corresponding constituents + valid_indices = filter(idx -> 1 <= idx <= length(constituents_collection), jet_indices) + return map(idx -> constituents_collection[idx], valid_indices) +end + +end # module JetConstituentBuilder diff --git a/ext/JetFlavourTagging/JetConstituentUtils.jl b/ext/JetFlavourTagging/JetConstituentUtils.jl new file mode 100644 index 00000000..cf3819bb --- /dev/null +++ b/ext/JetFlavourTagging/JetConstituentUtils.jl @@ -0,0 +1,2191 @@ +using EDM4hep +using JetReconstruction +using StructArrays: StructVector + +# Import physical constants +include("JetPhysicalConstants.jl") +using .JetPhysicalConstants + +const JetConstituents = StructVector{ReconstructedParticle, <:Any} +const JetConstituentsData = Vector{Float32} + +### Basic Kinematic (11) + +# get_pt - Transverse momentum +# get_p - Total momentum +# get_e - Energy +# get_mass - Mass +# get_type - Particle type +# get_charge - Electric charge +# get_theta - Polar angle +# get_phi - Azimuthal angle +# get_y - Rapidity +# get_eta - Pseudorapidity +# get_Bz - Magnetic field component + +""" + get_pt(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the transverse momentum of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents (each element contains particles for one jet) + +# Returns +A vector of vectors of transverse momentum values (sqrt(px^2 + py^2)) +""" +function get_pt(jets_constituents::Vector{<:JetConstituents}) + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + Float32[@inbounds sqrt(mom_x[i]^2 + mom_y[i]^2) + for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_p(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the momentum magnitude of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of momentum magnitudes (sqrt(px^2 + py^2 + pz^2)) +""" +function get_p(jets_constituents::Vector{<:JetConstituents}) + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + mom_z = jet_constituents.momentum.z + Float32[@inbounds sqrt(mom_x[i]^2 + mom_y[i]^2 + mom_z[i]^2) + for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_e(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the energy of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of energy values +""" +function get_e(jets_constituents::Vector{<:JetConstituents}) + return [jet_constituents.energy for jet_constituents in jets_constituents] +end + +""" + get_type(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the PDG type of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of particle types (PDG codes/Particle IDs) +""" +function get_type(jets_constituents::Vector{<:JetConstituents}) + return [jet_constituents.type for jet_constituents in jets_constituents] +end + +""" + get_mass(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the mass of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of mass values +""" +function get_mass(jets_constituents::Vector{<:JetConstituents}) + return [jet_constituents.mass for jet_constituents in jets_constituents] +end + +""" + get_charge(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the charge of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of charge values +""" +function get_charge(jets_constituents::Vector{<:JetConstituents}) + return [jet_constituents.charge for jet_constituents in jets_constituents] +end + +""" + get_theta(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the polar angle of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of polar angle values +""" +function get_theta(jets_constituents::Vector{<:JetConstituents}) + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + mom_z = jet_constituents.momentum.z + Float32[@inbounds(let x = mom_x[i], y = mom_y[i], z = mom_z[i] + (x == 0.0f0 && y == 0.0f0 && z == 0.0f0) ? 0.0f0 : + atan(sqrt(x^2 + y^2), z) + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_phi(jets_constituents::Vector{JetConstituents}) -> Vector{JetConstituentsData} + +Get the azimuthal angle of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of azimuthal angle values +""" +function get_phi(jets_constituents::Vector{<:JetConstituents}) + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + Float32[@inbounds(let x = mom_x[i], y = mom_y[i] + (x == 0.0f0 && y == 0.0f0) ? 0.0f0 : atan(y, x) + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_y(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the rapidity of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of rapidity values +""" +function get_y(jets_constituents::Vector{<:JetConstituents}) + return [begin + energies = jet_constituents.energy + mom_z = jet_constituents.momentum.z + Float32[@inbounds(let e = energies[i], pz = mom_z[i] + 0.5f0 * log((e + pz) / (e - pz)) + end) for i in eachindex(energies)] + end + for jet_constituents in jets_constituents] +end + +""" + get_eta(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Get the pseudorapidity of each particle in each jet. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +A vector of vectors of pseudorapidity values (eta = -ln(tan(theta/2))) +""" +function get_eta(jets_constituents::Vector{<:JetConstituents}) + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + mom_z = jet_constituents.momentum.z + Float32[@inbounds(let x = mom_x[i], y = mom_y[i], z = mom_z[i] + p = sqrt(x^2 + y^2 + z^2) + if p == 0.0f0 + 0.0f0 + elseif p == abs(z) # particle along beam axis + sign(z) * Inf32 + else + 0.5f0 * log((p + z) / (p - z)) + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_Bz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Calculate the magnetic field Bz for each particle based on track curvature and momentum. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: Vector of track states (used to get the omega value) + +# Returns +A vector of vectors of Bz values. +""" +function get_Bz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + a = C_LIGHT * MM_TO_M / FS_TO_S + n_tracks = length(tracks) + + # If tracks is a StructVector, we can access omega column directly + omega_values = tracks.omega + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + charges = jet_constituents.charge + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + pt = sqrt(mom_x[i]^2 + mom_y[i]^2) + omega_values[track_idx + 1] / a * pt * + copysign(1.0f0, charges[i]) + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +### Track Related Functions (5) + +## Track Parameter Transformations (XPtoPar) +# XPtoPar_dxy - Transformed transverse impact parameter +# XPtoPar_dz - Transformed longitudinal impact parameter +# XPtoPar_phi - Transformed azimuthal angle +# XPtoPar_C - Track curvature parameter +# XPtoPar_ct - c×tau parameter + +""" + get_dxy(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the transverse impact parameter dxy for each particle in each jet relative to vertex V. +Reference: FCCAnalyses c++ function XPtoPar_dxy, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects +- V: LorentzVector representing the primary vertex +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of dxy values (one vector per jet) +""" +function get_dxy(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) + cSpeed_Bz = C_LIGHT * NS_TO_S * Bz + n_tracks = length(tracks) + + Vx, Vy = Float32(V.x), Float32(V.y) + + D0_values = tracks.D0 + phi_values = tracks.phi + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + charges = jet_constituents.charge + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + idx = track_idx + 1 + D0 = D0_values[idx] + phi0 = phi_values[idx] + + sin_phi, cos_phi = sincos(phi0) + x1 = -D0 * sin_phi - Vx + x2 = D0 * cos_phi - Vy + + px = mom_x[i] + py = mom_y[i] + + a = -charges[i] * cSpeed_Bz + pt = sqrt(px^2 + py^2) + r2 = x1^2 + x2^2 + cross = x1 * py - x2 * px + + # Compute impact parameter + discriminant = pt^2 - 2 * a * cross + a^2 * r2 + if discriminant > 0 + t = sqrt(discriminant) + if pt < 10.0f0 + (t - pt) / a + else + (-2 * cross + a * r2) / (t + pt) + end + else + UNDEF_VAL + end + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_dz(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the longitudinal impact parameter dz for each particle in each jet relative to vertex V. +Reference: FCCAnalyses c++ function XPtoPar_dz, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects +- V: LorentzVector representing the primary vertex +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of dz values (one vector per jet) +""" +function get_dz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) + cSpeed_Bz = C_LIGHT * NS_TO_S * Bz + n_tracks = length(tracks) + + Vx, Vy, Vz = Float32(V.x), Float32(V.y), Float32(V.z) + + D0_values = tracks.D0 + Z0_values = tracks.Z0 + phi_values = tracks.phi + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + mom_z = jet_constituents.momentum.z + charges = jet_constituents.charge + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + idx = track_idx + 1 + D0 = D0_values[idx] + Z0 = Z0_values[idx] + phi0 = phi_values[idx] + + sin_phi, cos_phi = sincos(phi0) + + x1 = -D0 * sin_phi - Vx + x2 = D0 * cos_phi - Vy + x3 = Z0 - Vz + + px = mom_x[i] + py = mom_y[i] + pz = mom_z[i] + + # Compute intermediate values + a = -charges[i] * cSpeed_Bz + pt = sqrt(px^2 + py^2) + c = a / (2 * pt) + r2 = x1^2 + x2^2 + cross = x1 * py - x2 * px + t = sqrt(pt^2 - 2 * a * cross + a^2 * r2) + + d = if pt < 10.0f0 + (t - pt) / a + else + (-2 * cross + a * r2) / (t + pt) + end + + b_arg = max(r2 - d^2, 0.0f0) / (1 + 2 * c * d) + b = c * sqrt(b_arg) + if abs(b) > 1.0f0 + b = sign(b) + end + + # Calculate st and ct + st = asin(b) / c + ct = pz / pt + + # Calculate z0 + dot = x1 * px + x2 * py + if dot > 0.0f0 + x3 - ct * st + else + x3 + ct * st + end + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_phi0(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the phi angle at the point of closest approach for each particle relative to vertex V. +This is a Julia implementation of the C++ function XPtoPar_phi. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects +- V: LorentzVector representing the primary vertex +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of phi values (one vector per jet) +""" +function get_phi0(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + V::LorentzVector, Bz::Float32) + cSpeed_Bz = C_LIGHT * NS_TO_S * Bz + n_tracks = length(tracks) + + Vx, Vy = Float32(V.x), Float32(V.y) + + D0_values = tracks.D0 + phi_values = tracks.phi + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + charges = jet_constituents.charge + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + idx = track_idx + 1 + D0 = D0_values[idx] + phi0_track = phi_values[idx] + + sin_phi, cos_phi = sincos(phi0_track) + + x1 = -D0 * sin_phi - Vx + x2 = D0 * cos_phi - Vy + + px = mom_x[i] + py = mom_y[i] + + a = -charges[i] * cSpeed_Bz + + # Minimize redundant calculations + pt2 = px^2 + py^2 + r2 = x1^2 + x2^2 + cross = x1 * py - x2 * px + two_a_cross = 2 * a * cross + a2_r2 = a^2 * r2 + + t = sqrt(pt2 - two_a_cross + a2_r2) + inv_t = 1.0f0 / t + + a_x1 = a * x1 + a_x2 = a * x2 + atan((py - a_x1) * inv_t, (px + a_x2) * inv_t) + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_c(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the track curvature for each particle in each jet. +Reference: FCCAnalyses c++ function XPtoPar_C, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of C values (one vector per jet) +""" +function get_c(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + Bz::Float32) + cSpeed_Bz_half = C_LIGHT * MM_TO_M / FS_TO_S * Bz * 0.5f0 + n_tracks = length(tracks) + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + charges = jet_constituents.charge + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + px = mom_x[i] + py = mom_y[i] + inv_pt = 1.0f0 / sqrt(px^2 + py^2) + copysign(cSpeed_Bz_half * inv_pt, charges[i]) + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +""" + get_ct(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the c*tau for each particle in each jet. +Reference: FCCAnalyses c++ function XPtoPar_ct, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of ct values (one vector per jet) +""" +function get_ct(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}, + Bz::Float32) + n_tracks = length(tracks) + + return [begin + mom_x = jet_constituents.momentum.x + mom_y = jet_constituents.momentum.y + mom_z = jet_constituents.momentum.z + track_indices = jet_constituents.tracks + + Float32[@inbounds(let track_idx = track_indices[i].first + if track_idx < n_tracks + px = mom_x[i] + py = mom_y[i] + pz = mom_z[i] + pt = sqrt(px^2 + py^2) + pz / pt + else + UNDEF_VAL + end + end) for i in eachindex(mom_x)] + end + for jet_constituents in jets_constituents] +end + +### Covariance Matrix Elements (15) + +## Diagonal Elements +# get_omega_cov - Omega variance +# get_d0_cov - d0 variance +# get_z0_cov - z0 variance +# get_phi0_cov - phi0 variance +# get_tanlambda_cov - tanLambda variance + +## Off-diagonal Elements +# get_d0_z0_cov +# get_phi0_d0_cov +# get_phi0_z0_cov +# get_tanlambda_phi0_cov +# get_tanlambda_d0_cov +# get_tanlambda_z0_cov +# get_omega_tanlambda_cov +# get_omega_phi0_cov +# get_omega_d0_cov +# get_omega_z0_cov + +""" + get_dxydxy(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the d0 covariance (dxy/dxy) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dxydxy values +""" +function get_dxydxy(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[1], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dphidxy(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the phi0-d0 covariance (dphi/dxy) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dphidxy values +""" +function get_dphidxy(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[2], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dphidphi(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the phi covariance (dphi/dphi) for each particle in each jet from its associated track. +Reference: FCCAnalyses c++ function get_phi0_cov, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dphidphi values +""" +function get_dphidphi(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[3], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dxyc(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the d0-omega covariance (dxy/c) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dxyc values +""" +function get_dxyc(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[4], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_phic(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the phi0-omega covariance (phi/c) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of phiomega values +""" +function get_phic(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[5], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dptdpt(jets_constituents::Vector{JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the omega covariance (dpt/dpt) for each particle in each jet from its associated track. +Reference: FCCAnalyses c++ function get_omega_cov, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents (each element contains particles for one jet) +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dptdpt values +""" +function get_dptdpt(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[6], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dxydz(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the d0-z0 covariance (dxy/dz) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dxy/dz values +""" +function get_dxydz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[7], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_phidz(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the phi0-z0 covariance (dphi/dz) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of phidz values +""" +function get_phidz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[8], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_cdz(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the omega-z0 covariance (c/dz) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dxdz values +""" +function get_cdz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[9], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dzdz(jets_constituents::Vector{JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the z0 covariance (dz/dz) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of z0 covariance values +""" +function get_dzdz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[10], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dxyctgtheta(jets_constituents::Vector{<:JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the d0-tanLambda covariance (dxy/ctgtheta) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dxyctgtheta values +""" +function get_dxyctgtheta(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[11], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_phictgtheta(jets_constituents::Vector{<:JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the phi0-tanLambda covariance (phi/ctgtheta) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of phictgtheta values +""" +function get_phictgtheta(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[12], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_cctgtheta(jets_constituents::Vector{<:JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the omega-tanLambda covariance (c/ctgtheta) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of cctgtheta values +""" +function get_cctgtheta(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[13], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_dlambdadz(jets_constituents::Vector{<:JetConstituents}, tracks::Vector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the tanLambda-z0 covariance (dlambda/dz) for each particle. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of dlambdadz values +""" +function get_dlambdadz(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[14], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +""" + get_detadeta(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) -> Vector{JetConstituentsData} + +Get the tanLambda covariance (deta/deta) for each particle in each jet from its associated track. +Reference: FCCAnalyses c++ function get_tanlambda_cov, adapted for jet constituents. + +# Arguments +- jets_constituents: Vector of jet constituents +- tracks: StructVector of TrackState objects + +# Returns +Vector of vectors of detadeta values +""" +function get_detadeta(jets_constituents::Vector{<:JetConstituents}, + tracks::StructVector{EDM4hep.TrackState}) + n_jets = length(jets_constituents) + n_tracks = length(tracks) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + track_indices = jet_constituents.tracks + n_particles = length(track_indices) + + jet_result = Vector{Float32}(undef, n_particles) + + @simd ivdep for i in 1:n_particles + track_idx = track_indices[i].first + jet_result[i] = ifelse(track_idx < n_tracks, + tracks[track_idx + 1].covMatrix[15], + UNDEF_VAL) + end + + result[j] = jet_result + end + + return result +end + +### Particle Identification (5) + +# get_is_mu - Check if constituent is muon +# get_is_el - Check if constituent is electron +# get_is_charged_had - Check if constituent is charged hadron +# get_is_gamma - Check if constituent is photon +# get_is_neutral_had - Check if constituent is neutral hadron + +""" + get_is_mu(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Check if each constituent particle is a muon. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +Vector of vectors of is muon boolean values as Float32. +""" +function get_is_mu(jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + charges = jet_constituents.charge + masses = jet_constituents.mass + n_particles = length(charges) + + is_mu = Vector{Float32}(undef, n_particles) + + @simd for i in 1:n_particles + charge_check = abs(charges[i]) > 0 + mass_check = abs(masses[i] - MUON_MASS) < MUON_TOLERANCE + is_mu[i] = (charge_check & mass_check) ? 1.0f0 : 0.0f0 + end + + result[j] = is_mu + end + + return result +end + +""" + get_is_el(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Check if each constituent particle is an electron. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +Vector of vectors of is electron boolean values as Float32. +""" +function get_is_el(jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + charges = jet_constituents.charge + masses = jet_constituents.mass + n_particles = length(charges) + + is_mu = Vector{Float32}(undef, n_particles) + + @simd for i in 1:n_particles + charge_check = abs(charges[i]) > 0 + mass_check = abs(masses[i] - ELECTRON_MASS) < ELECTRON_TOLERANCE + is_mu[i] = (charge_check & mass_check) ? 1.0f0 : 0.0f0 + end + + result[j] = is_mu + end + + return result +end + +""" + get_is_charged_had(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Check if each constituent particle is a charged hadron. + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +Vector of vectors of is charged hadron boolean values as Float32. +""" +function get_is_charged_had(jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + charges = jet_constituents.charge + masses = jet_constituents.mass + n_particles = length(charges) + + is_mu = Vector{Float32}(undef, n_particles) + + @simd for i in 1:n_particles + charge_check = abs(charges[i]) > 0 + mass_check = abs(masses[i] - PION_MASS) < PION_TOLERANCE + is_mu[i] = (charge_check & mass_check) ? 1.0f0 : 0.0f0 + end + + result[j] = is_mu + end + + return result +end + +""" + get_is_gamma(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Check if each constituent particle is a photon (gamma) (PDG 22). + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +Vector of vectors of is photon boolean values as Float32. +""" +function get_is_gamma(jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + types = jet_constituents.type + n_particles = length(types) + + is_gamma = Vector{Float32}(undef, n_particles) + + @simd for i in 1:n_particles + is_gamma[i] = types[i] == PDG_PHOTON ? 1.0f0 : 0.0f0 + end + + result[j] = is_gamma + end + + return result +end + +""" + get_is_neutral_had(jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Check if each constituent particle is a neutral hadron (PDG 130). + +# Arguments +- jets_constituents: Vector of jet constituents + +# Returns +Vector of vectors of is neutral hadron boolean values as Float32. +""" +function get_is_neutral_had(jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for j in 1:n_jets + jet_constituents = jets_constituents[j] + types = jet_constituents.type + n_particles = length(types) + + is_gamma = Vector{Float32}(undef, n_particles) + + @simd for i in 1:n_particles + is_gamma[i] = types[i] == PDG_K_LONG ? 1.0f0 : 0.0f0 + end + + result[j] = is_gamma + end + + return result +end + +### Relative Kinematics (5) + +# get_erel_cluster - Relative energy for clustered jets +# get_erel_log_cluster - Log of relative energy for clustered jets +# get_thetarel_cluster - Relative polar angle for clustered jets +# get_phirel_cluster - Relative azimuthal angle for clustered jets +# get_theta_phi_rel_cluster - Combined relative angles for clustered jets as they use the same logic + +""" + get_erel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) -> Vector{JetConstituentsData} + +Calculate relative energy (E_const/E_jet) for each constituent particle in clustered jets. + +# Arguments +- `jets::Vector{EEJet}`: Vector of clustered jets. +- `jets_constituents::Vector{<:JetConstituents}`: Vector of jet constituents corresponding to the jets. + +# Returns +Vector containing relative energy values for each constituent in the jets. +""" +function get_erel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets) + res = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for i in 1:n_jets + e_jet = jets[i].E + constituents_collection = jets_constituents[i] + energies = constituents_collection.energy + n_constituents = length(energies) + jet_constituents_collection = Vector{Float32}(undef, n_constituents) + + if e_jet > 0.0f0 + inv_e_jet = 1.0f0 / e_jet + @inbounds @simd for j in 1:n_constituents + jet_constituents_collection[j] = energies[j] * inv_e_jet + end + else + @inbounds @simd for j in 1:n_constituents + jet_constituents_collection[j] = 0.0f0 + end + end + + res[i] = jet_constituents_collection + end + + return res +end + +""" + get_erel_log_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{JetConstituents}) -> Vector{JetConstituentsData} + +Calculate log of relative energy (log(E_const/E_jet)) for each constituent particle in clustered jets. + +# Arguments +- `jets::Vector{EEJet}`: Vector of clustered jets +- `jets_constituents::Vector{JetConstituents}`: Vector of jet constituents corresponding to the jets + +# Returns +Vector containing log of relative energy values for each constituent in the jets +""" +function get_erel_log_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets) + res = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for i in 1:n_jets + e_jet = jets[i].E + constituents_collection = jets_constituents[i] + energies = constituents_collection.energy + n_constituents = length(energies) + jet_constituents_collection = Vector{Float32}(undef, n_constituents) + + if e_jet > 0.0f0 + inv_e_jet = 1.0f0 / e_jet + @inbounds @simd for j in 1:n_constituents + jet_constituents_collection[j] = log10(energies[j] * inv_e_jet) + end + else + @inbounds @simd for j in 1:n_constituents + jet_constituents_collection[j] = 0.0f0 + end + end + + res[i] = jet_constituents_collection + end + + return res +end + +""" + get_thetarel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{JetConstituents}) -> Vector{JetConstituentsData} + +Calculate relative theta angle between constituent particle and clustered jet axis. + +# Arguments +- `jets::Vector{EEJet}`: Vector of clustered jets +- `jets_constituents::Vector{JetConstituents}`: Vector of jet constituents corresponding to the jets + +# Returns +Vector containing relative theta angle values for each constituent in the jets +""" +function get_thetarel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for i in 1:n_jets + jet = jets[i] + px, py, pz = jet.px, jet.py, jet.pz + + # Pre-compute jet angles + pt_jet = sqrt(px^2 + py^2) + theta_jet = atan(pt_jet, pz) + phi_jet = atan(py, px) + + # Pre-compute trig values using sincos + sin_phi, cos_phi = sincos(-phi_jet) + sin_theta, cos_theta = sincos(-theta_jet) + + constituents_collection = jets_constituents[i] + mom_x = constituents_collection.momentum.x + mom_y = constituents_collection.momentum.y + mom_z = constituents_collection.momentum.z + n_constituents = length(mom_x) + jet_constituents_collection = Vector{Float32}(undef, n_constituents) + + @inbounds for j in 1:n_constituents + # First rotation + p_rot_x = mom_x[j] * cos_phi - mom_y[j] * sin_phi + p_rot_y = mom_x[j] * sin_phi + mom_y[j] * cos_phi + + # Second rotation + p_rot2_x = p_rot_x * cos_theta - mom_z[j] * sin_theta + p_rot2_z = p_rot_x * sin_theta + mom_z[j] * cos_theta + + pt_rot_sq = p_rot2_x^2 + p_rot_y^2 + pt_rot = sqrt(pt_rot_sq) + jet_constituents_collection[j] = atan(pt_rot, p_rot2_z) + end + + result[i] = jet_constituents_collection + end + + return result +end + +""" + get_phirel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{JetConstituents}) -> Vector{JetConstituentsData} + +Calculate relative phi angle between constituent particle and clustered jet axis. + +# Arguments +- `jets::Vector{EEJet}`: Vector of clustered jets +- `jets_constituents::Vector{JetConstituents}`: Vector of jet constituents corresponding to the jets + +# Returns +Vector containing relative phi angle values for each constituent in the jets +""" +function get_phirel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets) + result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for i in 1:n_jets + jet = jets[i] + px, py, pz = jet.px, jet.py, jet.pz + + # Pre-compute jet angles + pt_jet = sqrt(px^2 + py^2) + theta_jet = atan(pt_jet, pz) + phi_jet = atan(py, px) + + # Pre-compute trig values using sincos + sin_phi, cos_phi = sincos(-phi_jet) + sin_theta, cos_theta = sincos(-theta_jet) + + constituents_collection = jets_constituents[i] + mom_x = constituents_collection.momentum.x + mom_y = constituents_collection.momentum.y + mom_z = constituents_collection.momentum.z + n_constituents = length(mom_x) + jet_constituents_collection = Vector{Float32}(undef, n_constituents) + + @inbounds for j in 1:n_constituents + # First rotation around z-axis by -phi_jet + p_rot_x = mom_x[j] * cos_phi - mom_y[j] * sin_phi + p_rot_y = mom_x[j] * sin_phi + mom_y[j] * cos_phi + + # Second rotation around y-axis by -theta_jet + p_rot2_x = p_rot_x * cos_theta - mom_z[j] * sin_theta + + # Calculate phi in rotated frame + jet_constituents_collection[j] = atan(p_rot_y, p_rot2_x) + end + + result[i] = jet_constituents_collection + end + + return result +end + +""" + get_thetaphirel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{JetConstituents}) -> Vector{JetConstituentsData} + +Calculate relative theta and phi angles between constituent particles and clustered jet axis. + +# Arguments +- `jets::Vector{EEJet}`: Vector of clustered jets +- `jets_constituents::Vector{JetConstituents}`: Vector of jet constituents corresponding to the jets + +# Returns +Tuple of Vectors containing relative theta and phi angle values for each constituent in the jets +""" +function get_thetarel_phirel_cluster(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}) + n_jets = length(jets) + theta_result = Vector{Vector{Float32}}(undef, n_jets) + phi_result = Vector{Vector{Float32}}(undef, n_jets) + + @inbounds for i in 1:n_jets + jet = jets[i] + px, py, pz = jet.px, jet.py, jet.pz + + # Pre-compute jet angles + pt_jet = sqrt(px^2 + py^2) + theta_jet = atan(pt_jet, pz) + phi_jet = atan(py, px) + + # Pre-compute trig values using sincos + sin_phi, cos_phi = sincos(-phi_jet) + sin_theta, cos_theta = sincos(-theta_jet) + + constituents_collection = jets_constituents[i] + mom_x = constituents_collection.momentum.x + mom_y = constituents_collection.momentum.y + mom_z = constituents_collection.momentum.z + n_constituents = length(mom_x) + + jet_theta = Vector{Float32}(undef, n_constituents) + jet_phi = Vector{Float32}(undef, n_constituents) + + @inbounds for j in 1:n_constituents + # First rotation around z-axis by -phi_jet + p_rot_x = mom_x[j] * cos_phi - mom_y[j] * sin_phi + p_rot_y = mom_x[j] * sin_phi + mom_y[j] * cos_phi + + # Second rotation around y-axis by -theta_jet + p_rot2_x = p_rot_x * cos_theta - mom_z[j] * sin_theta + p_rot2_z = p_rot_x * sin_theta + mom_z[j] * cos_theta + p_rot2_y = p_rot_y + + # Calculate both theta and phi in rotated frame + pt_rot = sqrt(p_rot2_x^2 + p_rot2_y^2) + jet_theta[j] = atan(pt_rot, p_rot2_z) + jet_phi[j] = atan(p_rot2_y, p_rot2_x) + end + + theta_result[i] = jet_theta + phi_result[i] = jet_phi + end + + return (theta_result, phi_result) +end + +### Special Measurements (2) + +# get_dndx - dE/dx measurement (energy loss) +# get_mtof - Mass from time-of-flight measurement + +""" + get_dndx(jets_constituents::Vector{JetConstituents}, + dNdx::StructVector{EDM4hep.Quantity}, + trackdata::StructVector{EDM4hep.Track}, + JetsConstituents_isChargedHad::Vector{JetConstituentsData}) -> Vector{JetConstituentsData} + +Calculate dE/dx or dN/dx for each charged hadron in jets. Neutrals, muons, and electrons are set to 0. +Only charged hadrons have valid dN/dx values. + +# Arguments +- jets_constituents: Vector of jet constituents (each element contains particles for one jet) +- dNdx: StructVector of Quantity objects containing dN/dx measurements (EFlowTrack_2) +- trackdata: StructVector of Track objects (EFlowTrack) +- JetsConstituents_isChargedHad: Vector of vectors indicating which particles are charged hadrons + +# Returns +Vector of vectors of dN/dx values (one vector per jet, one value per constituent) +""" +function get_dndx(jets_constituents::Vector{<:JetConstituents}, + dNdx::StructVector{EDM4hep.Quantity}, + trackdata::StructVector{EDM4hep.Track}, + JetsConstituents_isChargedHad::Vector{Vector{Float32}}) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + tracks_len = length(trackdata) + + @inbounds for i in 1:n_jets + jet = jets_constituents[i] + tracks_first = jet.tracks + isChargedHad = JetsConstituents_isChargedHad[i] + n_constituents = length(jet) + tmp = Vector{Float32}(undef, n_constituents) + + @simd ivdep for j in 1:n_constituents + has_valid_track = tracks_first[j].first + 1 <= tracks_len + is_charged_had = isChargedHad[j] == 1.0f0 + tmp[j] = ifelse(has_valid_track & is_charged_had, -1.0f0, 0.0f0) + end + + result[i] = tmp + end + + return result +end + +function get_mtof(jets_constituents::Vector{<:JetConstituents}, + track_L::AbstractArray{T} where {T <: Float32}, + trackdata::StructVector{EDM4hep.Track}, + trackerhits::StructVector{EDM4hep.TrackerHit}, + gammadata::StructVector{EDM4hep.Cluster}, + nhdata::StructVector{EDM4hep.Cluster}, + calohits::StructVector{EDM4hep.CalorimeterHit}, + V::LorentzVector) + n_jets = length(jets_constituents) + result = Vector{Vector{Float32}}(undef, n_jets) + + # Pre-compute limits + tracks_len = length(trackdata) + gamma_len = length(gammadata) + nh_len = length(nhdata) + cluster_limit = nh_len + gamma_len + + # Pre-compute vertex values + vx, vy, vz = V.x, V.y, V.z + v_t_scaled = V.t * MM_TO_M * C_LIGHT_INV # Tin calculation + + @inbounds for i in 1:n_jets + single_jet_constituents = jets_constituents[i] + n_constituents = length(single_jet_constituents) + + # Pre-allocate for this jet + tmp = Vector{Float32}(undef, n_constituents) + result_idx = 0 + + # Access fields once + clusters_first = single_jet_constituents.clusters + tracks_first = single_jet_constituents.tracks + types = single_jet_constituents.type + charges = single_jet_constituents.charge + masses = single_jet_constituents.mass + energies = single_jet_constituents.energy + mom_x = single_jet_constituents.momentum.x + mom_y = single_jet_constituents.momentum.y + mom_z = single_jet_constituents.momentum.z + + for j in 1:n_constituents + cluster_idx = clusters_first[j].first + track_idx = tracks_first[j].first + particle_type = types[j] + + mass_calculated = INVALID_MASS # Invalid marker + + # Handle cluster-based particles + if cluster_idx < cluster_limit + if particle_type == PDG_K_LONG # Neutral hadron + nh_idx = cluster_idx + 1 - gamma_len + hit_idx = nhdata[nh_idx].hits.first + 1 + + # Get hit data + hit = calohits[hit_idx] + tof = hit.time + + # Calculate distance + dx = hit.position.x - vx + dy = hit.position.y - vy + dz = hit.position.z - vz + l = sqrt(dx^2 + dy^2 + dz^2) * MM_TO_M + + beta = l / (tof * C_LIGHT) + + if 0.0f0 < beta < 1.0f0 + energy = energies[j] + mass_calculated = energy * sqrt(1.0f0 - beta^2) + else + mass_calculated = INVALID_TOF_MASS # Invalid measurement + end + + elseif particle_type == PDG_PHOTON # Photon + mass_calculated = 0.0f0 + end + end + + # Handle track-based particles (only if not already calculated) + if mass_calculated < 0.0f0 && track_idx < tracks_len + charge = charges[j] + if abs(charge) > 0.0f0 + mass = masses[j] + + # Check for known particles + if abs(mass - ELECTRON_MASS) < ELECTRON_TOLERANCE + mass_calculated = ELECTRON_MASS + elseif abs(mass - MUON_MASS) < MUON_TOLERANCE + mass_calculated = MUON_MASS + else + # Calculate from time of flight + track = trackdata[track_idx + 1] + last_hit_idx = track.trackerHits.last + Tout = trackerhits[last_hit_idx].time + tof = Tout - v_t_scaled + + l = track_L[track_idx + 1] * MM_TO_M + beta = l / (tof * C_LIGHT) + + if 0.0f0 < beta < 1.0f0 + # Calculate momentum magnitude + p = sqrt(mom_x[j]^2 + mom_y[j]^2 + mom_z[j]^2) + mass_calculated = p * sqrt(1.0f0 / (beta^2) - 1.0f0) + else + mass_calculated = PION_MASS # Default + end + end + end + end + + # Store result if we calculated a mass + if mass_calculated >= 0.0f0 + result_idx += 1 + tmp[result_idx] = mass_calculated + end + end + + # Resize to actual number of particles with calculated mass + result[i] = resize!(tmp, result_idx) + end + + return result +end + +### Impact Parameters and Jet Distance (12) + +## 2D Impact Parameter +# get_Sip2dVal_clusterV - Vectorized version 2D signed impact parameter value for clustered jets +# get_Sip2dSig - 2D impact parameter significance + +# 3D Impact Parameter +# get_Sip3dVal_clusterV - Vectorized version 3D signed impact parameter value for clustered jets +# get_Sip3dSig - 3D impact parameter significance + +# Jet Distance +# get_JetDistVal_clusterV - Vectorized version jet distance value for clustered jets +# get_JetDistSig - Jet distance significance + +""" + get_Sip2dVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + D0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the 2D signed impact parameter value for each particle relative to the jet axis. +This is a Julia implementation of the C++ function get_Sip2dVal_clusterV. + +# Arguments +- jets: Vector of EEJet objects representing jets +- D0: Vector of vectors containing D0 values (transverse impact parameters) +- phi0: Vector of vectors containing phi0 values (azimuthal angles at impact point) +- Bz: The magnetic field in Tesla + +# Returns +Vector of vectors of 2D signed impact parameter values (one vector per jet) +""" +function get_Sip2dVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + D0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) + n_jets = length(jets) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for i in 1:n_jets + px = Float32(jets[i].px) + py = Float32(jets[i].py) + d0_vals = D0[i] + phi_vals = phi0[i] + n_constituents = length(d0_vals) + + sip2d_values = Vector{Float32}(undef, n_constituents) + + @inbounds for j in 1:n_constituents + d0_val = d0_vals[j] + phi_val = phi_vals[j] + + sin_phi, cos_phi = sincos(phi_val) + + d0x = -d0_val * sin_phi + d0y = d0_val * cos_phi + + dot_product = d0x * px + d0y * py + + abs_d0 = abs(d0_val) + sign_dot = sign(dot_product) + signed_ip = sign_dot * abs_d0 + + is_valid = Float32(d0_val != UNDEF_VAL) + sip2d_values[j] = is_valid * signed_ip + (1.0f0 - is_valid) * UNDEF_VAL + end + + result[i] = sip2d_values + end + + return result +end + +""" + get_btagSip2dVal(jets::Vector{JetReconstruction.EEJet}, + pfcand_dxy::Vector{JetConstituentsData}, + pfcand_phi0::Vector{JetConstituentsData}, + Bz::Float32) -> Vector{JetConstituentsData} + +Call the implementation function get_Sip2dVal_clusterV +""" +function get_btagSip2dVal(jets::Vector{JetReconstruction.EEJet}, + pfcand_dxy::Vector{JetConstituentsData}, + pfcand_phi0::Vector{JetConstituentsData}, + Bz::Float32) + # Simply call the implementation function + return get_Sip2dVal_clusterV(jets, pfcand_dxy, pfcand_phi0, Bz) +end + +""" + get_Sip2dSig(Sip2dVals::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}) -> Vector{JetConstituentsData} + +Calculate the 2D signed impact parameter significance for each particle. +This is a Julia implementation of the C++ function get_Sip2dSig. + +# Arguments +- Sip2dVals: Vector of vectors containing 2D signed impact parameter values +- err2_D0: Vector of vectors containing squared errors of the D0 values + +# Returns +Vector of vectors of 2D signed impact parameter significances (one vector per jet) +""" +function get_Sip2dSig(Sip2dVals::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}) + n_jets = length(Sip2dVals) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for i in 1:n_jets + n_constituents = length(Sip2dVals[i]) + sig_values = Vector{Float32}(undef, n_constituents) + sip_vals = Sip2dVals[i] + err_vals = err2_D0[i] + + @simd for j in 1:n_constituents + err_val = err_vals[j] + sip_val = sip_vals[j] + + valid = err_val > 0.0f0 + sqrt_err = sqrt(max(err_val, eps(Float32))) # Avoid sqrt of negative + sig = sip_val / sqrt_err + + sig_values[j] = valid ? sig : UNDEF_VAL + end + + result[i] = sig_values + end + + return result +end + +""" + get_btagSip2dSig(pfcand_btagSip2dVal::Vector{JetConstituentsData}, + pfcand_dxydxy::Vector{JetConstituentsData}) -> Vector{JetConstituentsData} + +Call the implementation function get_Sip2dSig +""" +function get_btagSip2dSig(pfcand_btagSip2dVal::Vector{JetConstituentsData}, + pfcand_dxydxy::Vector{JetConstituentsData}) + # Simply call the implementation function + return get_Sip2dSig(pfcand_btagSip2dVal, pfcand_dxydxy) +end + +""" + get_Sip3dVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + D0::Vector{JetConstituentsData}, + Z0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the 3D signed impact parameter value for each particle relative to the jet axis. +""" +function get_Sip3dVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + D0::Vector{JetConstituentsData}, + Z0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) + n_jets = length(jets) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for i in 1:n_jets + px = Float32(jets[i].px) + py = Float32(jets[i].py) + pz = Float32(jets[i].pz) + d0_vals = D0[i] + z0_vals = Z0[i] + phi_vals = phi0[i] + n_constituents = length(d0_vals) + + cprojs = Vector{Float32}(undef, n_constituents) + + @inbounds for j in 1:n_constituents + d0_val = d0_vals[j] + z0_val = z0_vals[j] + phi_val = phi_vals[j] + + sin_phi = sin(phi_val) + cos_phi = cos(phi_val) + + dx = -d0_val * sin_phi + dy = d0_val * cos_phi + dz = z0_val + + dot_prod = dx * px + dy * py + dz * pz + + magnitude = sqrt(d0_val * d0_val + z0_val * z0_val) + sign_dot = sign(dot_prod) + signed_ip = sign_dot * magnitude + + is_valid = Float32(d0_val != UNDEF_VAL) + cprojs[j] = is_valid * signed_ip + (1.0f0 - is_valid) * UNDEF_VAL + end + + result[i] = cprojs + end + + return result +end + +function get_btagSip3dVal(jets::Vector{JetReconstruction.EEJet}, + pfcand_dxy::Vector{JetConstituentsData}, + pfcand_dz::Vector{JetConstituentsData}, + pfcand_phi0::Vector{JetConstituentsData}, + Bz::Float32) + # Simply call the implementation function + return get_Sip3dVal_clusterV(jets, pfcand_dxy, pfcand_dz, pfcand_phi0, Bz) +end + +""" + get_Sip3dSig(Sip3dVals::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) -> Vector{JetConstituentsData} + +Calculate the 3D signed impact parameter significance (value/error) for each particle. +""" +function get_Sip3dSig(Sip3dVals::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) + n_jets = length(Sip3dVals) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for i in 1:n_jets + n_constituents = length(Sip3dVals[i]) + sigs = Vector{Float32}(undef, n_constituents) + sip_vals = Sip3dVals[i] + err_d0 = err2_D0[i] + err_z0 = err2_Z0[i] + + @simd for j in 1:n_constituents + err_d0_val = err_d0[j] + err_z0_val = err_z0[j] + sip_val = sip_vals[j] + + # Branchless computation + valid = err_d0_val > 0.0f0 + err_sum = err_d0_val + err_z0_val + sqrt_err = sqrt(max(err_sum, eps(Float32))) + sig = sip_val / sqrt_err + + sigs[j] = valid ? sig : UNDEF_VAL + end + + result[i] = sigs + end + + return result +end + +function get_btagSip3dSig(Sip3dVals::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) + # Simply call the implementation function + return get_Sip3dSig(Sip3dVals, err2_D0, err2_Z0) +end + +""" + get_JetDistVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + jets_constituents::Vector{<:JetConstituents}, + D0::Vector{JetConstituentsData}, + Z0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) -> Vector{JetConstituentsData} + +Calculate the jet distance value for each particle, measuring the distance between +the point of closest approach and the jet axis. +""" +function get_JetDistVal_clusterV(jets::Vector{JetReconstruction.EEJet}, + jets_constituents::Vector{<:JetConstituents}, + D0::Vector{JetConstituentsData}, + Z0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) + n_jets = length(jets) + result = Vector{JetConstituentsData}(undef, n_jets) + + for i in 1:n_jets + px_jet, py_jet, pz_jet = jets[i].px, jets[i].py, jets[i].pz + single_jet_constituents = jets_constituents[i] + n_constituents = length(single_jet_constituents) + tmp = Vector{Float32}(undef, n_constituents) + + for j in 1:n_constituents + if D0[i][j] != UNDEF_VAL + d0_val = D0[i][j] + z0_val = Z0[i][j] + + # Use sincos for efficiency + sin_phi, cos_phi = sincos(phi0[i][j]) + + # Impact parameter vector + dx = -d0_val * sin_phi + dy = d0_val * cos_phi + dz = z0_val + + # Constituent momentum + px_ct = single_jet_constituents[j].momentum.x + py_ct = single_jet_constituents[j].momentum.y + pz_ct = single_jet_constituents[j].momentum.z + + # Cross product: n = p_ct × p_jet + nx = py_ct * pz_jet - pz_ct * py_jet + ny = pz_ct * px_jet - px_ct * pz_jet + nz = px_ct * py_jet - py_ct * px_jet + + # Normalize + n_mag = sqrt(nx^2 + ny^2 + nz^2) + inv_n_mag = 1.0f0 / max(n_mag, eps(Float32)) + nx *= inv_n_mag + ny *= inv_n_mag + nz *= inv_n_mag + + # Distance (r_jet = [0,0,0], so we just need n·d) + tmp[j] = nx * dx + ny * dy + nz * dz + else + tmp[j] = UNDEF_VAL + end + end + + result[i] = tmp + end + + return result +end + +function get_btagJetDistVal(jets::Vector{JetReconstruction.EEJet}, + jets_constituents::Vector{<:JetConstituents}, + D0::Vector{JetConstituentsData}, + Z0::Vector{JetConstituentsData}, + phi0::Vector{JetConstituentsData}, + Bz::Float32) + # Simply call the implementation function + return get_JetDistVal_clusterV(jets, jets_constituents, D0, Z0, phi0, Bz) +end + +""" + get_JetDistSig(JetDistVal::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) -> Vector{JetConstituentsData} + +Calculate the jet distance significance (value/error) for each particle. +""" +function get_JetDistSig(JetDistVal::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) + n_jets = length(JetDistVal) + result = Vector{JetConstituentsData}(undef, n_jets) + + @inbounds for i in 1:n_jets + n_constituents = length(JetDistVal[i]) + tmp = Vector{Float32}(undef, n_constituents) + jet_vals = JetDistVal[i] + err_d0 = err2_D0[i] + err_z0 = err2_Z0[i] + + @simd for j in 1:n_constituents + err_d0_val = err_d0[j] + err_z0_val = err_z0[j] + jet_val = jet_vals[j] + + # Branchless computation + valid = err_d0_val > 0.0f0 + err_sum = err_d0_val + err_z0_val + sqrt_err = sqrt(max(err_sum, eps(Float32))) + sig = jet_val / sqrt_err + + tmp[j] = valid ? sig : UNDEF_VAL + end + + result[i] = tmp + end + + return result +end + +function get_btagJetDistSig(JetDistVal::Vector{JetConstituentsData}, + err2_D0::Vector{JetConstituentsData}, + err2_Z0::Vector{JetConstituentsData}) + # Simply call the implementation function + return get_JetDistSig(JetDistVal, err2_D0, err2_Z0) +end diff --git a/ext/JetFlavourTagging/JetFlavourHelper.jl b/ext/JetFlavourTagging/JetFlavourHelper.jl new file mode 100644 index 00000000..4be13924 --- /dev/null +++ b/ext/JetFlavourTagging/JetFlavourHelper.jl @@ -0,0 +1,684 @@ +module JetFlavourHelper + +using JSON +using ONNXRunTime +using StructArrays: StructVector +using JetReconstruction +using EDM4hep +using LorentzVectorHEP + +# Utility functions for the jet features from EDM4hep +include("JetConstituentUtils.jl") + +""" + JetFlavourHelper + +A module for jet flavour identification using neural networks. +""" + +""" + setup_onnx_runtime(onnx_path::AbstractString, json_path::AbstractString) -> ONNXRunTime.InferenceSession + +Setup the ONNX model and preprocessing configuration for jet flavour tagging. + +# Arguments +- `onnx_path`: Path to the ONNX model file +- `json_path`: Path to the JSON configuration file + +# Returns +An ONNX inference session for the loaded model +""" +function setup_onnx_runtime(onnx_path::AbstractString, json_path::AbstractString) + # Load JSON configuration + config = JSON.parsefile(json_path) + model = ONNXRunTime.load_inference(onnx_path) + + return model, config +end + +""" + normalize_feature(value::Float32, info::Dict) -> Float32 + +Normalize a feature value based on the preprocessing information. + +# Arguments +- `value`: Raw feature value +- `info`: Dictionary containing normalization parameters + +# Returns +Normalized feature value +""" +function normalize_feature(value::Float32, info::Dict) + if value == -9.0f0 + return 0.0f0 # Replace -9.0 (missing value) with 0 + end + + # Apply normalization using median and norm_factor + normalized = (value - info["median"]) * info["norm_factor"] + + # Clamp to specified bounds + return clamp(normalized, info["lower_bound"], info["upper_bound"]) +end + +""" + prepare_input_tensor(jets_constituents::Vector{StructVector{EDM4hep.ReconstructedParticle}}, + jets::Vector{EEJet}, + config::Dict, + feature_data::Dict, + jet_index::Int=1) -> Dict{String, Array} + +Prepare input tensors for the neural network from jet constituents. + +# Arguments +- `jets_constituents`: Vector of jet constituents (structured as a vector of StructVector of ReconstructedParticle) +- `jets`: Vector of jets (EEJet) +- `config`: JSON configuration for preprocessing +- `feature_data`: Dictionary containing all extracted features +- `jet_index`: Index of the jet to process (default 1) + +# Returns +Dictionary of input tensors +""" +function prepare_input_tensor(jets_constituents::Vector{<:JetConstituents}, + jets::Vector{EEJet}, + config::Dict, + feature_data::Dict, + jet_index::Int = 1) + + # Get input names and variable info + input_names = config["input_names"] + + # Initialize input tensor dictionary + input_tensors = Dict{String, Array{Float32}}() + + # Get max length for padding + max_length = config["pf_points"]["var_length"] + + # Initialize tensors for single jet processing + for input_name in input_names + if input_name == "pf_features" + feature_vars = length(config[input_name]["var_names"]) + input_tensors[input_name] = zeros(Float32, 1, feature_vars, max_length) + elseif input_name == "pf_vectors" + vector_vars = length(config[input_name]["var_names"]) + input_tensors[input_name] = zeros(Float32, 1, vector_vars, max_length) + elseif input_name == "pf_mask" + input_tensors[input_name] = zeros(Float32, 1, 1, max_length) + end + end + + # Note: This function prepares tensors for a single jet at a time + # The caller can loop through jets and process them individually + # A batch processing version is planned for next version under the inference() function. + + # Process the specified jet + if length(jets) >= jet_index && jet_index > 0 + i = jet_index + jet = jets[1] # We're processing single jets, so always use first element + + # Fill each tensor for this jet + constituents = jets_constituents[1] # Single jet passed + num_constituents = min(length(constituents), max_length) + + # Fill mask (1 for valid constituents, 0 for padding) + if haskey(feature_data, "pf_mask") + for j in 1:num_constituents + input_tensors["pf_mask"][1, 1, j] = 1.0f0 + end + end + + # Fill points + if haskey(feature_data, "pf_points") && haskey(input_tensors, "pf_points") + for (var_idx, var_name) in enumerate(config["pf_points"]["var_names"]) + var_info = config["pf_points"]["var_infos"][var_name] + + for j in 1:num_constituents + if j <= length(feature_data["pf_points"][var_name][i]) + raw_value = feature_data["pf_points"][var_name][i][j] + norm_value = normalize_feature(raw_value, var_info) + input_tensors["pf_points"][1, var_idx, j] = norm_value + end + end + end + end + + # Fill features + if haskey(feature_data, "pf_features") && haskey(input_tensors, "pf_features") + for (var_idx, var_name) in enumerate(config["pf_features"]["var_names"]) + var_info = config["pf_features"]["var_infos"][var_name] + + for j in 1:num_constituents + if haskey(feature_data["pf_features"], var_name) && + j <= length(feature_data["pf_features"][var_name][i]) + raw_value = feature_data["pf_features"][var_name][i][j] + norm_value = normalize_feature(raw_value, var_info) + input_tensors["pf_features"][1, var_idx, j] = norm_value + end + end + end + end + + # Fill vectors (energies, momenta) + if haskey(feature_data, "pf_vectors") && haskey(input_tensors, "pf_vectors") + for (var_idx, var_name) in enumerate(config["pf_vectors"]["var_names"]) + for j in 1:num_constituents + if haskey(feature_data["pf_vectors"], var_name) && + j <= length(feature_data["pf_vectors"][var_name][i]) + input_tensors["pf_vectors"][1, var_idx, j] = feature_data["pf_vectors"][var_name][i][j] + end + end + end + end + end + + return input_tensors +end + +""" + get_weights(slot::Int, vars::Dict{String, Vector{Vector{Float32}}}, + jets::Vector{EEJet}, json_config::Dict, model::ONNXRunTime.InferenceSession) -> Vector{Vector{Float32}} + +Compute jet flavour probabilities for each jet. + +# Arguments +- `slot`: Threading slot +- `vars`: Dictionary containing all features for jet constituents +- `jets`: Vector of jets +- `json_config`: JSON configuration for preprocessing +- `model`: ONNX inference session + +# Returns +Vector of flavour probabilities for each jet +""" +function get_weights(slot::Int, vars::Dict{String, Dict{String, Vector{Vector{Float32}}}}, + jets::Vector{EEJet}, jets_constituents::Vector{<:JetConstituents}, + json_config::Dict, model::ONNXRunTime.InferenceSession) + + # The model processes one jet at a time + result = Vector{Vector{Float32}}() + + # Process each jet individually + for i in 1:length(jets) + # Create single-jet arrays + single_jet = [jets[i]] + single_constituents = [jets_constituents[i]] + + # Create single-jet feature data by extracting only features for this jet + single_jet_vars = Dict{String, Dict{String, Vector{Vector{Float32}}}}() + for (category, features) in vars + single_jet_vars[category] = Dict{String, Vector{Vector{Float32}}}() + for (fname, fvalues) in features + # Extract only the features for jet i + if i <= length(fvalues) + single_jet_vars[category][fname] = [fvalues[i]] + else + # If no features for this jet, create empty array + single_jet_vars[category][fname] = [Float32[]] + end + end + end + + # Prepare input tensor for this single jet with extracted features + input_tensors = prepare_input_tensor(single_constituents, single_jet, json_config, + single_jet_vars, 1) + + # Run inference + output = model(input_tensors) + + # Extract probabilities + probabilities = output["softmax"] + + # Get probabilities for this jet + num_classes = size(probabilities, 2) + jet_probs = Vector{Float32}(undef, num_classes) + for c in 1:num_classes + jet_probs[c] = probabilities[1, c] # Always index 1 since we process one jet at a time + end + push!(result, jet_probs) + end + + return result +end + +""" + get_weight(jet_weights::Vector{Vector{Float32}}, weight_idx::Int) -> Vector{Float32} + +Extract a specific weight/score from the jet weights. + +# Arguments +- `jet_weights`: Vector of weight vectors for each jet +- `weight_idx`: Index of the weight to extract + +# Returns +Vector of the specified weight for each jet +""" +function get_weight(jet_weights::Vector{Vector{Float32}}, weight_idx::Int) + if weight_idx < 0 + error("Invalid index requested for jet flavour weight.") + end + + result = Vector{Float32}() + + for jet_weight in jet_weights + if weight_idx >= length(jet_weight) + error("Flavour weight index exceeds the number of weights registered.") + end + + push!(result, jet_weight[weight_idx + 1]) # +1 for Julia's 1-based indexing + end + + return result +end + +""" + inference(json_config_path::AbstractString, onnx_model_path::AbstractString, df::DataFrame, + jets::Vector{EEJet}, jets_constituents::Vector{StructVector{EDM4hep.ReconstructedParticle}}, + feature_data::Dict) -> DataFrame + +Run flavour tagging inference on a collection of jets. + +# Arguments +- `json_config_path`: Path to the JSON configuration file +- `onnx_model_path`: Path to the ONNX model file +- `jets`: Vector of jets +- `jets_constituents`: Vector of jet constituents +- `feature_data`: Dictionary containing all extracted features + +# Returns +DataFrame with added flavour tagging scores +""" +function inference(json_config_path::AbstractString, onnx_model_path::AbstractString, + jets::Vector{EEJet}, + jets_constituents::Vector{StructVector{EDM4hep.ReconstructedParticle}}, + feature_data::Dict) + + # Extract input variables/score names from JSON file + initvars = String[] + variables = String[] + scores = String[] + + config = JSON.parsefile(json_config_path) + + # Extract feature names + for varname in config["pf_features"]["var_names"] + push!(initvars, varname) + push!(variables, varname) + end + + # Extract vector names + for varname in config["pf_vectors"]["var_names"] + push!(initvars, varname) + push!(variables, varname) + end + + # Extract output names + for scorename in config["output_names"] + push!(scores, scorename) + end + + # Setup model + model, _ = setup_onnx_runtime(onnx_model_path, json_config_path, initvars) + + # Run inference + weights = get_weights(0, feature_data, jets, jets_constituents, config, model) + + # Extract individual scores + jet_scores = Dict{String, Vector{Float32}}() + + for (i, scorename) in enumerate(scores) + jet_scores[scorename] = get_weight(weights, i - 1) # Adjust for 0-based indexing in get_weight + end + + return jet_scores +end + +# TODO: Add primary vertex as an argument (from MC Particle) +""" + extract_features(jets::Vector{EEJet}, jets_constituents::Vector{<:JetConstituents}, + tracks::AbstractVector{EDM4hep.TrackState}, bz::Float32, + track_L::AbstractArray{T} where T <: AbstractFloat, + json_config::Dict, + trackdata::AbstractVector{EDM4hep.Track}=AbstractVector{EDM4hep.Track}(), + trackerhits::AbstractVector{EDM4hep.TrackerHit}=AbstractVector{EDM4hep.TrackerHit}(), + gammadata::AbstractVector{EDM4hep.Cluster}=AbstractVector{EDM4hep.Cluster}(), + nhdata::AbstractVector{EDM4hep.Cluster}=AbstractVector{EDM4hep.Cluster}(), + calohits::AbstractVector{EDM4hep.CalorimeterHit}=AbstractVector{EDM4hep.CalorimeterHit}(), + dNdx::AbstractVector{EDM4hep.Quantity}=AbstractVector{EDM4hep.Quantity}()) -> Dict + +Extract features for jet flavour tagging based on JSON configuration. + +# Arguments +- `jets`: Vector of jets (EEJet) +- `jets_constituents`: Vector of jet constituents +- `tracks`: StructVector of track states +- `bz`: Magnetic field strength +- `track_L`: Array of track lengths +- `json_config`: JSON configuration dict specifying which features to extract (required) +- `trackdata`: Vector of track data (optional) +- `trackerhits`: Vector of tracker hits (optional) +- `gammadata`: Vector of gamma clusters (optional) +- `nhdata`: Vector of neutral hadron clusters (optional) +- `calohits`: Vector of calorimeter hits (optional) +- `dNdx`: Vector of dE/dx measurements (optional) +- `mc_vertices`: Vector of MC vertices for each reconstructed particle (optional) + +# Returns +Dictionary containing extracted features as specified in the JSON configuration. +""" +function extract_features(jets::Vector{EEJet}, jets_constituents::Vector{<:JetConstituents}, + tracks::AbstractVector{EDM4hep.TrackState}, bz::Float32, + track_L::AbstractArray{T} where {T <: AbstractFloat}, + json_config::Dict, + trackdata::AbstractVector{EDM4hep.Track} = AbstractVector{EDM4hep.Track}(), + trackerhits::AbstractVector{EDM4hep.TrackerHit} = AbstractVector{EDM4hep.TrackerHit}(), + gammadata::AbstractVector{EDM4hep.Cluster} = AbstractVector{EDM4hep.Cluster}(), + nhdata::AbstractVector{EDM4hep.Cluster} = AbstractVector{EDM4hep.Cluster}(), + calohits::AbstractVector{EDM4hep.CalorimeterHit} = AbstractVector{EDM4hep.CalorimeterHit}(), + dNdx::AbstractVector{EDM4hep.Quantity} = AbstractVector{EDM4hep.Quantity}(), + mc_vertices::Union{Nothing, Vector{LorentzVector{Float32}}} = nothing) + + # Primary vertex for displacement calculations + # Use provided MC vertices or default to (0,0,0,0) + # If mc_vertices are provided, find the most common vertex (primary vertex) + if isnothing(mc_vertices) || isempty(mc_vertices) + v_in = LorentzVector(0.0, 0.0, 0.0, 0.0) + else + # Find the most common vertex (likely the primary vertex) + # For simplicity, use the first non-zero vertex found + v_in = LorentzVector(0.0, 0.0, 0.0, 0.0) + for vertex in mc_vertices + if vertex.x != 0.0 || vertex.y != 0.0 || vertex.z != 0.0 + v_in = vertex + break + end + end + end + + # Initialize feature containers + features = Dict{String, Dict{String, Vector{Vector{Float32}}}}() + + # Cache for computed features to avoid redundant calculations + computed_features = Dict{String, Vector{Vector{Float32}}}() + + # Create kwargs dict with all available data + kwargs = Dict(:jets => jets, + :tracks => tracks, + :v_in => v_in, + :bz => bz, + :track_L => track_L, + :trackdata => trackdata, + :trackerhits => trackerhits, + :gammadata => gammadata, + :nhdata => nhdata, + :calohits => calohits, + :dNdx => dNdx) + + # Process each input type specified in the JSON + for input_name in json_config["input_names"] + if !haskey(json_config, input_name) || !haskey(json_config[input_name], "var_names") + continue + end + + features[input_name] = Dict{String, Vector{Vector{Float32}}}() + + # Extract each requested variable + for var_name in json_config[input_name]["var_names"] + # Remove "pfcand_" prefix if present + clean_name = replace(var_name, "pfcand_" => "") + + # Check if already computed + if haskey(computed_features, clean_name) + features[input_name][var_name] = computed_features[clean_name] + continue + end + + # Get the extractor function + extractor = get_feature_extractor(clean_name) + if extractor === nothing + # Feature not implemented, create empty arrays + features[input_name][var_name] = [Float32[] for _ in 1:length(jets)] + continue + end + + # Handle dependencies for certain features + deps = get_feature_dependencies(clean_name) + if !isempty(deps) + dep_kwargs = Dict{Symbol, Any}() + + for (dep_key, dep_name) in deps + if !haskey(computed_features, dep_name) + # Recursively compute dependency + dep_extractor = get_feature_extractor(dep_name) + if dep_extractor !== nothing + # Check if this dependency has its own dependencies + sub_deps = get_feature_dependencies(dep_name) + if !isempty(sub_deps) + sub_dep_kwargs = Dict{Symbol, Any}() + for (sub_dep_key, sub_dep_name) in sub_deps + if !haskey(computed_features, sub_dep_name) + sub_dep_extractor = get_feature_extractor(sub_dep_name) + if sub_dep_extractor !== nothing + computed_features[sub_dep_name] = sub_dep_extractor(jets_constituents; + kwargs...) + end + end + if haskey(computed_features, sub_dep_name) + sub_dep_kwargs[sub_dep_key] = computed_features[sub_dep_name] + end + end + # Compute dependency with its dependencies + temp_kwargs = merge(kwargs, sub_dep_kwargs) + computed_features[dep_name] = dep_extractor(jets_constituents; + temp_kwargs...) + else + # No sub-dependencies, compute directly + computed_features[dep_name] = dep_extractor(jets_constituents; + kwargs...) + end + end + end + if haskey(computed_features, dep_name) + dep_kwargs[dep_key] = computed_features[dep_name] + end + end + + # Merge dependency kwargs with main kwargs + merge!(kwargs, dep_kwargs) + end + + # Special handling for features that need isChargedHad + if clean_name == "dndx" && !haskey(kwargs, :jets_constituents_isChargedHad) + if !haskey(computed_features, "isChargedHad") + computed_features["isChargedHad"] = get_is_charged_had(jets_constituents) + end + kwargs[:jets_constituents_isChargedHad] = computed_features["isChargedHad"] + end + + # Extract the feature + try + feature_data = extractor(jets_constituents; kwargs...) + computed_features[clean_name] = feature_data + features[input_name][var_name] = feature_data + catch e + # If extraction fails, create empty arrays + @warn "Failed to extract feature $var_name: $e" + features[input_name][var_name] = [Float32[] for _ in 1:length(jets)] + end + end + end + + return features +end + +""" + get_feature_extractor(feature_name::String) + +Return the appropriate feature extraction function for a given feature name. + +# Arguments +- `feature_name`: Name of the feature (without "pfcand_" prefix) + +# Returns +A function that can extract the specified feature, or nothing if not found +""" +function get_feature_extractor(feature_name::String) + # Map feature names to extraction functions + feature_map = Dict( + # Basic kinematic features + "e" => (jets_constituents; kwargs...) -> get_e(jets_constituents), + "p" => (jets_constituents; kwargs...) -> get_p(jets_constituents), + "theta" => (jets_constituents; kwargs...) -> get_theta(jets_constituents), + "phi" => (jets_constituents; kwargs...) -> get_phi(jets_constituents), + "charge" => (jets_constituents; kwargs...) -> get_charge(jets_constituents), + "type" => (jets_constituents; kwargs...) -> get_type(jets_constituents), + + # Relative kinematic features + "erel_log" => (jets_constituents; jets = nothing, kwargs...) -> get_erel_log_cluster(jets, + jets_constituents), + "thetarel" => (jets_constituents; jets = nothing, kwargs...) -> get_thetarel_cluster(jets, + jets_constituents), + "phirel" => (jets_constituents; jets = nothing, kwargs...) -> get_phirel_cluster(jets, + jets_constituents), + + # Particle ID features + "isMu" => (jets_constituents; kwargs...) -> get_is_mu(jets_constituents), + "isEl" => (jets_constituents; kwargs...) -> get_is_el(jets_constituents), + "isChargedHad" => (jets_constituents; kwargs...) -> get_is_charged_had(jets_constituents), + "isGamma" => (jets_constituents; kwargs...) -> get_is_gamma(jets_constituents), + "isNeutralHad" => (jets_constituents; kwargs...) -> get_is_neutral_had(jets_constituents), + + # Track parameters + "dxy" => (jets_constituents; tracks = nothing, v_in = nothing, bz = nothing, kwargs...) -> get_dxy(jets_constituents, + tracks, + v_in, + bz), + "dz" => (jets_constituents; tracks = nothing, v_in = nothing, bz = nothing, kwargs...) -> get_dz(jets_constituents, + tracks, + v_in, + bz), + "phi0" => (jets_constituents; tracks = nothing, v_in = nothing, bz = nothing, kwargs...) -> get_phi0(jets_constituents, + tracks, + v_in, + bz), + + # Covariance matrix elements + "dptdpt" => (jets_constituents; tracks = nothing, kwargs...) -> get_dptdpt(jets_constituents, + tracks), + "detadeta" => (jets_constituents; tracks = nothing, kwargs...) -> get_detadeta(jets_constituents, + tracks), + "dphidphi" => (jets_constituents; tracks = nothing, kwargs...) -> get_dphidphi(jets_constituents, + tracks), + "dxydxy" => (jets_constituents; tracks = nothing, kwargs...) -> get_dxydxy(jets_constituents, + tracks), + "dzdz" => (jets_constituents; tracks = nothing, kwargs...) -> get_dzdz(jets_constituents, + tracks), + "dxydz" => (jets_constituents; tracks = nothing, kwargs...) -> get_dxydz(jets_constituents, + tracks), + "dphidxy" => (jets_constituents; tracks = nothing, kwargs...) -> get_dphidxy(jets_constituents, + tracks), + "dlambdadz" => (jets_constituents; tracks = nothing, kwargs...) -> get_dlambdadz(jets_constituents, + tracks), + "dxyc" => (jets_constituents; tracks = nothing, kwargs...) -> get_dxyc(jets_constituents, + tracks), + "dxyctgtheta" => (jets_constituents; tracks = nothing, kwargs...) -> get_dxyctgtheta(jets_constituents, + tracks), + "phic" => (jets_constituents; tracks = nothing, kwargs...) -> get_phic(jets_constituents, + tracks), + "phidz" => (jets_constituents; tracks = nothing, kwargs...) -> get_phidz(jets_constituents, + tracks), + "phictgtheta" => (jets_constituents; tracks = nothing, kwargs...) -> get_phictgtheta(jets_constituents, + tracks), + "cdz" => (jets_constituents; tracks = nothing, kwargs...) -> get_cdz(jets_constituents, + tracks), + "cctgtheta" => (jets_constituents; tracks = nothing, kwargs...) -> get_cctgtheta(jets_constituents, + tracks), + + # B-tagging features + "btagSip2dVal" => (jets_constituents; jets = nothing, dxy = nothing, phi0 = nothing, bz = nothing, kwargs...) -> get_btagSip2dVal(jets, + dxy, + phi0, + bz), + "btagSip2dSig" => (jets_constituents; btagSip2dVal = nothing, dxydxy = nothing, kwargs...) -> get_btagSip2dSig(btagSip2dVal, + dxydxy), + "btagSip3dVal" => (jets_constituents; jets = nothing, dxy = nothing, dz = nothing, phi0 = nothing, bz = nothing, kwargs...) -> get_btagSip3dVal(jets, + dxy, + dz, + phi0, + bz), + "btagSip3dSig" => (jets_constituents; btagSip3dVal = nothing, dxydxy = nothing, dzdz = nothing, kwargs...) -> get_btagSip3dSig(btagSip3dVal, + dxydxy, + dzdz), + "btagJetDistVal" => (jets_constituents; jets = nothing, dxy = nothing, dz = nothing, phi0 = nothing, bz = nothing, kwargs...) -> get_btagJetDistVal(jets, + jets_constituents, + dxy, + dz, + phi0, + bz), + "btagJetDistSig" => (jets_constituents; btagJetDistVal = nothing, dxydxy = nothing, dzdz = nothing, kwargs...) -> get_btagJetDistSig(btagJetDistVal, + dxydxy, + dzdz), + + # Special features + "mtof" => (jets_constituents; track_L = nothing, trackdata = nothing, trackerhits = nothing, gammadata = nothing, nhdata = nothing, calohits = nothing, v_in = nothing, kwargs...) -> get_mtof(jets_constituents, + track_L, + trackdata, + trackerhits, + gammadata, + nhdata, + calohits, + v_in), + "dndx" => (jets_constituents; dNdx = nothing, trackdata = nothing, jets_constituents_isChargedHad = nothing, kwargs...) -> get_dndx(jets_constituents, + dNdx, + trackdata, + jets_constituents_isChargedHad), + + # Mask + "mask" => (jets_constituents; kwargs...) -> [fill(1.0f0, + length(constituents)) + for constituents in jets_constituents]) + + return get(feature_map, feature_name, nothing) +end + +""" + get_feature_dependencies(feature_name::String) -> Dict{Symbol, String} + +Get the dependencies for features that require other computed features. + +# Arguments +- `feature_name`: Name of the feature + +# Returns +Dictionary mapping parameter names to feature names +""" +function get_feature_dependencies(feature_name::String) + deps = Dict{Symbol, String}() + + if feature_name == "btagSip2dSig" + deps[:btagSip2dVal] = "btagSip2dVal" + deps[:dxydxy] = "dxydxy" + elseif feature_name == "btagSip3dSig" + deps[:btagSip3dVal] = "btagSip3dVal" + deps[:dxydxy] = "dxydxy" + deps[:dzdz] = "dzdz" + elseif feature_name == "btagJetDistSig" + deps[:btagJetDistVal] = "btagJetDistVal" + deps[:dxydxy] = "dxydxy" + deps[:dzdz] = "dzdz" + elseif feature_name == "btagSip2dVal" + deps[:dxy] = "dxy" + deps[:phi0] = "phi0" + elseif feature_name == "btagSip3dVal" + deps[:dxy] = "dxy" + deps[:dz] = "dz" + deps[:phi0] = "phi0" + elseif feature_name == "btagJetDistVal" + deps[:dxy] = "dxy" + deps[:dz] = "dz" + deps[:phi0] = "phi0" + end + + return deps +end + +end # module diff --git a/ext/JetFlavourTagging/JetFlavourTagging.jl b/ext/JetFlavourTagging/JetFlavourTagging.jl new file mode 100644 index 00000000..3de4c0a1 --- /dev/null +++ b/ext/JetFlavourTagging/JetFlavourTagging.jl @@ -0,0 +1,177 @@ +module JetFlavourTagging + +using JetReconstruction +using EDM4hep +using JSON +using ONNXRunTime +using StructArrays: StructVector +using LorentzVectorHEP + +const JetConstituents = StructVector{ReconstructedParticle, <:Any} +const JetConstituentsData = Vector{Float32} + +# Building jet constituents from the EDM4hep event +include("JetConstituentBuilder.jl") + +""" + build_constituents_cluster(reco_particles::JetConstituents, + indices::Vector{Vector{Int}}) -> Vector{JetConstituents} + +Build the collection of constituents using cluster indices. + +# Arguments +- reco_particles: a vector of `JetConstituents` representing reconstructed particles. +- indices: a vector of vectors, where each inner vector contains indices of particles for a specific cluster. + +# Returns +A vector of JetConstituents, each containing the constituents for a specific cluster. +""" +function JetReconstruction.build_constituents_cluster(reco_particles::JetConstituents, + indices::Vector{Vector{Int64}}) + JetConstituentBuilder.build_constituents_cluster(reco_particles, indices) +end + +# TODO: Restore get_jets_constituents, get_constituents, build_constituents function. + +# Helper functions that talk with the ONNX model +include("JetFlavourHelper.jl") + +""" + extract_features(jets::Vector{EEJet}, jets_constituents::Vector{<:JetConstituents}, + tracks::AbstractVector{EDM4hep.TrackState}, bz::Float32, + track_L::AbstractArray{T} where T <: AbstractFloat, + json_config::Dict, + trackdata::AbstractVector{EDM4hep.Track}=AbstractVector{EDM4hep.Track}(), + trackerhits::AbstractVector{EDM4hep.TrackerHit}=AbstractVector{EDM4hep.TrackerHit}(), + gammadata::AbstractVector{EDM4hep.Cluster}=AbstractVector{EDM4hep.Cluster}(), + nhdata::AbstractVector{EDM4hep.Cluster}=AbstractVector{EDM4hep.Cluster}(), + calohits::AbstractVector{EDM4hep.CalorimeterHit}=AbstractVector{EDM4hep.CalorimeterHit}(), + dNdx::AbstractVector{EDM4hep.Quantity}=AbstractVector{EDM4hep.Quantity}()) -> Dict + +Extract features for jet flavour tagging based on JSON configuration. + +# Arguments +- `jets`: Vector of jets (EEJet) +- `jets_constituents`: Vector of jet constituents +- `tracks`: StructVector of track states +- `bz`: Magnetic field strength +- `track_L`: Array of track lengths +- `json_config`: JSON configuration dict specifying which features to extract (required) +- `trackdata`: Vector of track data (optional) +- `trackerhits`: Vector of tracker hits (optional) +- `gammadata`: Vector of gamma clusters (optional) +- `nhdata`: Vector of neutral hadron clusters (optional) +- `calohits`: Vector of calorimeter hits (optional) +- `dNdx`: Vector of dE/dx measurements (optional) +- `mc_vertices`: Vector of MC vertices for each reconstructed particle (optional) + +# Returns +Dictionary containing extracted features as specified in the JSON configuration. +""" +function JetReconstruction.extract_features(jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}, + tracks::AbstractVector{EDM4hep.TrackState}, + bz::Float32, + track_L::AbstractArray{T} where {T <: + AbstractFloat}, + json_config::Dict, + trackdata::AbstractVector{EDM4hep.Track} = AbstractVector{EDM4hep.Track}(), + trackerhits::AbstractVector{EDM4hep.TrackerHit} = AbstractVector{EDM4hep.TrackerHit}(), + gammadata::AbstractVector{EDM4hep.Cluster} = AbstractVector{EDM4hep.Cluster}(), + nhdata::AbstractVector{EDM4hep.Cluster} = AbstractVector{EDM4hep.Cluster}(), + calohits::AbstractVector{EDM4hep.CalorimeterHit} = AbstractVector{EDM4hep.CalorimeterHit}(), + dNdx::AbstractVector{EDM4hep.Quantity} = AbstractVector{EDM4hep.Quantity}(), + mc_vertices::Union{Nothing, + Vector{LorentzVector{Float32}}} = nothing) + return JetFlavourHelper.extract_features(jets, jets_constituents, tracks, bz, track_L, + json_config, trackdata, + trackerhits, gammadata, nhdata, calohits, dNdx, + mc_vertices) +end + +""" + setup_onnx_runtime(onnx_path::AbstractString, json_path::AbstractString) -> ONNXRunTime.InferenceSession + +Setup the ONNX model and preprocessing configuration for jet flavour tagging. + +# Arguments +- `onnx_path`: Path to the ONNX model file +- `json_path`: Path to the JSON configuration file + +# Returns +An ONNX inference session for the loaded model +""" +function JetReconstruction.setup_onnx_runtime(onnx_path::AbstractString, + json_path::AbstractString) + return JetFlavourHelper.setup_onnx_runtime(onnx_path, json_path) +end + +""" + prepare_input_tensor(jets_constituents::Vector{StructVector{EDM4hep.ReconstructedParticle}}, + jets::Vector{EEJet}, + config::Dict, + feature_data::Dict) -> Dict{String, Array} + +Prepare input tensors for the neural network from jet constituents. + +# Arguments +- `jets_constituents`: Vector of jet constituents (structured as a vector of StructVector of ReconstructedParticle) +- `jets`: Vector of jets (EEJet) +- `config`: JSON configuration for preprocessing +- `feature_data`: Dictionary containing all extracted features + +# Returns +Dictionary of input tensors +""" +function JetReconstruction.prepare_input_tensor(jets_constituents::Vector{<:JetConstituents}, + jets::Vector{EEJet}, + config::Dict, + feature_data::Dict) + return JetFlavourHelper.prepare_input_tensor(jets_constituents, jets, config, + feature_data) +end + +""" + get_weights(slot::Int, vars::Dict{String, Vector{Vector{Float32}}}, + jets::Vector{EEJet}, json_config::Dict, model::ONNXRunTime.InferenceSession) -> Vector{Vector{Float32}} + +Compute jet flavour probabilities for each jet. + +# Arguments +- `slot`: Threading slot +- `vars`: Dictionary containing all features for jet constituents +- `jets`: Vector of jets +- `json_config`: JSON configuration for preprocessing +- `model`: ONNX inference session + +# Returns +Vector of flavor probabilities for each jet +""" +function JetReconstruction.get_weights(slot::Int, + vars::Dict{String, + Dict{String, Vector{Vector{Float32}}}}, + jets::Vector{EEJet}, + jets_constituents::Vector{<:JetConstituents}, + json_config::Dict, + model::ONNXRunTime.InferenceSession) + return JetFlavourHelper.get_weights(slot, vars, jets, jets_constituents, json_config, + model) +end + +""" + get_weight(jet_weights::Vector{Vector{Float32}}, weight_idx::Int) -> Vector{Float32} + +Extract a specific weight/score from the jet weights. + +# Arguments +- `jet_weights`: Vector of weight vectors for each jet +- `weight_idx`: Index of the weight to extract + +# Returns +Vector of the specified weight for each jet +""" +function JetReconstruction.get_weight(jet_weights::Vector{Vector{Float32}}, weight_idx::Int) + return JetFlavourHelper.get_weight(jet_weights, weight_idx) +end + +end # module JetFlavourTagging diff --git a/ext/JetFlavourTagging/JetPhysicalConstants.jl b/ext/JetFlavourTagging/JetPhysicalConstants.jl new file mode 100644 index 00000000..cfa726aa --- /dev/null +++ b/ext/JetFlavourTagging/JetPhysicalConstants.jl @@ -0,0 +1,56 @@ +""" +Physical constants and special values used in the JetFlavourTagging extension. + +This module contains all physical constants, particle properties, tolerance values, +and special markers used throughout the jet flavour tagging algorithms. +""" +module JetPhysicalConstants + +using PhysicalConstants.CODATA2018: c_0, m_e, ħ, e, k_B, N_A, R + +# Note: PhysicalConstants.jl provides values with units and uncertainties +# To extract just the numerical value, use .val +# Example: c_0.val gives the speed of light value without units + +# Physical Constants from PhysicalConstants.jl +# Extract the numerical value and convert to Float32 +const C_LIGHT = Float32(c_0.val) # Speed of light in m/s from CODATA2018 +const C_LIGHT_INV = 1.0f0 / C_LIGHT # Inverse speed of light + +# Particle Masses (in GeV/c²) +# Electron mass from CODATA2018 (convert from kg to GeV/c²) +# 1 GeV/c² = 1.78266192e-27 kg +const ELECTRON_MASS = Float32(m_e.val / 1.78266192e-27) # Electron mass in GeV/c² +# Note: Muon and pion masses are not in CODATA2018, using PDG values +const MUON_MASS = 0.105658f0 # Muon mass in GeV/c² (PDG value) +const PION_MASS = 0.13957039f0 # Charged pion mass (π±) in GeV/c² (PDG value) + +# Mass Comparison Tolerances +const ELECTRON_TOLERANCE = 1.0f-5 # Tolerance for electron mass comparison +const MUON_TOLERANCE = 1.0f-3 # Tolerance for muon mass comparison +const PION_TOLERANCE = 1.0f-3 # Tolerance for pion mass comparison + +# PDG Particle ID Codes +const PDG_PHOTON = 22 # Photon (γ) +const PDG_K_LONG = 130 # K⁰_L (K-long neutral hadron) + +# Special/Undefined Values +const UNDEF_VAL = -9.0f0 # Sentinel value for missing/invalid data +const INVALID_TOF_MASS = 9.0f0 # Invalid mass from time-of-flight +const INVALID_MASS = -1.0f0 # Invalid mass calculation marker + +# Unit Conversion Factors +const MM_TO_M = 0.001f0 # Millimeter to meter conversion +const NS_TO_S = 1.0f-9 # Nanosecond to second conversion +const PS_TO_S = 1.0f-12 # Picosecond to second conversion +const FS_TO_S = 1.0f-15 # Femtosecond to second conversion + +# Export all constants +export C_LIGHT, C_LIGHT_INV +export ELECTRON_MASS, MUON_MASS, PION_MASS +export ELECTRON_TOLERANCE, MUON_TOLERANCE, PION_TOLERANCE +export PDG_PHOTON, PDG_K_LONG +export UNDEF_VAL, INVALID_TOF_MASS, INVALID_MASS +export MM_TO_M, NS_TO_S, PS_TO_S, FS_TO_S + +end # module JetPhysicalConstants diff --git a/src/JetReconstruction.jl b/src/JetReconstruction.jl index 5f9461d1..52901cec 100644 --- a/src/JetReconstruction.jl +++ b/src/JetReconstruction.jl @@ -93,6 +93,18 @@ function jetsplot() end function animatereco() end export jetsplot, animatereco +# Jet flavour tagging as an extension: Constituents +function build_constituents_cluster() end +export build_constituents_cluster + +# Jet flavour tagging as an extension: ONNX RunTime +function extract_features() end +function setup_onnx_runtime() end +function prepare_input_tensor() end +function get_weights() end +function get_weight() end +export extract_features, setup_onnx_runtime, prepare_input_tensor, get_weights, get_weight + # JSON results include("JSONresults.jl") export FinalJet, FinalJets diff --git a/test/data/jet-flavour-utils-event15.json b/test/data/jet-flavour-utils-event15.json new file mode 100644 index 00000000..2b0e4523 --- /dev/null +++ b/test/data/jet-flavour-utils-event15.json @@ -0,0 +1,3993 @@ +{ + "get_btagJetDistSig": [ + [ + -67.947235, + 86.78783, + -9.0 + ], + [ + -9.0, + -9.0, + 20.286007, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 33.79303, + -9.0, + -9.0, + -73.43549, + 17.754808, + -9.0, + -9.0, + -9.0, + 10.082263, + -9.0, + 27.870596, + -9.0, + 18.321705, + -9.0, + 7.818067, + -9.0, + -5.3476977, + -9.0, + -9.0, + -9.0, + -9.0, + 19.269228, + 13.5011215, + 9.805508, + -9.0, + -9.0, + 16.275623, + 32.900364, + -9.0, + -222.94427, + -22.069105, + -9.0, + -20.039652, + -9.0, + -9.0, + 161.43437, + -39.78606, + -9.0, + -72.438614, + -106.28761, + -9.0, + -9.0, + -9.0, + -87.568756, + -9.0, + -80.52304, + -87.215324, + -25.640558, + -9.0, + -49.46312, + -9.0, + -80.44129, + -9.0, + -9.0 + ] + ], + "get_dphidxy": [ + [ + -1.4331519e-7, + -3.790623e-8, + -9.0 + ], + [ + -9.0, + -9.0, + -4.2596348e-5, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -0.00022184216, + -9.0, + -9.0, + -2.4982919e-5, + -1.1713913e-5, + -9.0, + -9.0, + -9.0, + -0.00010989412, + -9.0, + -8.637991e-6, + -9.0, + -2.970423e-6, + -9.0, + -4.2645508e-5, + -9.0, + -6.020899e-6, + -9.0, + -9.0, + -9.0, + -9.0, + -7.210957e-5, + -0.00011014773, + -2.5094352e-5, + -9.0, + -9.0, + -1.445844e-5, + -1.4781508e-5, + -9.0, + -6.8675786e-6, + -3.823138e-5, + -9.0, + -1.4150387e-5, + -9.0, + -9.0, + -5.530079e-6, + -1.1464187e-5, + -9.0, + -1.52565e-6, + -3.306231e-7, + -9.0, + -9.0, + -9.0, + -7.545557e-7, + -9.0, + -1.713476e-6, + -1.4087543e-6, + -6.771888e-5, + -9.0, + -1.4865961e-5, + -9.0, + -2.1571765e-5, + -9.0, + -9.0 + ] + ], + "get_is_neutral_had": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + ], + "get_cctgtheta": [ + [ + -6.449425e-13, + -1.2061634e-13, + -9.0 + ], + [ + -9.0, + -9.0, + 1.1841313e-10, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 1.0515157e-9, + -9.0, + -9.0, + 2.5744823e-11, + 1.3882086e-10, + -9.0, + -9.0, + -9.0, + 1.6032972e-11, + -9.0, + 5.0662374e-11, + -9.0, + 4.1482436e-11, + -9.0, + 6.557174e-9, + -9.0, + -1.5103817e-12, + -9.0, + -9.0, + -9.0, + -9.0, + -8.270623e-11, + -1.3924242e-10, + -1.0142376e-11, + -9.0, + -9.0, + -3.4599637e-12, + -4.1310397e-12, + -9.0, + -3.699432e-12, + -2.8240157e-11, + -9.0, + -3.3806753e-12, + -9.0, + -9.0, + -2.5002214e-12, + -4.217762e-12, + -9.0, + -1.2024328e-12, + -8.898259e-13, + -9.0, + -9.0, + -9.0, + -1.0994897e-12, + -9.0, + -1.2317936e-12, + -1.2215513e-12, + -7.714124e-11, + -9.0, + -3.5868556e-12, + -9.0, + -1.1031687e-11, + -9.0, + -9.0 + ] + ], + "get_thetarel_phirel_cluster_phi": [ + [ + -0.5208537, + 3.0646446, + 3.0244973 + ], + [ + 0.72659594, + 0.5165751, + 0.38651362, + 0.40995017, + 0.4025254, + 0.37651336, + 0.3608677, + 0.36050427, + 0.21223237, + -0.09424443, + 0.18831973, + 0.288004, + 0.29260573, + 0.3138442, + 0.29361147, + 0.29616594, + 0.60098714, + 0.7948116, + 0.4454165, + 0.45674273, + 0.34565127, + 0.38426408, + 0.27704844, + 0.13654359, + -0.049244475, + -0.17645209, + 0.4505026, + -0.24221085, + -0.056248356, + 0.9263052, + 0.6869299, + -0.021149734, + -0.019146517, + -0.07003536, + 0.24315244, + 0.12440243, + 0.1031176, + -0.49736837, + -0.14890999, + -0.35360157, + -0.2957282, + -0.30756757, + -0.25922143, + -0.38812438, + -0.2925794, + -0.28040218, + -0.5224518, + -0.52352357, + -0.48899177, + -0.38702706, + -0.4361887, + -0.46316266, + -0.48044658, + -0.5772573, + -0.554203, + -0.5133457, + -0.50817966, + -1.0648339, + -0.91453534, + -0.77644795, + -0.7406752, + -0.6503955 + ] + ], + "get_dlambdadz": [ + [ + -7.0471465e-8, + -3.0397544e-8, + -9.0 + ], + [ + -9.0, + -9.0, + -5.389441e-6, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -0.0007103638, + -9.0, + -9.0, + -0.00012935487, + -2.6112128e-5, + -9.0, + -9.0, + -9.0, + -8.0461305e-6, + -9.0, + -1.0155118e-5, + -9.0, + -2.3524648e-5, + -9.0, + -0.00073939556, + -9.0, + -6.6659354e-8, + -9.0, + -9.0, + -9.0, + -9.0, + -3.4784377e-6, + -7.044444e-6, + -9.30439e-8, + -9.0, + -9.0, + -6.9575755e-8, + -7.140907e-8, + -9.0, + -1.3527776e-7, + -4.011601e-7, + -9.0, + -6.945535e-8, + -9.0, + -9.0, + -1.230852e-7, + -7.313976e-8, + -9.0, + -6.592762e-8, + -5.940042e-8, + -9.0, + -9.0, + -9.0, + -6.4185365e-8, + -9.0, + -6.613856e-8, + -6.599732e-8, + -3.027601e-6, + -9.0, + -6.9780754e-8, + -9.0, + -1.0467242e-7, + -9.0, + -9.0 + ] + ], + "get_detadeta": [ + [ + 3.9663663e-9, + 1.2429885e-9, + -9.0 + ], + [ + -9.0, + -9.0, + 5.025091e-7, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 3.8369876e-6, + -9.0, + -9.0, + 6.577294e-7, + 2.4365427e-6, + -9.0, + -9.0, + -9.0, + 7.522306e-7, + -9.0, + 9.363379e-7, + -9.0, + 2.1560136e-6, + -9.0, + 6.974637e-5, + -9.0, + 5.2411213e-9, + -9.0, + -9.0, + -9.0, + -9.0, + 3.2654916e-7, + 6.611711e-7, + 7.986393e-9, + -9.0, + -9.0, + 5.69052e-9, + 5.864933e-9, + -9.0, + 6.0152203e-9, + 3.653371e-8, + -9.0, + 5.6734226e-9, + -9.0, + -9.0, + 5.624221e-9, + 5.9808887e-9, + -9.0, + 4.8908837e-9, + 4.008018e-9, + -9.0, + -9.0, + -9.0, + 4.593706e-9, + -9.0, + 4.941001e-9, + 4.9028723e-9, + 2.843563e-7, + -9.0, + 5.7172747e-9, + -9.0, + 9.068475e-9, + -9.0, + -9.0 + ] + ], + "get_dxyctgtheta": [ + [ + 6.292053e-9, + 3.3675938e-9, + -9.0 + ], + [ + -9.0, + -9.0, + -6.7511223e-7, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -4.8680708e-5, + -9.0, + -9.0, + 4.158769e-8, + -8.708319e-8, + -9.0, + -9.0, + -9.0, + -3.6949543e-6, + -9.0, + -4.5584798e-8, + -9.0, + -1.0822876e-8, + -9.0, + -1.5915326e-6, + -9.0, + -4.698429e-9, + -9.0, + -9.0, + -9.0, + -9.0, + -2.5969312e-6, + -3.5149278e-6, + -6.833127e-7, + -9.0, + -9.0, + -9.057484e-8, + -1.215067e-7, + -9.0, + -6.6410564e-8, + -1.1264813e-6, + -9.0, + -8.6954394e-8, + -9.0, + -9.0, + -1.7921314e-8, + -1.2226397e-7, + -9.0, + 1.1045052e-9, + 1.968187e-9, + -9.0, + -9.0, + -9.0, + 1.2744004e-9, + -9.0, + 8.9180335e-10, + 7.8586343e-10, + -2.5856657e-6, + -9.0, + -9.6392625e-8, + -9.0, + -6.0772203e-7, + -9.0, + -9.0 + ] + ], + "get_pt": [ + [ + 19.585016, + 39.434834, + 5.6807175 + ], + [ + 0.030583913, + 0.044588927, + 0.9466278, + 0.21867794, + 0.7071525, + 0.32061985, + 0.5851007, + 1.846262, + 0.6913725, + 0.03469129, + 0.20054317, + 2.5339751, + 2.2439353, + 0.22262803, + 6.182514, + 1.7882537, + 0.61830574, + 0.26268315, + 2.4979172, + 1.0739576, + 5.542606, + 0.15381713, + 1.5448015, + 0.05124782, + 1.8302935, + 0.39949065, + 0.21064934, + 0.28746507, + 1.7741588, + 0.44911113, + 0.3913513, + 0.71315163, + 0.20376468, + 0.27531433, + 1.042198, + 1.0560232, + 1.5965091, + 1.3944374, + 0.8002477, + 2.1503873, + 1.0711482, + 1.963839, + 0.1535457, + 1.4954293, + 1.4522746, + 0.18812735, + 3.8077383, + 8.995785, + 1.4198884, + 3.3923564, + 2.0927284, + 5.7730374, + 1.150375, + 3.6240416, + 4.069206, + 0.45670453, + 0.33336213, + 1.0059193, + 0.05143729, + 0.8816397, + 0.6298998, + 0.2377641 + ] + ], + "get_phictgtheta": [ + [ + -1.2330542e-10, + -6.218373e-11, + -9.0 + ], + [ + -9.0, + -9.0, + 6.172346e-8, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 3.641078e-7, + -9.0, + -9.0, + 2.499853e-9, + 8.330528e-9, + -9.0, + -9.0, + -9.0, + 3.4265395e-7, + -9.0, + 4.3107646e-9, + -9.0, + 1.7081615e-9, + -9.0, + 1.7558588e-7, + -9.0, + 4.1850798e-10, + -9.0, + -9.0, + -9.0, + -9.0, + 2.420236e-7, + 3.2854544e-7, + 6.203719e-8, + -9.0, + -9.0, + 7.973657e-9, + 1.0720772e-8, + -9.0, + 2.9957257e-9, + 1.0282898e-7, + -9.0, + 7.654111e-9, + -9.0, + -9.0, + 7.910896e-10, + 1.0800605e-8, + -9.0, + -6.9361954e-11, + -6.1310214e-11, + -9.0, + -9.0, + -9.0, + -6.1018225e-11, + -9.0, + -5.7562583e-11, + -4.6775226e-11, + 2.4079992e-7, + -9.0, + 8.4872465e-9, + -9.0, + 5.4962324e-8, + -9.0, + -9.0 + ] + ], + "get_Sip3dVal_clusterV": [ + [ + 1.2500484, + 1.2302116, + -9.0 + ], + [ + -9.0, + -9.0, + -1.221479, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -31.435162, + -9.0, + -9.0, + 49.14793, + -1.2351213, + -9.0, + -9.0, + -9.0, + -1.243018, + -9.0, + -1.2365547, + -9.0, + -1.2731507, + -9.0, + -3.0375164, + -9.0, + -1.2371324, + -9.0, + -9.0, + -9.0, + -9.0, + -1.2283102, + -1.2468963, + -1.3215363, + -9.0, + -9.0, + -1.2372044, + -1.2729056, + -9.0, + -3.7006345, + -1.2084202, + -9.0, + -1.1800656, + -9.0, + -9.0, + 1.947983, + -1.4918803, + -9.0, + -1.2342032, + -1.218618, + -9.0, + -9.0, + -9.0, + -1.2404118, + -9.0, + -1.2354509, + -1.2281592, + -1.2978495, + -9.0, + -1.1241578, + -9.0, + -1.6519567, + -9.0, + -9.0 + ] + ], + "get_is_gamma": [ + [ + 0.0, + 0.0, + 1.0 + ], + [ + 1.0, + 1.0, + 0.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 1.0 + ] + ], + "get_phirel_cluster": [ + [ + -0.5208537, + 3.0646446, + 3.0244973 + ], + [ + 0.72659594, + 0.5165751, + 0.38651362, + 0.40995017, + 0.4025254, + 0.37651336, + 0.3608677, + 0.36050427, + 0.21223237, + -0.09424443, + 0.18831973, + 0.288004, + 0.29260573, + 0.3138442, + 0.29361147, + 0.29616594, + 0.60098714, + 0.7948116, + 0.4454165, + 0.45674273, + 0.34565127, + 0.38426408, + 0.27704844, + 0.13654359, + -0.049244475, + -0.17645209, + 0.4505026, + -0.24221085, + -0.056248356, + 0.9263052, + 0.6869299, + -0.021149734, + -0.019146517, + -0.07003536, + 0.24315244, + 0.12440243, + 0.1031176, + -0.49736837, + -0.14890999, + -0.35360157, + -0.2957282, + -0.30756757, + -0.25922143, + -0.38812438, + -0.2925794, + -0.28040218, + -0.5224518, + -0.52352357, + -0.48899177, + -0.38702706, + -0.4361887, + -0.46316266, + -0.48044658, + -0.5772573, + -0.554203, + -0.5133457, + -0.50817966, + -1.0648339, + -0.91453534, + -0.77644795, + -0.7406752, + -0.6503955 + ] + ], + "get_e": [ + [ + 20.46865, + 41.66906, + 5.97301 + ], + [ + 0.038685087, + 0.07068448, + 2.3770719, + 0.49115527, + 2.2526124, + 1.1003908, + 1.757357, + 5.843931, + 2.5155842, + 0.10352973, + 0.6454202, + 8.573691, + 7.3090167, + 0.68363595, + 21.805687, + 7.0275903, + 1.7399287, + 0.50034, + 7.1612535, + 3.5086353, + 21.835117, + 0.674085, + 6.9506555, + 0.33211777, + 1.8667521, + 0.41225448, + 0.21637824, + 0.28761873, + 1.7807349, + 0.5055865, + 0.43000504, + 0.8081286, + 0.21984208, + 0.30141035, + 1.1128494, + 1.1604855, + 1.7245221, + 1.8815017, + 1.3515927, + 2.365801, + 1.1421924, + 2.1421323, + 0.17058831, + 1.793193, + 1.7349396, + 0.2294921, + 3.973371, + 9.240602, + 1.4701939, + 3.5380654, + 2.2533035, + 6.112487, + 1.2084657, + 3.8613276, + 4.4815288, + 0.5223787, + 0.3654142, + 1.0771525, + 0.055122558, + 1.157116, + 0.7516057, + 0.29473728 + ] + ], + "get_Sip3dSig": [ + [ + 339.9819, + 454.42355, + -9.0 + ], + [ + -9.0, + -9.0, + -51.253014, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -77.00076, + -9.0, + -9.0, + 283.22318, + -58.696545, + -9.0, + -9.0, + -9.0, + -33.636562, + -9.0, + -80.866066, + -9.0, + -72.00806, + -9.0, + -33.132137, + -9.0, + -132.91104, + -9.0, + -9.0, + -9.0, + -9.0, + -42.123447, + -34.58784, + -76.005974, + -9.0, + -9.0, + -91.24579, + -92.93252, + -9.0, + -288.07434, + -56.679485, + -9.0, + -87.88222, + -9.0, + -9.0, + 167.02583, + -121.84979, + -9.0, + -222.53693, + -339.40958, + -9.0, + -9.0, + -9.0, + -280.50995, + -9.0, + -213.76668, + -226.90367, + -45.933792, + -9.0, + -81.871735, + -9.0, + -101.52174, + -9.0, + -9.0 + ] + ], + "get_eta": [ + [ + 0.29927474, + -0.3350398, + -0.31943128 + ], + [ + 0.71267295, + 1.0350659, + 1.5697458, + 1.4485981, + 1.8261456, + 1.9043777, + 1.7639877, + 1.8194458, + 1.9636836, + 1.7571852, + 1.8369585, + 1.8893272, + 1.8493899, + 1.7874366, + 1.93259, + 2.0425565, + 1.6911292, + 1.2601175, + 1.7119268, + 1.8420545, + 2.0476682, + 2.157461, + 2.1842961, + 2.5559576, + 0.18448064, + 0.2521173, + -0.23269734, + -0.032695387, + 0.086073026, + 0.4022625, + 0.27938148, + 0.47743002, + 0.39467838, + 0.4320314, + 0.34287453, + 0.42331508, + 0.39782882, + 0.80910134, + 1.1085049, + 0.44394884, + 0.339929, + 0.4229579, + 0.4669032, + 0.61550313, + 0.60825956, + 0.6515521, + 0.29172704, + 0.22644348, + 0.2654127, + 0.2545968, + 0.38927722, + 0.3408125, + 0.31647375, + 0.35802734, + 0.4452957, + 0.44838184, + 0.43507576, + 0.34974533, + 0.3763145, + 0.7598866, + 0.61203706, + 0.679145 + ] + ], + "get_dndx": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -1.0, + 0.0, + 0.0, + -1.0, + -1.0, + 0.0, + 0.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + -1.0, + -1.0, + -1.0, + 0.0, + 0.0, + -1.0, + -1.0, + 0.0, + -1.0, + -1.0, + 0.0, + -1.0, + 0.0, + 0.0, + -1.0, + -1.0, + 0.0, + -1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -1.0, + -1.0, + -1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + 0.0 + ] + ], + "get_mtof": [ + [ + 0.00051099894, + 0.105658, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12318668, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.4024081, + 0.0, + 0.0, + 4.6633806, + 0.42040208, + 0.0, + 9.0, + 9.0, + 0.1424469, + 0.0, + 0.13957039, + 9.0, + 0.13957039, + 0.0, + 0.13957039, + 0.0, + 0.23157507, + 0.0, + 0.0, + 0.0, + 0.0, + 0.13760842, + 0.14033855, + 0.12590039, + 0.0, + 0.0, + 0.13731903, + 0.078510016, + 0.0, + 0.2713576, + 0.07135989, + 0.0, + 0.13521484, + 0.0, + 0.0, + 0.25415704, + 0.05268342, + 0.0, + 0.44820195, + 0.13957039, + 0.0, + 9.0, + 0.0, + 0.105658, + 0.0, + 0.21945497, + 0.44456804, + 0.1405827, + 0.0, + 0.12969306, + 0.0, + 0.15140556, + 0.0, + 0.0 + ] + ], + "get_btagSip3dSig": [ + [ + 339.9819, + 454.42355, + -9.0 + ], + [ + -9.0, + -9.0, + -51.253014, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -77.00076, + -9.0, + -9.0, + 283.22318, + -58.696545, + -9.0, + -9.0, + -9.0, + -33.636562, + -9.0, + -80.866066, + -9.0, + -72.00806, + -9.0, + -33.132137, + -9.0, + -132.91104, + -9.0, + -9.0, + -9.0, + -9.0, + -42.123447, + -34.58784, + -76.005974, + -9.0, + -9.0, + -91.24579, + -92.93252, + -9.0, + -288.07434, + -56.679485, + -9.0, + -87.88222, + -9.0, + -9.0, + 167.02583, + -121.84979, + -9.0, + -222.53693, + -339.40958, + -9.0, + -9.0, + -9.0, + -280.50995, + -9.0, + -213.76668, + -226.90367, + -45.933792, + -9.0, + -81.871735, + -9.0, + -101.52174, + -9.0, + -9.0 + ] + ], + "get_Sip2dVal_clusterV": [ + [ + -0.0023730414, + 0.0070432117, + -9.0 + ], + [ + -9.0, + -9.0, + -0.029755257, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -13.5313015, + -9.0, + -9.0, + -4.109446, + -0.028878165, + -9.0, + -9.0, + -9.0, + -0.061065387, + -9.0, + -0.013139597, + -9.0, + -0.023180787, + -9.0, + -0.25246602, + -9.0, + 0.012407542, + -9.0, + -9.0, + -9.0, + -9.0, + -0.030000899, + 0.0011837726, + -0.19509333, + -9.0, + -9.0, + 0.0040815696, + 0.3277208, + -9.0, + 1.7234792, + -0.0054963706, + -9.0, + 0.022943396, + -9.0, + -9.0, + -1.7221894, + 0.070499085, + -9.0, + 0.0029449174, + 0.009541756, + -9.0, + -9.0, + -9.0, + 0.0062310095, + -9.0, + 0.012660839, + 0.0027262159, + 0.28730667, + -9.0, + 0.15754971, + -9.0, + 0.7631629, + -9.0, + -9.0 + ] + ], + "get_erel_log_cluster": [ + [ + -0.52212626, + -0.21340178, + -1.0570222 + ], + [ + -3.641111, + -3.3793304, + -1.8526123, + -2.5374358, + -1.8759681, + -2.1871076, + -1.9837946, + -1.4619495, + -1.8280157, + -3.2135894, + -2.418812, + -1.2954867, + -1.3647956, + -2.3938296, + -0.89008474, + -1.3818481, + -1.988123, + -2.5293894, + -1.3736655, + -1.6835163, + -0.889499, + -2.3999398, + -1.3866287, + -2.7073624, + -1.9575678, + -2.6134892, + -2.893441, + -2.7698374, + -1.9780552, + -2.524859, + -2.595181, + -2.3211741, + -2.8865438, + -2.7494965, + -2.182218, + -2.1640148, + -1.9919858, + -1.95415, + -2.0978086, + -1.8546762, + -2.1709154, + -1.8978082, + -2.9967053, + -1.9750274, + -1.9893701, + -2.8678868, + -1.6294954, + -1.2629542, + -2.06128, + -1.6798886, + -1.8758348, + -1.4424366, + -2.1464202, + -1.6419178, + -1.5772283, + -2.510669, + -2.665869, + -2.1963773, + -3.4873252, + -2.1652775, + -2.3526645, + -2.7592194 + ] + ], + "get_btagSip3dVal": [ + [ + 1.2500484, + 1.2302116, + -9.0 + ], + [ + -9.0, + -9.0, + -1.221479, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -31.435162, + -9.0, + -9.0, + 49.14793, + -1.2351213, + -9.0, + -9.0, + -9.0, + -1.243018, + -9.0, + -1.2365547, + -9.0, + -1.2731507, + -9.0, + -3.0375164, + -9.0, + -1.2371324, + -9.0, + -9.0, + -9.0, + -9.0, + -1.2283102, + -1.2468963, + -1.3215363, + -9.0, + -9.0, + -1.2372044, + -1.2729056, + -9.0, + -3.7006345, + -1.2084202, + -9.0, + -1.1800656, + -9.0, + -9.0, + 1.947983, + -1.4918803, + -9.0, + -1.2342032, + -1.218618, + -9.0, + -9.0, + -9.0, + -1.2404118, + -9.0, + -1.2354509, + -1.2281592, + -1.2978495, + -9.0, + -1.1241578, + -9.0, + -1.6519567, + -9.0, + -9.0 + ] + ], + "get_thetarel_cluster": [ + [ + 2.966602, + 2.6673548, + 2.6804037 + ], + [ + 1.2688109, + 1.104194, + 0.8537709, + 0.90194684, + 0.7381785, + 0.7217353, + 0.7827701, + 0.7615588, + 0.768493, + 0.855986, + 0.8140161, + 0.7669116, + 0.77883875, + 0.79305065, + 0.7503496, + 0.7141931, + 0.6717117, + 0.8081033, + 0.76337284, + 0.6997838, + 0.6875741, + 0.62219524, + 0.681343, + 0.37576556, + 1.9060644, + 1.8329093, + 2.2823558, + 2.1103423, + 2.0034032, + 1.4725348, + 1.695954, + 1.6296993, + 1.7051743, + 1.669573, + 1.7384486, + 1.6747093, + 1.6995277, + 1.2766685, + 1.148414, + 1.6260409, + 1.7339329, + 1.6537427, + 1.6207951, + 1.4655788, + 1.4909564, + 1.4565848, + 1.7335267, + 1.798572, + 1.7681736, + 1.8011085, + 1.6599228, + 1.7002355, + 1.7198068, + 1.6526482, + 1.5753503, + 1.5842625, + 1.5982518, + 1.4586637, + 1.5046724, + 1.2052091, + 1.3525188, + 1.3295034 + ] + ], + "get_dxyc": [ + [ + -9.806078e-12, + -1.1884923e-11, + -9.0 + ], + [ + -9.0, + -9.0, + 2.2361994e-9, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -3.9790885e-7, + -9.0, + -9.0, + -4.8361343e-8, + -1.04284394e-10, + -9.0, + -9.0, + -9.0, + 1.2184227e-8, + -9.0, + -4.238671e-11, + -9.0, + -1.6290509e-10, + -9.0, + -3.5651507e-10, + -9.0, + 2.9299937e-10, + -9.0, + -9.0, + -9.0, + -9.0, + 1.2219179e-8, + 2.1435525e-8, + 2.6029148e-9, + -9.0, + -9.0, + 1.2349527e-9, + 1.2548483e-9, + -9.0, + 3.0484518e-10, + 4.2269774e-9, + -9.0, + 1.2003353e-9, + -9.0, + -9.0, + 2.1060016e-10, + 8.364162e-10, + -9.0, + 7.593243e-12, + -2.0008912e-12, + -9.0, + -9.0, + -9.0, + -2.4154103e-13, + -9.0, + 1.0620326e-11, + 4.863143e-12, + 1.10811476e-8, + -9.0, + 1.2803522e-9, + -9.0, + 1.9449595e-9, + -9.0, + -9.0 + ] + ], + "get_theta": [ + [ + 1.2758915, + 1.8997383, + 1.8849299 + ], + [ + 0.91176635, + 0.6826049, + 0.41033983, + 0.46143308, + 0.31932476, + 0.2956577, + 0.3394235, + 0.32143468, + 0.27885988, + 0.34169573, + 0.3159475, + 0.3000747, + 0.31210768, + 0.3317022, + 0.2875477, + 0.25795338, + 0.36453143, + 0.5527261, + 0.35718852, + 0.31436795, + 0.25665262, + 0.23021467, + 0.22417025, + 0.15492517, + 1.3873533, + 1.3213083, + 1.8014216, + 1.6034858, + 1.4848294, + 1.1789633, + 1.2949799, + 1.1105342, + 1.1859827, + 1.1516099, + 1.234449, + 1.1595855, + 1.1830645, + 0.83780706, + 0.6375889, + 1.1407508, + 1.2372309, + 1.1599128, + 1.1199875, + 0.99083775, + 0.9969089, + 0.96098363, + 1.2831212, + 1.3462636, + 1.3084459, + 1.3189063, + 1.1909939, + 1.2363962, + 1.2594769, + 1.2201816, + 1.1395268, + 1.1367251, + 1.1488307, + 1.2279707, + 1.2030616, + 0.87498266, + 0.99373984, + 0.9385439 + ] + ], + "get_Sip2dSig": [ + [ + -0.8752023, + 3.6409461, + -9.0 + ], + [ + -9.0, + -9.0, + -1.3618871, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -73.746315, + -9.0, + -9.0, + -64.52851, + -2.42035, + -9.0, + -9.0, + -9.0, + -1.7638716, + -9.0, + -1.2636894, + -9.0, + -3.5361004, + -9.0, + -11.560585, + -9.0, + 1.403294, + -9.0, + -9.0, + -9.0, + -9.0, + -1.0657609, + 0.03416786, + -11.497127, + -9.0, + -9.0, + 0.31063297, + 24.691864, + -9.0, + 138.90831, + -0.2649694, + -9.0, + 1.7638229, + -9.0, + -9.0, + -153.34438, + 5.979673, + -9.0, + 0.59213567, + 3.2991476, + -9.0, + -9.0, + -9.0, + 1.6434427, + -9.0, + 2.4296033, + 0.56557465, + 10.523689, + -9.0, + 11.835416, + -9.0, + 48.2779, + -9.0, + -9.0 + ] + ], + "get_Bz": [ + [ + 2.0e-24, + 2.0e-24, + -9.0 + ], + [ + -9.0, + -9.0, + 2.0e-24, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 2.0e-24, + -9.0, + -9.0, + 1.9999998e-24, + 1.9999998e-24, + -9.0, + -9.0, + -9.0, + 2.0e-24, + -9.0, + 2.0e-24, + -9.0, + 1.9999998e-24, + -9.0, + 2.0e-24, + -9.0, + 2.0e-24, + -9.0, + -9.0, + -9.0, + -9.0, + 2.0e-24, + 2.0000002e-24, + 2.0e-24, + -9.0, + -9.0, + 2.0e-24, + 2.0000002e-24, + -9.0, + 2.0000002e-24, + 2.0e-24, + -9.0, + 2.0000002e-24, + -9.0, + -9.0, + 2.0e-24, + 2.0e-24, + -9.0, + 2.0e-24, + 2.0000002e-24, + -9.0, + -9.0, + -9.0, + 2.0e-24, + -9.0, + 1.9999998e-24, + 2.0000002e-24, + 2.0e-24, + -9.0, + 2.0e-24, + -9.0, + 2.0e-24, + -9.0, + -9.0 + ] + ], + "get_phi": [ + [ + -3.042015, + -2.9142451, + -2.8966546 + ], + [ + 1.5166088, + 1.3598826, + 1.3782432, + 1.3641411, + 1.5823449, + 1.5713365, + 1.4304261, + 1.4650902, + 1.1463447, + 0.37179676, + 1.0389442, + 1.3154137, + 1.3061527, + 1.3270807, + 1.3551353, + 1.4326339, + 1.9952759, + 1.9709144, + 1.6051311, + 1.7500311, + 1.5960102, + 1.8636805, + 1.472166, + -2.8857496, + 0.5382197, + 0.40965012, + 0.931081, + 0.3781444, + 0.53426325, + 1.6223304, + 1.2982714, + 0.56194866, + 0.5650433, + 0.5092027, + 0.83972263, + 0.72056067, + 0.6960077, + -0.07613915, + 0.35617736, + 0.19530092, + 0.27626678, + 0.25017342, + 0.29710987, + 0.11882726, + 0.23609786, + 0.24349344, + 0.046236947, + 0.062421393, + 0.08835016, + 0.19631808, + 0.11527469, + 0.09733063, + 0.08470749, + -0.03215172, + -0.032359812, + 0.013598185, + 0.023169989, + -0.59008586, + -0.42515403, + -0.43553087, + -0.3188946, + -0.231122 + ] + ], + "get_phic": [ + [ + 1.6819156e-13, + 1.9235148e-13, + -9.0 + ], + [ + -9.0, + -9.0, + -2.013012e-10, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 3.696659e-9, + -9.0, + -9.0, + 3.9276318e-10, + 1.1435488e-11, + -9.0, + -9.0, + -9.0, + -1.1212453e-9, + -9.0, + 3.413153e-12, + -9.0, + 7.9943205e-12, + -9.0, + 6.6458526e-11, + -9.0, + -2.5687372e-11, + -9.0, + -9.0, + -9.0, + -9.0, + -1.1460632e-9, + -2.0148294e-9, + -2.375275e-10, + -9.0, + -9.0, + -1.1009648e-10, + -1.1200412e-10, + -9.0, + -1.4127822e-11, + -3.8827352e-10, + -9.0, + -1.0699264e-10, + -9.0, + -9.0, + -9.669027e-12, + -7.467223e-11, + -9.0, + -6.7238326e-13, + 6.675754e-14, + -9.0, + -9.0, + -9.0, + -4.0132147e-14, + -9.0, + -9.197352e-13, + -4.5189755e-13, + -1.038851e-9, + -9.0, + -1.1417029e-10, + -9.0, + -1.7658444e-10, + -9.0, + -9.0 + ] + ], + "get_dxydxy": [ + [ + 7.351801e-6, + 3.7420793e-6, + -9.0 + ], + [ + -9.0, + -9.0, + 0.00047735905, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 0.033666536, + -9.0, + -9.0, + 0.0040556756, + 0.0001423583, + -9.0, + -9.0, + -9.0, + 0.0011985505, + -9.0, + 0.000108114364, + -9.0, + 4.2974134e-5, + -9.0, + 0.00047692077, + -9.0, + 7.8176134e-5, + -9.0, + -9.0, + -9.0, + -9.0, + 0.00079240825, + 0.0012003309, + 0.00028794282, + -9.0, + -9.0, + 0.00017264711, + 0.00017615712, + -9.0, + 0.00015394147, + 0.00043028934, + -9.0, + 0.00016920199, + -9.0, + -9.0, + 0.00012613223, + 0.00013899914, + -9.0, + 2.4734536e-5, + 8.364752e-6, + -9.0, + -9.0, + -9.0, + 1.4374995e-5, + -9.0, + 2.7155276e-5, + 2.323489e-5, + 0.0007453416, + -9.0, + 0.00017720183, + -9.0, + 0.0002498836, + -9.0, + -9.0 + ] + ], + "get_ct": [ + [ + 0.30376217, + -0.34134316, + -9.0 + ], + [ + -9.0, + -9.0, + 2.298664, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 3.4925916, + -9.0, + -9.0, + 3.2318728, + 3.0993044, + -9.0, + -9.0, + -9.0, + 2.620647, + -9.0, + 2.679554, + -9.0, + 3.8103883, + -9.0, + 4.38592, + -9.0, + 0.18552884, + -9.0, + -9.0, + -9.0, + -9.0, + 0.41319934, + 0.28303018, + 0.4957753, + -9.0, + -9.0, + 0.34963235, + 0.43607152, + -9.0, + 0.90031534, + 1.3498868, + -9.0, + 0.34651348, + -9.0, + -9.0, + 0.6551093, + 0.6464668, + -9.0, + 0.29588258, + 0.22838365, + -9.0, + -9.0, + -9.0, + 0.34744865, + -9.0, + 0.36572534, + 0.46015847, + 0.46355784, + -9.0, + 0.35691932, + -9.0, + 0.83515716, + -9.0, + -9.0 + ] + ], + "get_JetDistVal_clusterV": [ + [ + -0.24982898, + 0.23495126, + -9.0 + ], + [ + -9.0, + -9.0, + 0.48346293, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 13.795829, + -9.0, + -9.0, + -12.743316, + 0.3736053, + -9.0, + -9.0, + -9.0, + 0.37258372, + -9.0, + 0.42618024, + -9.0, + 0.32394, + -9.0, + 0.71675146, + -9.0, + -0.049776226, + -9.0, + -9.0, + -9.0, + -9.0, + 0.56188637, + 0.48671722, + 0.170491, + -9.0, + -9.0, + 0.22068168, + 0.45063943, + -9.0, + -2.8639665, + -0.4705186, + -9.0, + -0.26908863, + -9.0, + -9.0, + 1.8827711, + -0.48712465, + -9.0, + -0.40174893, + -0.38161564, + -9.0, + -9.0, + -9.0, + -0.38722804, + -9.0, + -0.46537778, + -0.4720695, + -0.7244685, + -9.0, + -0.6791642, + -9.0, + -1.3089366, + -9.0, + -9.0 + ] + ], + "get_phi0": [ + [ + -3.042015, + -2.9142451, + -9.0 + ], + [ + -9.0, + -9.0, + 1.3782432, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -1.995248, + -9.0, + -9.0, + -1.8261789, + 1.3061527, + -9.0, + -9.0, + -9.0, + 1.9952759, + -9.0, + 1.6051311, + -9.0, + 1.5960102, + -9.0, + 1.472166, + -9.0, + 0.5382197, + -9.0, + -9.0, + -9.0, + -9.0, + 1.6223304, + 1.2982714, + 0.56194866, + -9.0, + -9.0, + 0.8397226, + 0.72056067, + -9.0, + -0.07613915, + 0.35617736, + -9.0, + 0.27626678, + -9.0, + -9.0, + 0.11882725, + 0.23609786, + -9.0, + 0.046236947, + 0.062421393, + -9.0, + -9.0, + -9.0, + 0.09733063, + -9.0, + -0.03215172, + -0.032359816, + 0.013598185, + -9.0, + -0.59008586, + -9.0, + -0.43553087, + -9.0, + -9.0 + ] + ], + "get_c": [ + [ + -1.5307236e19, + 7.6022246e18, + -9.0 + ], + [ + -9.0, + -9.0, + -3.1669517e20, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 4.3361927e20, + -9.0, + -9.0, + -1.1830916e20, + -1.336012e20, + -9.0, + -9.0, + -9.0, + -4.8486116e20, + -9.0, + 1.2001697e20, + -9.0, + -5.408872e19, + -9.0, + 1.9406537e20, + -9.0, + 1.6379474e20, + -9.0, + -9.0, + -9.0, + -9.0, + 6.67524e20, + -7.6604435e20, + 4.2037686e20, + -9.0, + -9.0, + 2.8765404e20, + -2.8388813e20, + -9.0, + 2.1499168e20, + 3.7462455e20, + -9.0, + 2.7987954e20, + -9.0, + -9.0, + -2.0047251e20, + 2.064296e20, + -9.0, + 7.873242e19, + -3.332588e19, + -9.0, + -9.0, + -9.0, + 5.1929763e19, + -9.0, + -8.2723235e19, + -7.367345e19, + -6.564254e20, + -9.0, + -2.980283e20, + -9.0, + -3.4003965e20, + -9.0, + -9.0 + ] + ], + "input_info": { + "n_jets": 2, + "event_id": 15, + "bz": 2.0, + "n_constituents_per_jet": [ + 3, + 62 + ], + "vertex": { + "t": 0.0, + "x": 0.0, + "z": 0.0, + "y": 0.0 + } + }, + "get_btagJetDistVal": [ + [ + -0.24982898, + 0.23495126, + -9.0 + ], + [ + -9.0, + -9.0, + 0.48346293, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 13.795829, + -9.0, + -9.0, + -12.743316, + 0.3736053, + -9.0, + -9.0, + -9.0, + 0.37258372, + -9.0, + 0.42618024, + -9.0, + 0.32394, + -9.0, + 0.71675146, + -9.0, + -0.049776226, + -9.0, + -9.0, + -9.0, + -9.0, + 0.56188637, + 0.48671722, + 0.170491, + -9.0, + -9.0, + 0.22068168, + 0.45063943, + -9.0, + -2.8639665, + -0.4705186, + -9.0, + -0.26908863, + -9.0, + -9.0, + 1.8827711, + -0.48712465, + -9.0, + -0.40174893, + -0.38161564, + -9.0, + -9.0, + -9.0, + -0.38722804, + -9.0, + -0.46537778, + -0.4720695, + -0.7244685, + -9.0, + -0.6791642, + -9.0, + -1.3089366, + -9.0, + -9.0 + ] + ], + "get_p": [ + [ + 20.46865, + 41.668926, + 5.97301 + ], + [ + 0.038685087, + 0.07068448, + 2.3729708, + 0.4911553, + 2.2526124, + 1.1003909, + 1.7573572, + 5.843931, + 2.5117095, + 0.10352973, + 0.6454202, + 8.572556, + 7.3076844, + 0.68363595, + 21.800007, + 7.009951, + 1.7343218, + 0.50034, + 7.1442223, + 3.4731693, + 21.834673, + 0.674085, + 6.949254, + 0.3321178, + 1.8615272, + 0.41225448, + 0.21637826, + 0.28761873, + 1.7807349, + 0.48594025, + 0.40672418, + 0.7959849, + 0.21984208, + 0.30141032, + 1.1040623, + 1.1520619, + 1.7245221, + 1.8763179, + 1.3443673, + 2.365801, + 1.1336329, + 2.1421325, + 0.17058831, + 1.7877531, + 1.7293165, + 0.22949208, + 3.970919, + 9.227408, + 1.4701939, + 3.5028973, + 2.2533035, + 6.111573, + 1.2084657, + 3.8588045, + 4.479355, + 0.5033882, + 0.3654142, + 1.068072, + 0.055122558, + 1.1486678, + 0.7516057, + 0.29473728 + ] + ], + "get_y": [ + [ + 0.29927474, + -0.33503863, + -0.31943128 + ], + [ + 0.71267295, + 1.0350659, + 1.559894, + 1.4485984, + 1.8261456, + 1.9043789, + 1.7639889, + 1.8194458, + 1.9444977, + 1.7571852, + 1.8369585, + 1.8878807, + 1.8475533, + 1.7874366, + 1.9294934, + 2.0065458, + 1.6679502, + 1.2601175, + 1.6940047, + 1.7500317, + 2.0473628, + 2.157461, + 2.1803334, + 2.5559614, + 0.18395254, + 0.2521173, + -0.23269743, + -0.032695387, + 0.086073026, + 0.38500237, + 0.2635235, + 0.4691451, + 0.39467838, + 0.43203136, + 0.33995277, + 0.41986668, + 0.39782882, + 0.8057711, + 1.0964984, + 0.44394884, + 0.33718327, + 0.42295802, + 0.4669032, + 0.6131304, + 0.60576797, + 0.6515521, + 0.29153663, + 0.22610903, + 0.2654127, + 0.25195706, + 0.38927722, + 0.3407575, + 0.31647375, + 0.35777283, + 0.44505003, + 0.42994905, + 0.43507576, + 0.34655395, + 0.3763145, + 0.7519823, + 0.61203706, + 0.679145 + ] + ], + "get_dzdz": [ + [ + 6.167121e-6, + 3.5868065e-6, + -9.0 + ], + [ + -9.0, + -9.0, + 9.062119e-5, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 0.13299727, + -9.0, + -9.0, + 0.026057245, + 0.00030042796, + -9.0, + -9.0, + -9.0, + 0.00016707527, + -9.0, + 0.00012571232, + -9.0, + 0.0002696319, + -9.0, + 0.007928092, + -9.0, + 8.462221e-6, + -9.0, + -9.0, + -9.0, + -9.0, + 5.7884205e-5, + 9.928106e-5, + 1.4374335e-5, + -9.0, + -9.0, + 1.1200187e-5, + 1.1453289e-5, + -9.0, + 1.1081046e-5, + 2.4263236e-5, + -9.0, + 1.11040545e-5, + -9.0, + -9.0, + 9.887867e-6, + 1.0906642e-5, + -9.0, + 6.024251e-6, + 4.5262614e-6, + -9.0, + -9.0, + -9.0, + 5.178986e-6, + -9.0, + 6.246609e-6, + 6.062323e-6, + 5.299152e-5, + -9.0, + 1.1330817e-5, + -9.0, + 1.4892712e-5, + -9.0, + -9.0 + ] + ], + "get_is_charged_had": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 1.0, + 1.0, + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 1.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 1.0, + 1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0 + ] + ], + "get_charge": [ + [ + -1.0, + 1.0, + 0.0 + ], + [ + 0.0, + 0.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + -1.0, + -1.0, + 0.0, + 0.0, + 0.0, + -1.0, + 0.0, + 1.0, + 0.0, + -1.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + -1.0, + 1.0, + 0.0, + 0.0, + 1.0, + -1.0, + 0.0, + 1.0, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + -1.0, + 1.0, + 0.0, + 1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + -1.0, + -1.0, + -1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + 0.0 + ] + ], + "get_dxy": [ + [ + -0.0023730414, + -0.0070432117, + -9.0 + ], + [ + -9.0, + -9.0, + 0.029755257, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -13.5313015, + -9.0, + -9.0, + -4.109446, + 0.028878165, + -9.0, + -9.0, + -9.0, + 0.061065387, + -9.0, + 0.013139597, + -9.0, + 0.023180787, + -9.0, + 0.25246602, + -9.0, + 0.012407542, + -9.0, + -9.0, + -9.0, + -9.0, + 0.030000899, + -0.0011837726, + -0.19509333, + -9.0, + -9.0, + -0.0040815696, + -0.3277208, + -9.0, + 1.7234792, + -0.0054963706, + -9.0, + 0.022943396, + -9.0, + -9.0, + -1.7221894, + 0.070499085, + -9.0, + 0.0029449174, + 0.009541756, + -9.0, + -9.0, + -9.0, + 0.0062310095, + -9.0, + 0.012660839, + 0.0027262159, + 0.28730667, + -9.0, + 0.15754971, + -9.0, + 0.7631629, + -9.0, + -9.0 + ] + ], + "get_phidz": [ + [ + 1.8795894e-9, + 1.0298999e-9, + -9.0 + ], + [ + -9.0, + -9.0, + -2.968296e-7, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -9.764265e-5, + -9.0, + -9.0, + 4.526342e-6, + -3.136453e-8, + -9.0, + -9.0, + -9.0, + -2.0056502e-6, + -9.0, + -1.3099254e-8, + -9.0, + -1.1173231e-8, + -9.0, + -1.2900277e-6, + -9.0, + -1.7138808e-9, + -9.0, + -9.0, + -9.0, + -9.0, + -2.22538e-6, + -3.0627493e-6, + -6.2787063e-7, + -9.0, + -9.0, + -7.2231025e-8, + -9.783879e-8, + -9.0, + -5.401044e-7, + -8.821527e-7, + -9.0, + -6.922028e-8, + -9.0, + -9.0, + 3.0399463e-7, + -9.7178024e-8, + -9.0, + 1.3670471e-9, + 8.823004e-10, + -9.0, + -9.0, + -9.0, + 9.97594e-10, + -9.0, + 1.4261355e-9, + 1.2249998e-9, + -2.208832e-6, + -9.0, + -7.707604e-8, + -9.0, + -5.329016e-7, + -9.0, + -9.0 + ] + ], + "get_is_mu": [ + [ + 0.0, + 1.0, + 0.0 + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + ], + "get_JetDistSig": [ + [ + -67.947235, + 86.78783, + -9.0 + ], + [ + -9.0, + -9.0, + 20.286007, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 33.79303, + -9.0, + -9.0, + -73.43549, + 17.754808, + -9.0, + -9.0, + -9.0, + 10.082263, + -9.0, + 27.870596, + -9.0, + 18.321705, + -9.0, + 7.818067, + -9.0, + -5.3476977, + -9.0, + -9.0, + -9.0, + -9.0, + 19.269228, + 13.5011215, + 9.805508, + -9.0, + -9.0, + 16.275623, + 32.900364, + -9.0, + -222.94427, + -22.069105, + -9.0, + -20.039652, + -9.0, + -9.0, + 161.43437, + -39.78606, + -9.0, + -72.438614, + -106.28761, + -9.0, + -9.0, + -9.0, + -87.568756, + -9.0, + -80.52304, + -87.215324, + -25.640558, + -9.0, + -49.46312, + -9.0, + -80.44129, + -9.0, + -9.0 + ] + ], + "get_erel_cluster": [ + [ + 0.30052024, + 0.61178416, + 0.08769559 + ], + [ + 0.0002285015, + 0.00041751255, + 0.0140406685, + 0.0029011106, + 0.0133055225, + 0.0064996867, + 0.0103801945, + 0.034518395, + 0.0148588205, + 0.00061151985, + 0.0038123084, + 0.05064229, + 0.043172225, + 0.0040380377, + 0.12879983, + 0.041509923, + 0.010277251, + 0.002955362, + 0.04229943, + 0.020724483, + 0.12897366, + 0.0039816233, + 0.04105549, + 0.0019617225, + 0.011026359, + 0.0024350667, + 0.0012780831, + 0.0016988798, + 0.010518281, + 0.0029863517, + 0.002539914, + 0.0047733793, + 0.0012985428, + 0.0017803428, + 0.0065732757, + 0.0068546487, + 0.010186248, + 0.011113481, + 0.007983463, + 0.013974096, + 0.0067465967, + 0.012652949, + 0.0010076154, + 0.010591866, + 0.010247781, + 0.0013555428, + 0.023469541, + 0.054581534, + 0.008684006, + 0.020898318, + 0.013309605, + 0.03610467, + 0.007138054, + 0.022807734, + 0.02647108, + 0.0030855383, + 0.0021583948, + 0.006362425, + 0.00032559285, + 0.0068347463, + 0.004439515, + 0.001740927 + ] + ], + "get_btagSip2dSig": [ + [ + -0.8752023, + 3.6409461, + -9.0 + ], + [ + -9.0, + -9.0, + -1.3618871, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -73.746315, + -9.0, + -9.0, + -64.52851, + -2.42035, + -9.0, + -9.0, + -9.0, + -1.7638716, + -9.0, + -1.2636894, + -9.0, + -3.5361004, + -9.0, + -11.560585, + -9.0, + 1.403294, + -9.0, + -9.0, + -9.0, + -9.0, + -1.0657609, + 0.03416786, + -11.497127, + -9.0, + -9.0, + 0.31063297, + 24.691864, + -9.0, + 138.90831, + -0.2649694, + -9.0, + 1.7638229, + -9.0, + -9.0, + -153.34438, + 5.979673, + -9.0, + 0.59213567, + 3.2991476, + -9.0, + -9.0, + -9.0, + 1.6434427, + -9.0, + 2.4296033, + 0.56557465, + 10.523689, + -9.0, + 11.835416, + -9.0, + 48.2779, + -9.0, + -9.0 + ] + ], + "get_dxydz": [ + [ + -9.620052e-8, + -5.600003e-8, + -9.0 + ], + [ + -9.0, + -9.0, + 3.1596717e-6, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 0.012508242, + -9.0, + -9.0, + -0.00079761405, + 2.3387707e-7, + -9.0, + -9.0, + -9.0, + 2.1446223e-5, + -9.0, + 7.7303774e-8, + -9.0, + 1.14505605e-8, + -9.0, + 1.0531564e-5, + -9.0, + 1.501868e-8, + -9.0, + -9.0, + -9.0, + -9.0, + 2.3817522e-5, + 3.2701962e-5, + 6.888741e-6, + -9.0, + -9.0, + 8.107701e-7, + 1.0965776e-6, + -9.0, + 1.20913555e-5, + 9.601748e-6, + -9.0, + 7.7685826e-7, + -9.0, + -9.0, + -6.9364123e-6, + 1.0851261e-6, + -9.0, + -2.1804137e-8, + -2.8028191e-8, + -9.0, + -9.0, + -9.0, + -2.036327e-8, + -9.0, + -2.2226248e-8, + -2.0300162e-8, + 2.3654367e-5, + -9.0, + 8.6529775e-7, + -9.0, + 5.8564556e-6, + -9.0, + -9.0 + ] + ], + "get_dptdpt": [ + [ + 1.4440912e-15, + 4.9895355e-16, + -9.0 + ], + [ + -9.0, + -9.0, + 1.2686038e-12, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 1.626925e-11, + -9.0, + -9.0, + 1.5390333e-12, + 1.1786605e-12, + -9.0, + -9.0, + -9.0, + 4.0698586e-12, + -9.0, + 5.2170356e-13, + -9.0, + 3.7192593e-13, + -9.0, + 7.962417e-12, + -9.0, + 1.2781154e-14, + -9.0, + -9.0, + -9.0, + -9.0, + 3.1703e-13, + 6.082442e-13, + 3.9173085e-14, + -9.0, + -9.0, + 3.083924e-14, + 3.145163e-14, + -9.0, + 2.0135442e-14, + 1.9287401e-13, + -9.0, + 3.018092e-14, + -9.0, + -9.0, + 1.6187638e-14, + 2.4226936e-14, + -9.0, + 3.9582684e-15, + 1.7403378e-15, + -9.0, + -9.0, + -9.0, + 2.5870614e-15, + -9.0, + 4.2983945e-15, + 3.7578257e-15, + 2.803512e-13, + -9.0, + 3.1707503e-14, + -9.0, + 3.8093006e-14, + -9.0, + -9.0 + ] + ], + "get_mass": [ + [ + 0.00051099894, + 0.105658375, + 0.00072206435 + ], + [ + -7.2657667e-6, + -1.211525e-5, + 0.13957039, + -6.145881e-5, + -0.00067013147, + -0.00021608449, + -0.0007735479, + -0.0012316033, + 0.13957039, + 2.5103373e-5, + 0.0001599524, + 0.13957039, + 0.13957039, + 0.00018150113, + 0.49770576, + 0.4976104, + 0.13957039, + 8.2336395e-5, + 0.13957039, + 0.49761188, + 0.13957039, + -0.00021552334, + 0.13957039, + -0.00011616132, + 0.13957039, + -5.0581315e-5, + -6.269215e-5, + -2.5549514e-5, + 0.0002591667, + 0.13957039, + 0.13957039, + 0.13957039, + -4.007986e-5, + 9.465513e-5, + 0.13957039, + 0.13957039, + -0.00023487065, + 0.13957039, + 0.13957039, + -0.0006808691, + 0.13957039, + -0.0005811269, + -4.9970153e-5, + 0.13957039, + 0.13957039, + 7.594635e-5, + 0.13957039, + 0.13957039, + -0.0002462741, + 0.4976121, + -0.0009334438, + 0.105658375, + 9.1004724e-5, + 0.13957039, + 0.13957039, + 0.13957039, + 3.5387056e-5, + 0.13957039, + 1.1097479e-5, + 0.13957039, + -6.0317438e-5, + 7.500272e-5 + ] + ], + "get_dphidphi": [ + [ + 8.289226e-9, + 1.5574163e-9, + -9.0 + ], + [ + -9.0, + -9.0, + 3.914325e-6, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + 1.5631548e-6, + -9.0, + -9.0, + 1.6434498e-7, + 1.0319517e-6, + -9.0, + -9.0, + -9.0, + 1.02482445e-5, + -9.0, + 7.497372e-7, + -9.0, + 2.4019647e-7, + -9.0, + 3.925059e-6, + -9.0, + 5.138348e-7, + -9.0, + -9.0, + -9.0, + -9.0, + 6.7253104e-6, + 1.0306365e-5, + 2.2855952e-6, + -9.0, + -9.0, + 1.2857342e-6, + 1.3158515e-6, + -9.0, + 3.2016317e-7, + 3.5064627e-6, + -9.0, + 1.2575318e-6, + -9.0, + -9.0, + 2.5531864e-7, + 1.0124714e-6, + -9.0, + 1.180763e-7, + 2.3032962e-8, + -9.0, + -9.0, + -9.0, + 5.5896784e-8, + -9.0, + 1.336964e-7, + 1.0831965e-7, + 6.312202e-6, + -9.0, + 1.3230573e-6, + -9.0, + 1.9539762e-6, + -9.0, + -9.0 + ] + ], + "get_type": [ + [ + 0, + 0, + 22 + ], + [ + 22, + 22, + 0, + 22, + 22, + 22, + 22, + 22, + 0, + 22, + 22, + 0, + 0, + 22, + 130, + 130, + 0, + 22, + 0, + 130, + 0, + 22, + 0, + 22, + 0, + 22, + 22, + 22, + 22, + 0, + 0, + 0, + 22, + 22, + 0, + 0, + 22, + 0, + 0, + 22, + 0, + 22, + 22, + 0, + 0, + 22, + 0, + 0, + 22, + 130, + 22, + 0, + 22, + 0, + 0, + 0, + 22, + 0, + 22, + 0, + 22, + 22 + ] + ], + "get_cdz": [ + [ + 9.835033e-12, + 1.9992014e-12, + -9.0 + ], + [ + -9.0, + -9.0, + -1.2846249e-9, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -3.117526e-7, + -9.0, + -9.0, + 5.859213e-9, + -1.4880788e-9, + -9.0, + -9.0, + -9.0, + -3.693833e-10, + -9.0, + -5.488911e-10, + -9.0, + -4.5163775e-10, + -9.0, + -6.9510165e-8, + -9.0, + 1.7312186e-11, + -9.0, + -9.0, + -9.0, + -9.0, + 8.115266e-10, + 1.3901854e-9, + 1.05917344e-10, + -9.0, + -9.0, + 3.7050852e-11, + 4.403962e-11, + -9.0, + 9.8432554e-11, + 2.775145e-10, + -9.0, + 3.6246027e-11, + -9.0, + -9.0, + 4.1132903e-11, + 4.5157645e-11, + -9.0, + 1.5176778e-11, + 1.2333312e-11, + -9.0, + -9.0, + -9.0, + 1.4484883e-11, + -9.0, + 1.540259e-11, + 1.5397723e-11, + 7.532449e-10, + -9.0, + 3.8348703e-11, + -9.0, + 1.14603015e-10, + -9.0, + -9.0 + ] + ], + "get_btagSip2dVal": [ + [ + -0.0023730414, + 0.0070432117, + -9.0 + ], + [ + -9.0, + -9.0, + -0.029755257, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -13.5313015, + -9.0, + -9.0, + -4.109446, + -0.028878165, + -9.0, + -9.0, + -9.0, + -0.061065387, + -9.0, + -0.013139597, + -9.0, + -0.023180787, + -9.0, + -0.25246602, + -9.0, + 0.012407542, + -9.0, + -9.0, + -9.0, + -9.0, + -0.030000899, + 0.0011837726, + -0.19509333, + -9.0, + -9.0, + 0.0040815696, + 0.3277208, + -9.0, + 1.7234792, + -0.0054963706, + -9.0, + 0.022943396, + -9.0, + -9.0, + -1.7221894, + 0.070499085, + -9.0, + 0.0029449174, + 0.009541756, + -9.0, + -9.0, + -9.0, + 0.0062310095, + -9.0, + 0.012660839, + 0.0027262159, + 0.28730667, + -9.0, + 0.15754971, + -9.0, + 0.7631629, + -9.0, + -9.0 + ] + ], + "get_is_el": [ + [ + 1.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + ], + "get_dz": [ + [ + -1.2500461, + -1.2301915, + -9.0 + ], + [ + -9.0, + -9.0, + -1.2211165, + -9.0, + -9.0, + -9.0, + -9.0, + -9.0, + -28.373814, + -9.0, + -9.0, + 48.975826, + -1.2347836, + -9.0, + -9.0, + -9.0, + -1.2415172, + -9.0, + -1.2364849, + -9.0, + -1.2729397, + -9.0, + -3.0270064, + -9.0, + -1.2370702, + -9.0, + -9.0, + -9.0, + -9.0, + -1.2279438, + -1.2468957, + -1.3070565, + -9.0, + -9.0, + -1.2371976, + -1.229995, + -9.0, + -3.2748, + -1.2084076, + -9.0, + -1.1798426, + -9.0, + -9.0, + 0.9103305, + -1.4902136, + -9.0, + -1.2341998, + -1.2185807, + -9.0, + -9.0, + -9.0, + -1.240396, + -9.0, + -1.2353859, + -1.2281562, + -1.2656493, + -9.0, + -1.1130629, + -9.0, + -1.4651086, + -9.0, + -9.0 + ] + ], + "get_thetarel_phirel_cluster_theta": [ + [ + 2.966602, + 2.6673548, + 2.6804037 + ], + [ + 1.2688109, + 1.104194, + 0.8537709, + 0.90194684, + 0.7381785, + 0.7217353, + 0.7827701, + 0.7615588, + 0.768493, + 0.855986, + 0.8140161, + 0.7669116, + 0.77883875, + 0.79305065, + 0.7503496, + 0.7141931, + 0.6717117, + 0.8081033, + 0.76337284, + 0.6997838, + 0.6875741, + 0.62219524, + 0.681343, + 0.37576556, + 1.9060644, + 1.8329093, + 2.2823558, + 2.1103423, + 2.0034032, + 1.4725348, + 1.695954, + 1.6296993, + 1.7051743, + 1.669573, + 1.7384486, + 1.6747093, + 1.6995277, + 1.2766685, + 1.148414, + 1.6260409, + 1.7339329, + 1.6537427, + 1.6207951, + 1.4655788, + 1.4909564, + 1.4565848, + 1.7335267, + 1.798572, + 1.7681736, + 1.8011085, + 1.6599228, + 1.7002355, + 1.7198068, + 1.6526482, + 1.5753503, + 1.5842625, + 1.5982518, + 1.4586637, + 1.5046724, + 1.2052091, + 1.3525188, + 1.3295034 + ] + ] +} diff --git a/test/runtests.jl b/test/runtests.jl index ece4ae6c..8610e13e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -45,6 +45,9 @@ function main() # SoftKiller tests include("test-softkiller.jl") + # Jet flavour tests + include("test-jet-flavour-utils.jl") + # Test with Aqua (https://juliatesting.github.io/Aqua.jl/stable/) include("test-aqua.jl") end diff --git a/test/test-jet-flavour-tagging.jl b/test/test-jet-flavour-tagging.jl new file mode 100644 index 00000000..8e0b9589 --- /dev/null +++ b/test/test-jet-flavour-tagging.jl @@ -0,0 +1,238 @@ +# Tests for Jet Flavour Tagging Extension + +include("common.jl") + +# Check if the extension is available +const HAS_FLAVOUR_TAGGING = try + using EDM4hep + using ONNXRunTime + using JSON + true +catch + false +end + +if !HAS_FLAVOUR_TAGGING + @warn "Jet Flavour Tagging extension dependencies not available, skipping tests" +else + using EDM4hep + using EDM4hep.RootIO + using StaticArrays + using LorentzVectorHEP + using JetReconstruction + using LoopVectorization + using JSON + using ONNXRunTime + using StructArrays: StructVector + + @testset "Jet Flavour Tagging Extension" begin + @testset "Extension Functions Defined" begin + # Verify all exported functions are defined + @test isdefined(JetReconstruction, :build_constituents_cluster) + @test isdefined(JetReconstruction, :extract_features) + @test isdefined(JetReconstruction, :setup_onnx_runtime) + @test isdefined(JetReconstruction, :prepare_input_tensor) + @test isdefined(JetReconstruction, :get_weights) + @test isdefined(JetReconstruction, :get_weight) + end + + # Navigate to the examples folder from the test folder + package_root = dirname(@__DIR__) # Go up from test/ to JetReconstruction.jl/ + onnx_data_dir = joinpath(package_root, "examples", "flavour-tagging", "data", + "wc_pt_7classes_12_04_2023") + edm4hep_dir = joinpath(package_root, "examples", "flavour-tagging", "data") + + # Paths to the actual files + onnx_path = joinpath(onnx_data_dir, "fccee_flavtagging_edm4hep_wc_v1.onnx") + json_path = joinpath(onnx_data_dir, "fccee_flavtagging_edm4hep_wc_v1.json") + edm4hep_example_file = joinpath(edm4hep_dir, "events_080263084.root") + + # Check if files exist + has_onnx_file = isfile(onnx_path) + has_json_file = isfile(json_path) + has_edm4hep_file = isfile(edm4hep_example_file) + has_test_files = has_onnx_file && has_json_file && has_edm4hep_file + + if !has_test_files + @warn "Test files not found in examples folder. Missing files:" * + (!has_onnx_file ? "\n - $onnx_path" : "") * + (!has_json_file ? "\n - $json_path" : "") * + (!has_edm4hep_file ? "\n - $edm4hep_example_file" : "") + end + + if has_test_files + @testset "Full Flavour Tagging Pipeline" begin + # Load the EDM4hep data + reader = RootIO.Reader(edm4hep_example_file) + events = RootIO.get(reader, "events") + @test length(events) > 0 + + # Use event #16 as in the notebook + event_id = 16 + evt = events[event_id] + + # Get reconstructed particles and tracks + recps = RootIO.get(reader, evt, "ReconstructedParticles") + @test length(recps) > 0 + + tracks = RootIO.get(reader, evt, "EFlowTrack_1") + @test length(tracks) >= 0 + + # Get needed collections for feature extraction + bz = RootIO.get(reader, evt, "magFieldBz", register = false)[1] + @test isa(bz, Float32) + + trackdata = RootIO.get(reader, evt, "EFlowTrack") + trackerhits = RootIO.get(reader, evt, "TrackerHits") + gammadata = RootIO.get(reader, evt, "EFlowPhoton") + nhdata = RootIO.get(reader, evt, "EFlowNeutralHadron") + calohits = RootIO.get(reader, evt, "CalorimeterHits") + dNdx = RootIO.get(reader, evt, "EFlowTrack_2") + track_L = RootIO.get(reader, evt, "EFlowTrack_L", register = false) + + @testset "Jet Reconstruction" begin + # Cluster jets using the EEkt algorithm + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, + algorithm = JetAlgorithm.EEKt) + @test isa(cs, ClusterSequence) + + # Get 2 exclusive jets + jets = exclusive_jets(cs; njets = 2, T = EEJet) + @test length(jets) == 2 + @test isa(jets[1], EEJet) + end + + @testset "build_constituents_cluster" begin + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, + algorithm = JetAlgorithm.EEKt) + jets = exclusive_jets(cs; njets = 2, T = EEJet) + + # Get constituent indices for each jet + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + @test length(constituent_indices) == 2 + @test all(indices -> length(indices) > 0, constituent_indices) + + # Build constituents + jet_constituents = JetReconstruction.build_constituents_cluster(recps, + constituent_indices) + @test length(jet_constituents) == 2 + @test all(jc -> isa(jc, StructVector{ReconstructedParticle}), + jet_constituents) + end + + @testset "extract_features" begin + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, + algorithm = JetAlgorithm.EEKt) + jets = exclusive_jets(cs; njets = 2, T = EEJet) + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + jet_constituents = JetReconstruction.build_constituents_cluster(recps, + constituent_indices) + + # Extract features + feature_data = JetReconstruction.extract_features(jets, + jet_constituents, + tracks, + bz, + track_L, + trackdata, + trackerhits, + gammadata, + nhdata, + calohits, + dNdx) + + @test isa(feature_data, Dict) + @test haskey(feature_data, "Particles") + @test haskey(feature_data, "Tracks") + end + + @testset "setup_onnx_runtime" begin + # Test ONNX runtime setup + model = JetReconstruction.setup_onnx_runtime(onnx_path, json_path) + @test isa(model, ONNXRunTime.InferenceSession) + + # Also test loading config separately + config = JSON.parsefile(json_path) + @test isa(config, Dict) + @test haskey(config, "output_names") + @test length(config["output_names"]) == 5 # 5 flavor classes + end + + @testset "prepare_input_tensor" begin + # Setup everything needed + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, + algorithm = JetAlgorithm.EEKt) + jets = exclusive_jets(cs; njets = 2, T = EEJet) + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + jet_constituents = JetReconstruction.build_constituents_cluster(recps, + constituent_indices) + + feature_data = JetReconstruction.extract_features(jets, + jet_constituents, + tracks, bz, track_L, + trackdata, + trackerhits, + gammadata, nhdata, + calohits, dNdx) + + config = JSON.parsefile(json_path) + + # Prepare input tensors + input_tensors = JetReconstruction.prepare_input_tensor(jet_constituents, + jets, config, + feature_data) + + @test isa(input_tensors, Dict) + @test all(v -> isa(v, Array), values(input_tensors)) + end + + @testset "get_weights and get_weight" begin + # Full pipeline to get weights + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, + algorithm = JetAlgorithm.EEKt) + jets = exclusive_jets(cs; njets = 2, T = EEJet) + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + jet_constituents = JetReconstruction.build_constituents_cluster(recps, + constituent_indices) + + feature_data = JetReconstruction.extract_features(jets, + jet_constituents, + tracks, bz, track_L, + trackdata, + trackerhits, + gammadata, nhdata, + calohits, dNdx) + + config = JSON.parsefile(json_path) + model = JetReconstruction.setup_onnx_runtime(onnx_path, json_path) + + # Get weights + weights = JetReconstruction.get_weights(0, # Thread slot + feature_data, + jets, + jet_constituents, + config, + model) + + @test isa(weights, Vector{Vector{Float32}}) + @test length(weights) == 2 # One weight vector per jet + @test all(w -> length(w) == 5, weights) # 5 flavor classes per jet + + # Test get_weight for extracting individual scores + for i in 0:4 # 5 flavor classes (0-indexed) + score = JetReconstruction.get_weight(weights, i) + @test isa(score, Vector{Float32}) + @test length(score) == 2 # One score per jet + @test all(s -> 0.0 <= s <= 1.0, score) # Probabilities should be in [0,1] + end + + # Test that probabilities sum to approximately 1 + for jet_idx in 1:2 + prob_sum = sum(weights[jet_idx]) + @test isapprox(prob_sum, 1.0, atol = 1e-4) + end + end + end + end + end +end diff --git a/test/test-jet-flavour-utils.jl b/test/test-jet-flavour-utils.jl new file mode 100644 index 00000000..bc85399a --- /dev/null +++ b/test/test-jet-flavour-utils.jl @@ -0,0 +1,326 @@ +# Tests for Jet Flavour Tagging Constituent Utilities + +include("common.jl") + +# Check if the extension is available +const HAS_FLAVOUR_TAGGING = try + using EDM4hep + using ONNXRunTime + using PhysicalConstants + using JSON + true +catch + false +end + +if !HAS_FLAVOUR_TAGGING + @warn "Jet Flavour Tagging extension dependencies not available, skipping tests" +else + using EDM4hep + using EDM4hep.RootIO + using LorentzVectorHEP + using ONNXRunTime + using PhysicalConstants + using JSON + using StructArrays + using JetReconstruction + + @testset "Jet Flavour Constituent Utilities" begin + # Navigate to find the test reference file + reference_file = joinpath(@__DIR__, "data", "jet-flavour-utils-event15.json") + + # Check if reference file exists + if !isfile(reference_file) + @warn "Reference output file not found at: $reference_file" * + "\nRun generate-jet-constituent-utils-outputs.jl first to generate it." + @test_skip "Reference outputs not available" + else + # Load reference outputs + reference_outputs = JSON.parsefile(reference_file) + + # Load the same data used to generate reference + package_root = dirname(@__DIR__) + edm4hep_dir = joinpath(package_root, "examples", "flavour-tagging", "data") + edm4hep_example_file = joinpath(edm4hep_dir, "events_080263084.root") + + if !isfile(edm4hep_example_file) + @test_skip "EDM4hep example file not found" + else + # Load data + reader = RootIO.Reader(edm4hep_example_file) + events = RootIO.get(reader, "events") + event_id = reference_outputs["input_info"]["event_id"] + evt = events[event_id] + + # Get all necessary data + recps = RootIO.get(reader, evt, "ReconstructedParticles") + tracks = RootIO.get(reader, evt, "EFlowTrack_1") + bz = Float32(reference_outputs["input_info"]["bz"]) + trackdata = RootIO.get(reader, evt, "EFlowTrack") + trackerhits = RootIO.get(reader, evt, "TrackerHits") + gammadata = RootIO.get(reader, evt, "EFlowPhoton") + nhdata = RootIO.get(reader, evt, "EFlowNeutralHadron") + calohits = RootIO.get(reader, evt, "CalorimeterHits") + dNdx = RootIO.get(reader, evt, "EFlowTrack_2") + track_L = RootIO.get(reader, evt, "EFlowTrack_L", register = false) + + # Create vertex from reference + v_info = reference_outputs["input_info"]["vertex"] + V = LorentzVector(v_info["x"], v_info["y"], v_info["z"], v_info["t"]) + + # Access the extension module first + ext_mod = Base.get_extension(JetReconstruction, :JetFlavourTagging) + @test !isnothing(ext_mod) + + # Access the helper module within the extension + helper_mod = ext_mod.JetFlavourHelper + + # Reconstruct jets + cs = jet_reconstruct(recps; p = 1.0, R = 2.0, algorithm = JetAlgorithm.EEKt) + jets = exclusive_jets(cs; njets = 2, T = EEJet) + + # Get jet constituents + constituent_indices = [constituent_indexes(jet, cs) for jet in jets] + jet_constituents = ext_mod.build_constituents_cluster(recps, + constituent_indices) + + # Helper function to compare outputs + function compare_outputs(actual, reference, name; atol = 1e-5, rtol = 1e-5) + # Handle special float values + function normalize_value(x) + if isa(x, String) + if x == "NaN" + return NaN + elseif x == "Inf" + return Inf + elseif x == "-Inf" + return -Inf + end + end + return x + end + + # Recursively normalize arrays + function normalize_array(arr) + if isa(arr, AbstractArray) + return [normalize_array(x) for x in arr] + else + return normalize_value(arr) + end + end + + ref_normalized = normalize_array(reference) + + @testset "$name" begin + if isa(actual, Vector{Vector{T}} where {T}) + @test length(actual) == length(ref_normalized) + for i in 1:length(actual) + @test length(actual[i]) == length(ref_normalized[i]) + for j in 1:length(actual[i]) + act_val = actual[i][j] + ref_val = ref_normalized[i][j] + if isnan(ref_val) + @test isnan(act_val) + elseif isinf(ref_val) + @test isinf(act_val) && + sign(act_val) == sign(ref_val) + else + @test isapprox(act_val, ref_val; atol = atol, + rtol = rtol) + end + end + end + else + @test false # Unexpected output type + end + end + end + + @testset "Basic Kinematics" begin + compare_outputs(helper_mod.get_pt(jet_constituents), + reference_outputs["get_pt"], "get_pt") + compare_outputs(helper_mod.get_p(jet_constituents), + reference_outputs["get_p"], "get_p") + compare_outputs(helper_mod.get_e(jet_constituents), + reference_outputs["get_e"], "get_e") + compare_outputs(helper_mod.get_type(jet_constituents), + reference_outputs["get_type"], "get_type") + compare_outputs(helper_mod.get_mass(jet_constituents), + reference_outputs["get_mass"], "get_mass") + compare_outputs(helper_mod.get_charge(jet_constituents), + reference_outputs["get_charge"], "get_charge") + compare_outputs(helper_mod.get_theta(jet_constituents), + reference_outputs["get_theta"], "get_theta") + compare_outputs(helper_mod.get_phi(jet_constituents), + reference_outputs["get_phi"], "get_phi") + compare_outputs(helper_mod.get_y(jet_constituents), + reference_outputs["get_y"], "get_y") + compare_outputs(helper_mod.get_eta(jet_constituents), + reference_outputs["get_eta"], "get_eta") + compare_outputs(helper_mod.get_Bz(jet_constituents, tracks), + reference_outputs["get_Bz"], "get_Bz") + end + + @testset "Track Parameters" begin + compare_outputs(helper_mod.get_dxy(jet_constituents, tracks, V, bz), + reference_outputs["get_dxy"], "get_dxy") + compare_outputs(helper_mod.get_dz(jet_constituents, tracks, V, bz), + reference_outputs["get_dz"], "get_dz") + compare_outputs(helper_mod.get_phi0(jet_constituents, tracks, V, bz), + reference_outputs["get_phi0"], "get_phi0") + compare_outputs(helper_mod.get_c(jet_constituents, tracks, bz), + reference_outputs["get_c"], "get_c") + compare_outputs(helper_mod.get_ct(jet_constituents, tracks, bz), + reference_outputs["get_ct"], "get_ct") + end + + @testset "Covariance Matrix Elements" begin + compare_outputs(helper_mod.get_dxydxy(jet_constituents, tracks), + reference_outputs["get_dxydxy"], "get_dxydxy") + compare_outputs(helper_mod.get_dphidxy(jet_constituents, tracks), + reference_outputs["get_dphidxy"], "get_dphidxy") + compare_outputs(helper_mod.get_dphidphi(jet_constituents, tracks), + reference_outputs["get_dphidphi"], "get_dphidphi") + compare_outputs(helper_mod.get_dxyc(jet_constituents, tracks), + reference_outputs["get_dxyc"], "get_dxyc") + compare_outputs(helper_mod.get_phic(jet_constituents, tracks), + reference_outputs["get_phic"], "get_phic") + compare_outputs(helper_mod.get_dptdpt(jet_constituents, tracks), + reference_outputs["get_dptdpt"], "get_dptdpt") + compare_outputs(helper_mod.get_dxydz(jet_constituents, tracks), + reference_outputs["get_dxydz"], "get_dxydz") + compare_outputs(helper_mod.get_phidz(jet_constituents, tracks), + reference_outputs["get_phidz"], "get_phidz") + compare_outputs(helper_mod.get_cdz(jet_constituents, tracks), + reference_outputs["get_cdz"], "get_cdz") + compare_outputs(helper_mod.get_dzdz(jet_constituents, tracks), + reference_outputs["get_dzdz"], "get_dzdz") + compare_outputs(helper_mod.get_dxyctgtheta(jet_constituents, tracks), + reference_outputs["get_dxyctgtheta"], "get_dxyctgtheta") + compare_outputs(helper_mod.get_phictgtheta(jet_constituents, tracks), + reference_outputs["get_phictgtheta"], "get_phictgtheta") + compare_outputs(helper_mod.get_cctgtheta(jet_constituents, tracks), + reference_outputs["get_cctgtheta"], "get_cctgtheta") + compare_outputs(helper_mod.get_dlambdadz(jet_constituents, tracks), + reference_outputs["get_dlambdadz"], "get_dlambdadz") + compare_outputs(helper_mod.get_detadeta(jet_constituents, tracks), + reference_outputs["get_detadeta"], "get_detadeta") + end + + @testset "Particle Identification" begin + compare_outputs(helper_mod.get_is_mu(jet_constituents), + reference_outputs["get_is_mu"], "get_is_mu") + compare_outputs(helper_mod.get_is_el(jet_constituents), + reference_outputs["get_is_el"], "get_is_el") + compare_outputs(helper_mod.get_is_charged_had(jet_constituents), + reference_outputs["get_is_charged_had"], + "get_is_charged_had") + compare_outputs(helper_mod.get_is_gamma(jet_constituents), + reference_outputs["get_is_gamma"], "get_is_gamma") + compare_outputs(helper_mod.get_is_neutral_had(jet_constituents), + reference_outputs["get_is_neutral_had"], + "get_is_neutral_had") + end + + @testset "Relative Kinematics" begin + compare_outputs(helper_mod.get_erel_cluster(jets, jet_constituents), + reference_outputs["get_erel_cluster"], + "get_erel_cluster") + compare_outputs(helper_mod.get_erel_log_cluster(jets, jet_constituents), + reference_outputs["get_erel_log_cluster"], + "get_erel_log_cluster") + compare_outputs(helper_mod.get_thetarel_cluster(jets, jet_constituents), + reference_outputs["get_thetarel_cluster"], + "get_thetarel_cluster") + compare_outputs(helper_mod.get_phirel_cluster(jets, jet_constituents), + reference_outputs["get_phirel_cluster"], + "get_phirel_cluster") + + # Combined function returns tuple + theta_phi_rel = helper_mod.get_thetarel_phirel_cluster(jets, + jet_constituents) + compare_outputs(theta_phi_rel[1], + reference_outputs["get_thetarel_phirel_cluster_theta"], + "get_thetarel_phirel_cluster_theta") + compare_outputs(theta_phi_rel[2], + reference_outputs["get_thetarel_phirel_cluster_phi"], + "get_thetarel_phirel_cluster_phi") + end + + @testset "Special Measurements" begin + is_charged_had = helper_mod.get_is_charged_had(jet_constituents) + compare_outputs(helper_mod.get_dndx(jet_constituents, dNdx, trackdata, + is_charged_had), + reference_outputs["get_dndx"], "get_dndx") + compare_outputs(helper_mod.get_mtof(jet_constituents, track_L, + trackdata, trackerhits, + gammadata, nhdata, calohits, V), + reference_outputs["get_mtof"], "get_mtof") + end + + @testset "Impact Parameters and Jet Distance" begin + # Calculate base values + D0 = helper_mod.get_dxy(jet_constituents, tracks, V, bz) + Z0 = helper_mod.get_dz(jet_constituents, tracks, V, bz) + phi0 = helper_mod.get_phi0(jet_constituents, tracks, V, bz) + err2_D0 = helper_mod.get_dxydxy(jet_constituents, tracks) + err2_Z0 = helper_mod.get_dzdz(jet_constituents, tracks) + + # 2D impact parameter + sip2d_val = helper_mod.get_Sip2dVal_clusterV(jets, D0, phi0, bz) + compare_outputs(sip2d_val, reference_outputs["get_Sip2dVal_clusterV"], + "get_Sip2dVal_clusterV") + + btag_sip2d_val = helper_mod.get_btagSip2dVal(jets, D0, phi0, bz) + compare_outputs(btag_sip2d_val, reference_outputs["get_btagSip2dVal"], + "get_btagSip2dVal") + + compare_outputs(helper_mod.get_Sip2dSig(sip2d_val, err2_D0), + reference_outputs["get_Sip2dSig"], "get_Sip2dSig") + compare_outputs(helper_mod.get_btagSip2dSig(btag_sip2d_val, err2_D0), + reference_outputs["get_btagSip2dSig"], + "get_btagSip2dSig") + + # 3D impact parameter + sip3d_val = helper_mod.get_Sip3dVal_clusterV(jets, D0, Z0, phi0, bz) + compare_outputs(sip3d_val, reference_outputs["get_Sip3dVal_clusterV"], + "get_Sip3dVal_clusterV") + + btag_sip3d_val = helper_mod.get_btagSip3dVal(jets, D0, Z0, phi0, bz) + compare_outputs(btag_sip3d_val, reference_outputs["get_btagSip3dVal"], + "get_btagSip3dVal") + + compare_outputs(helper_mod.get_Sip3dSig(sip3d_val, err2_D0, err2_Z0), + reference_outputs["get_Sip3dSig"], "get_Sip3dSig") + compare_outputs(helper_mod.get_btagSip3dSig(btag_sip3d_val, err2_D0, + err2_Z0), + reference_outputs["get_btagSip3dSig"], + "get_btagSip3dSig") + + # Jet distance + jet_dist_val = helper_mod.get_JetDistVal_clusterV(jets, + jet_constituents, D0, + Z0, phi0, bz) + compare_outputs(jet_dist_val, + reference_outputs["get_JetDistVal_clusterV"], + "get_JetDistVal_clusterV") + + btag_jet_dist_val = helper_mod.get_btagJetDistVal(jets, + jet_constituents, D0, + Z0, phi0, bz) + compare_outputs(btag_jet_dist_val, + reference_outputs["get_btagJetDistVal"], + "get_btagJetDistVal") + + compare_outputs(helper_mod.get_JetDistSig(jet_dist_val, err2_D0, + err2_Z0), + reference_outputs["get_JetDistSig"], "get_JetDistSig") + compare_outputs(helper_mod.get_btagJetDistSig(btag_jet_dist_val, + err2_D0, err2_Z0), + reference_outputs["get_btagJetDistSig"], + "get_btagJetDistSig") + end + end + end + end +end