Add protection against nan inputs for DeepMET

cms-sw · May 16, 2024 · 5c341f4 · 5c341f4
1 parent 43944b8
commit 5c341f4
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 11 deletions.
diff --git a/RecoMET/METPUSubtraction/interface/DeepMETHelp.h b/RecoMET/METPUSubtraction/interface/DeepMETHelp.h
@@ -6,6 +6,7 @@
 
 namespace deepmet_helper {
   float scale_and_rm_outlier(float val, float scale);
+  float rm_outlier(float val);
 
   static const std::unordered_map<int, int32_t> charge_embedding{{-1, 0}, {0, 1}, {1, 2}};
   static const std::unordered_map<int, int32_t> pdg_id_embedding{

diff --git a/RecoMET/METPUSubtraction/plugins/DeepMETProducer.cc b/RecoMET/METPUSubtraction/plugins/DeepMETProducer.cc
@@ -81,12 +81,12 @@ void DeepMETProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
     // fill the tensor
     // PF keys [b'PF_dxy', b'PF_dz', b'PF_eta', b'PF_mass', b'PF_pt', b'PF_puppiWeight', b'PF_px', b'PF_py']
     float* ptr = &input_.tensor<float, 3>()(0, i_pf, 0);
-    *ptr = pf.dxy();
-    *(++ptr) = pf.dz();
-    *(++ptr) = pf.eta();
-    *(++ptr) = pf.mass();
+    *ptr = rm_outlier(pf.dxy());
+    *(++ptr) = rm_outlier(pf.dz());
+    *(++ptr) = rm_outlier(pf.eta());
+    *(++ptr) = rm_outlier(pf.mass());
     *(++ptr) = scale_and_rm_outlier(pf.pt(), scale);
-    *(++ptr) = pf.puppiWeight();
+    *(++ptr) = rm_outlier(pf.puppiWeight());
     *(++ptr) = scale_and_rm_outlier(pf.px(), scale);
     *(++ptr) = scale_and_rm_outlier(pf.py(), scale);
     input_cat0_.tensor<float, 3>()(0, i_pf, 0) = charge_embedding.at(pf.charge());

diff --git a/RecoMET/METPUSubtraction/plugins/DeepMETSonicProducer.cc b/RecoMET/METPUSubtraction/plugins/DeepMETSonicProducer.cc
@@ -73,12 +73,12 @@ void DeepMETSonicProducer::acquire(edm::Event const& iEvent, edm::EventSetup con
     }
 
     // PF keys [b'PF_dxy', b'PF_dz', b'PF_eta', b'PF_mass', b'PF_pt', b'PF_puppiWeight', b'PF_px', b'PF_py']
-    vpfdata.push_back(pf.dxy());
-    vpfdata.push_back(pf.dz());
-    vpfdata.push_back(pf.eta());
-    vpfdata.push_back(pf.mass());
+    vpfdata.push_back(rm_outlier(pf.dxy()));
+    vpfdata.push_back(rm_outlier(pf.dz()));
+    vpfdata.push_back(rm_outlier(pf.eta()));
+    vpfdata.push_back(rm_outlier(pf.mass()));
     vpfdata.push_back(scale_and_rm_outlier(pf.pt(), scale_));
-    vpfdata.push_back(pf.puppiWeight());
+    vpfdata.push_back(rm_outlier(pf.puppiWeight()));
     vpfdata.push_back(scale_and_rm_outlier(pf.px(), scale_));
     vpfdata.push_back(scale_and_rm_outlier(pf.py(), scale_));
 

diff --git a/RecoMET/METPUSubtraction/src/DeepMETHelper.cc b/RecoMET/METPUSubtraction/src/DeepMETHelper.cc
@@ -1,11 +1,17 @@
 #include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
+#include <cmath>
 
 namespace deepmet_helper {
   float scale_and_rm_outlier(float val, float scale) {
     float ret_val = val * scale;
-    if (ret_val > 1e6 || ret_val < -1e6)
+    if (std::isnan(ret_val) || ret_val > 1e6 || ret_val < -1e6)
       return 0.;
     return ret_val;
   }
 
+  float rm_outlier(float val) {
+    if (std::isnan(val) || val > 1e6 || val < -1e6)
+      return 0.;
+    return val;
+  }
 }  // namespace deepmet_helper