Adding ProbNNs for Downstream, Upstream and ProbNNghost

Hlt/RecoConf/python/RecoConf/protoparticles.py

@@ -27,12 +27,20 @@ from PyConf.Tools import (
     ChargedProtoParticleAddCombineDLLs,
 )
 
-from PyConf.Tools import (ChargedProtoParticleAddGlobalPID_Electron,
-                          ChargedProtoParticleAddGlobalPID_Muon,
-                          ChargedProtoParticleAddGlobalPID_Pion,
-                          ChargedProtoParticleAddGlobalPID_Kaon,
-                          ChargedProtoParticleAddGlobalPID_Proton,
-                          ChargedProtoParticleAddGlobalPID_Ghost)
+from PyConf.Tools import (
+    ChargedProtoParticleAddGlobalPID_Electron,
+    ChargedProtoParticleAddGlobalPID_Muon,
+    ChargedProtoParticleAddGlobalPID_Pion,
+    ChargedProtoParticleAddGlobalPID_Kaon,
+    ChargedProtoParticleAddGlobalPID_Proton,
+    ChargedProtoParticleAddGlobalPID_Ghost,
+    ChargedProtoParticleAddGlobalPID_Electron_Upstream,
+    ChargedProtoParticleAddGlobalPID_Muon_Upstream,
+    ChargedProtoParticleAddGlobalPID_Pion_Upstream,
+    ChargedProtoParticleAddGlobalPID_Kaon_Upstream,
+    ChargedProtoParticleAddGlobalPID_Proton_Upstream,
+    ChargedProtoParticleAddGlobalPID_Ghost_Upstream,
+)
 
 import Functors as F
 
@@ -53,20 +61,35 @@ def hypo_estimator(calo_pids):
 
 def getProbNNModel(PIDType, TrackType):
     # TODO add differentiation for TrackType (probably only for Upstream)
-    if PIDType == 'Electron':
-        return ChargedProtoParticleAddGlobalPID_Electron
-    elif PIDType == 'Muon':
-        return ChargedProtoParticleAddGlobalPID_Muon
-    elif PIDType == 'Pion':
-        return ChargedProtoParticleAddGlobalPID_Pion
-    elif PIDType == 'Kaon':
-        return ChargedProtoParticleAddGlobalPID_Kaon
-    elif PIDType == 'Proton':
-        return ChargedProtoParticleAddGlobalPID_Proton
-    elif PIDType == 'Ghost':
-        return ChargedProtoParticleAddGlobalPID_Ghost
-    else:
-        raise ValueError(f"not a configured PID type: '{PIDType}'")
+    if TrackType in ['Long', 'Downstream']:
+        if PIDType == 'Electron':
+            return ChargedProtoParticleAddGlobalPID_Electron
+        elif PIDType == 'Muon':
+            return ChargedProtoParticleAddGlobalPID_Muon
+        elif PIDType == 'Pion':
+            return ChargedProtoParticleAddGlobalPID_Pion
+        elif PIDType == 'Kaon':
+            return ChargedProtoParticleAddGlobalPID_Kaon
+        elif PIDType == 'Proton':
+            return ChargedProtoParticleAddGlobalPID_Proton
+        elif PIDType == 'Ghost':
+            return ChargedProtoParticleAddGlobalPID_Ghost
+    elif TrackType == 'Upstream':
+        if PIDType == 'Electron':
+            return ChargedProtoParticleAddGlobalPID_Electron_Upstream
+        elif PIDType == 'Muon':
+            return ChargedProtoParticleAddGlobalPID_Muon_Upstream
+        elif PIDType == 'Pion':
+            return ChargedProtoParticleAddGlobalPID_Pion_Upstream
+        elif PIDType == 'Kaon':
+            return ChargedProtoParticleAddGlobalPID_Kaon_Upstream
+        elif PIDType == 'Proton':
+            return ChargedProtoParticleAddGlobalPID_Proton_Upstream
+        elif PIDType == 'Ghost':
+            return ChargedProtoParticleAddGlobalPID_Ghost_Upstream
+    raise ValueError(
+        f"not a configured PID type: '{PIDType}' for track type '{TrackType}'")
+    return None
 
 
 @configurable
@@ -213,23 +236,37 @@ def make_charged_protoparticles(tracks,
                 'Proton':
                 'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_proton_long.json',
                 'Ghost':
-                ''
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_ghost_long.json'
             },
+            # use long for downstream as PID features are same, even used RICH1 gas
+            # only minor difference is for brem base pid, typical brem emission chances are lower for electrons
             'Downstream': {
-                'Electron': '',
-                'Muon': '',
-                'Pion': '',
-                'Kaon': '',
-                'Proton': '',
-                'Ghost': ''
+                'Electron':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_electron_long.json',
+                'Muon':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_muon_long.json',
+                'Pion':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_pion_long.json',
+                'Kaon':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_kaon_long.json',
+                'Proton':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_proton_long.json',
+                'Ghost':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_ghost_long.json'
             },
             'Upstream': {
-                'Electron': '',
-                'Muon': '',
-                'Pion': '',
-                'Kaon': '',
-                'Proton': '',
-                'Ghost': ''
+                'Electron':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_electron_upstream.json',
+                'Muon':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_muon_upstream.json',
+                'Pion':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_pion_upstream.json',
+                'Kaon':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_kaon_upstream.json',
+                'Proton':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_proton_upstream.json',
+                'Ghost':
+                'paramfile://data/ProbNNs/2024/probnn_2024tune1_pp_ghost_upstream.json'
             },
         }
         for PIDType in ["Electron", "Muon", "Pion", "Kaon", "Proton", "Ghost"]: