1.3.1 remove ps and brem (#2)

jpata · May 2, 2022 · 05ac2ef · 05ac2ef
1 parent 6b63284
commit 05ac2ef
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 14 deletions.
diff --git a/heptfds/cms_pf/qcd.py b/heptfds/cms_pf/qcd.py
@@ -19,14 +19,15 @@
 _CITATION = """
 """
 
-PADDED_NUM_ELEM_SIZE = 12000
+PADDED_NUM_ELEM_SIZE = 6400
 
 class CmsPfQcd(tfds.core.GeneratorBasedBuilder):
     """DatasetBuilder for cms_pf dataset."""
 
-    VERSION = tfds.core.Version("1.3.0")
+    VERSION = tfds.core.Version("1.3.1")
     RELEASE_NOTES = {
         "1.3.0": "12_2_0_pre2 generation with updated caloparticle/trackingparticle",
+        "1.3.1": "Remove PS again"
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     FIXME

diff --git a/heptfds/cms_pf/ttbar.py b/heptfds/cms_pf/ttbar.py
@@ -19,17 +19,18 @@
 _CITATION = """
 """
 
-PADDED_NUM_ELEM_SIZE = 12000
+PADDED_NUM_ELEM_SIZE = 6400
 
 class CmsPfTtbar(tfds.core.GeneratorBasedBuilder):
     """DatasetBuilder for cms_pf dataset."""
 
-    VERSION = tfds.core.Version("1.3.0")
+    VERSION = tfds.core.Version("1.3.1")
     RELEASE_NOTES = {
         "1.0.0": "Initial release.",
         "1.1.0": "Add muon type, fix electron GSF association",
         "1.2.0": "12_1_0_pre3 generation, add corrected energy, cluster flags, 20k events",
         "1.3.0": "12_2_0_pre2 generation with updated caloparticle/trackingparticle",
+        "1.3.1": "Remove PS again"
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     mkdir -p data

diff --git a/heptfds/cms_pf/ztt.py b/heptfds/cms_pf/ztt.py
@@ -19,17 +19,15 @@
 _CITATION = """
 """
 
-PADDED_NUM_ELEM_SIZE = 12000
+PADDED_NUM_ELEM_SIZE = 6400
 
 class CmsPfZtt(tfds.core.GeneratorBasedBuilder):
     """DatasetBuilder for cms_pf dataset."""
 
-    VERSION = tfds.core.Version("1.3.0")
+    VERSION = tfds.core.Version("1.3.1")
     RELEASE_NOTES = {
-        "1.0.0": "Initial release.",
-        "1.1.0": "Add muon type, fix electron GSF association",
-        "1.2.0": "12_1_0_pre3 generation, add corrected energy, cluster flags, 20k events",
         "1.3.0": "12_2_0_pre2 generation with updated caloparticle/trackingparticle",
+        "1.3.1": "Remove PS again"
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     mkdir -p data

diff --git a/heptfds/cms_utils.py b/heptfds/cms_utils.py
@@ -52,12 +52,12 @@ def prepare_data_cms(fn, padded_num_elem_size):
         ygen = event["ygen"]
         ycand = event["ycand"]
 
-        # remove PS from inputs, they don't seem to be very useful
-        #msk_ps = (Xelem["typ"] == 2) | (Xelem["typ"] == 3)
+        # remove PS and BREM from inputs
+        msk_ps = (Xelem["typ"] == 2) | (Xelem["typ"] == 3) | (Xelem["typ"] == 7)
 
-        Xelem = Xelem
-        ygen = ygen
-        ycand = ycand
+        Xelem = Xelem[~msk_ps]
+        ygen = ygen[~msk_ps]
+        ycand = ycand[~msk_ps]
 
         Xelem = append_fields(
             Xelem, "typ_idx", np.array([ELEM_LABELS_CMS.index(int(i)) for i in Xelem["typ"]], dtype=np.float32)