From 741e4bbe6adc287175d7a460b87262f3d6e57c57 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 17 Dec 2024 11:17:13 +0900
Subject: [PATCH] add table and charge filters

---
 .../icecube/utilities/i3_filters.py           | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/src/graphnet/data/extractors/icecube/utilities/i3_filters.py b/src/graphnet/data/extractors/icecube/utilities/i3_filters.py
index dfbfd06c0..4b4534adb 100644
--- a/src/graphnet/data/extractors/icecube/utilities/i3_filters.py
+++ b/src/graphnet/data/extractors/icecube/utilities/i3_filters.py
@@ -155,3 +155,66 @@ def _keep_frame(self, frame: "icetray.I3Frame") -> bool:
                 "FilterMask filters will not be applied."
             )
             return True
+
+
+class TableFilter(I3Filter):
+    """A filter that checks if a table is present in the frame."""
+
+    def __init__(self, table_name: str):
+        """Initialize TableFilter.
+
+        Args:
+            table_name: str
+                The name of the table to check for.
+        """
+        self._table_name = table_name
+
+    def _keep_frame(self, frame: "icetray.I3Frame") -> bool:
+        """Check that the frame has a table.
+
+        Args:
+            frame: I3-frame
+                The I3-frame to check.
+        """
+        return frame.Has(self._table_name)
+
+
+class ChargeFilter(I3Filter):
+    """A filter that checks the recorded charge and requires a lower limit.
+
+    This also requires that the charge table is present in the frame.
+    """
+
+    def __init__(
+        self, min_charge: float, table_name: str = "Homogenized_QTot"
+    ):
+        """Initialize ChargeFilter.
+
+        Args:
+            min_charge: The minimum charge required to keep the frame.
+            table_name: The name of the charge table.
+        """
+        self._min_charge = min_charge
+        self._table_name = table_name
+
+    def _keep_frame(self, frame: "icetray.I3Frame") -> bool:
+        """Check that the frame has a charge and that it is within the limits.
+
+        Args:
+            frame: I3-frame
+        """
+        if frame.Has(self._table_name):
+            try:
+                charge = frame[self._table_name].value
+                return charge >= self._min_charge
+            except AttributeError:
+                try:
+                    charge = frame[self._table_name].charge
+                    return charge >= self._min_charge
+                except AttributeError:
+                    self.warning_once(
+                        f"Charge table {self._table_name} has no attribute\
+                          'value' or 'charge'."
+                    )
+                    return False
+        return False