Skip to content

Commit

Permalink
add table and charge filters
Browse files Browse the repository at this point in the history
  • Loading branch information
Aske-Rosted committed Dec 17, 2024
1 parent 5fb93f6 commit 741e4bb
Showing 1 changed file with 63 additions and 0 deletions.
63 changes: 63 additions & 0 deletions src/graphnet/data/extractors/icecube/utilities/i3_filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 741e4bb

Please sign in to comment.