Merge pull request #801 from Aske-Rosted/i3filtermapextractor

I3filtermapextractor

Author: Aske-Rosted

src/graphnet/data/extractors/icecube/__init__.py

@@ -18,3 +18,4 @@
 from .i3pisaextractor import I3PISAExtractor
 from .i3ntmuonlabelsextractor import I3NTMuonLabelExtractor
 from .i3quesoextractor import I3QUESOExtractor
+from .i3filtermapextractor import I3FilterMapExtractor

src/graphnet/data/extractors/icecube/i3filtermapextractor.py

@@ -0,0 +1,37 @@
+"""I3Extractor for extracting the boolean condition of the I3FilterMask."""
+
+from typing import TYPE_CHECKING, Dict
+
+from graphnet.data.extractors.icecube import I3Extractor
+
+if TYPE_CHECKING:
+    from icecube import icetray  # pyright: reportMissingImports=false
+
+
+class I3FilterMapExtractor(I3Extractor):
+    """Class for extracting I3FilterMap properties.
+
+    This class extracts the boolean condition of the I3FilterMask from the
+    I3FilterMap in the frame.
+    """
+
+    def __init__(
+        self,
+        key: str = "FilterMask",
+        extractor_name: str = "I3FilterMap",
+        exclude: list = [None],
+    ) -> None:
+        """Construct I3FilterMapExtractor."""
+        # Base class constructor
+        super().__init__(extractor_name=extractor_name, exclude=exclude)
+        self._key = key
+
+    def __call__(self, frame: "icetray.I3Frame") -> Dict[str, bool]:
+        """Extract I3FilterMap properties from I3FilterMap in frame."""
+        output = {}
+        if self._key not in frame:
+            raise KeyError(f"Key {self._key} not found in frame.")
+        items = frame[self._key].items()
+        for item in items:
+            output[item[0]] = item[1].condition_passed
+        return output