From 05ac2ef1cd1f8a3715954c8f3ccfddd8e5d5826d Mon Sep 17 00:00:00 2001
From: Joosep Pata <joosep.pata@gmail.com>
Date: Mon, 2 May 2022 14:45:43 +0300
Subject: [PATCH] 1.3.1 remove ps and brem (#2)

---
 heptfds/cms_pf/qcd.py   |  5 +++--
 heptfds/cms_pf/ttbar.py |  5 +++--
 heptfds/cms_pf/ztt.py   |  8 +++-----
 heptfds/cms_utils.py    | 10 +++++-----
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/heptfds/cms_pf/qcd.py b/heptfds/cms_pf/qcd.py
index 7371f59..7f1f2ce 100644
--- 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
index 7225370..ae6c9a2 100644
--- 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
index 9605257..b3e7817 100644
--- 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
index 35990da..b3b3cd7 100644
--- 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)