Fixing the candidate flagging (#853)

Co-authored-by: Luigi Dello Stritto <[email protected]>
alisw · Jan 25, 2024 · 48b531b · 48b531b
1 parent 478b97b
commit 48b531b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/machine_learning_hep/bitwise.py b/machine_learning_hep/bitwise.py
@@ -21,9 +21,11 @@
 
 from .logger import get_logger
 
-def tag_bit_df(dfin, namebitmap, activatedbit):
+def tag_bit_df(dfin, namebitmap, activatedbit, absval = False):
     try:
         ar = dfin[namebitmap].to_numpy(dtype='int')
+        if absval:
+            ar = abs(ar)
         mask_on = reduce(operator.or_, ((1 << bit) for bit in activatedbit[0]), 0)
         mask_off = reduce(operator.or_, ((1 << bit) for bit in activatedbit[1]), 0)
         return np.logical_and(np.bitwise_and(ar, mask_on) == mask_on,

diff --git a/machine_learning_hep/processer.py b/machine_learning_hep/processer.py
@@ -427,7 +427,7 @@ def read_df(tree, df_base, var):
         # needs to be revisited for Run 3
         if self.mcordata == "mc":
             dfreco[self.v_ismcsignal] = np.array(tag_bit_df(dfreco, self.v_bitvar,
-                                                            self.b_mcsig), dtype=int)
+                                                            self.b_mcsig, True), dtype=int)
             dfreco[self.v_ismcprompt] = np.array(tag_bit_df(dfreco, self.v_bitvar_origrec,
                                                             self.b_mcsigprompt), dtype=int)
             dfreco[self.v_ismcfd] = np.array(tag_bit_df(dfreco, self.v_bitvar_origrec,
@@ -440,7 +440,7 @@ def read_df(tree, df_base, var):
                 dfreco[self.v_ismcfd] = np.logical_and(dfreco[self.v_ismcfd] == 1, mydf)
 
             dfreco[self.v_ismcbkg] = np.array(tag_bit_df(dfreco, self.v_bitvar,
-                                                         self.b_mcbkg), dtype=int)
+                                                         self.b_mcbkg, True), dtype=int)
 
         pickle.dump(dfreco, openfile(self.l_reco[file_index], "wb"), protocol=4)
 
@@ -450,13 +450,13 @@ def read_df(tree, df_base, var):
             dfgen[self.v_isstd] = np.array(tag_bit_df(dfgen, self.v_bitvar,
                                                       self.b_std), dtype=int)
             dfgen[self.v_ismcsignal] = np.array(tag_bit_df(dfgen, self.v_bitvar,
-                                                           self.b_mcsig), dtype=int)
+                                                           self.b_mcsig, True), dtype=int)
             dfgen[self.v_ismcprompt] = np.array(tag_bit_df(dfgen, self.v_bitvar_origgen,
                                                            self.b_mcsigprompt), dtype=int)
             dfgen[self.v_ismcfd] = np.array(tag_bit_df(dfgen, self.v_bitvar_origgen,
                                                        self.b_mcsigfd), dtype=int)
             dfgen[self.v_ismcbkg] = np.array(tag_bit_df(dfgen, self.v_bitvar,
-                                                        self.b_mcbkg), dtype=int)
+                                                        self.b_mcbkg, True), dtype=int)
             dfgen = dfgen.reset_index(drop=True)
 
             if dfjetgen is not None: