Add IceTop support to I3TruthExtractor

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

@@ -32,6 +32,7 @@ def __init__(
         mctree: Optional[str] = "I3MCTree",
         extend_boundary: Optional[float] = 0.0,
         exclude: list = [None],
+        ice_top: bool = False,
     ):
         """Construct I3TruthExtractor.
 
@@ -45,6 +46,7 @@ def __init__(
             extend_boundary: Distance to extend the convex hull of the detector
                 for defining starting events.
             exclude: List of keys to exclude from the extracted data.
+            ice_top: Whether the data being extracted is IceTop data. Defaults to False.
         """
         # Base class constructor
         super().__init__(name, exclude=exclude)
@@ -89,6 +91,7 @@ def __init__(
 
         self._extend_boundary = extend_boundary
         self._mctree = mctree
+        self._ice_top = ice_top
 
     def set_gcd(self, i3_file: str, gcd_file: Optional[str] = None) -> None:
         """Extract GFrame and CFrame from i3/gcd-file pair.
@@ -130,8 +133,8 @@ def __call__(
         self, frame: "icetray.I3Frame", padding_value: Any = -1
     ) -> Dict[str, Any]:
         """Extract truth-level information."""
-        is_mc = frame_is_montecarlo(frame, self._mctree)
-        is_noise = frame_is_noise(frame, self._mctree)
+        is_mc = frame_is_montecarlo(frame, self._mctree, self._ice_top)
+        is_noise = frame_is_noise(frame, self._mctree, self._ice_top)
         sim_type = self._find_data_type(is_mc, self._i3_file, frame)
 
         output = {
@@ -172,10 +175,11 @@ def __call__(
         # for example I3RecoPulseSeriesMap,  etc.
         # At low levels i3 files contain several other P frame splits
         # (e.g NullSplit). We remove those here.
-        if frame["I3EventHeader"].sub_event_stream not in [
-            "InIceSplit",
-            "Final",
-        ]:
+        allowed_streams = ["InIceSplit", "Final"]
+        if self._ice_top:
+            allowed_streams.append("ice_top")
+
+        if frame["I3EventHeader"].sub_event_stream not in allowed_streams:
             return output
 
         if "FilterMask" in frame:
@@ -250,8 +254,12 @@ def __call__(
                     "pid": MCInIcePrimary.pdg_encoding,
                     "interaction_type": interaction_type,
                     "elasticity": elasticity,
-                    "dbang_decay_length": self._extract_dbang_decay_length(
-                        frame, padding_value
+                    "dbang_decay_length": (
+                        padding_value
+                        if self._ice_top
+                        else self._extract_dbang_decay_length(
+                            frame, padding_value
+                        )
                     ),
                     "energy_track": energy_track,
                     "energy_cascade": energy_cascade,
@@ -415,17 +423,20 @@ def _get_primary_particle_interaction_type_and_elasticity(
                 or 0 (neither); and the elasticity in the range ]0,1[.
         """
         if sim_type != "noise":
-            try:
-                MCInIcePrimary = frame["MCInIcePrimary"]
-            except KeyError:
-                MCInIcePrimary = frame[self._mctree][0]
-            if (
-                MCInIcePrimary.energy != MCInIcePrimary.energy
-            ):  # This is a nan check. Only happens for some muons
-                # where second item in MCTree is primary. Weird!
-                MCInIcePrimary = frame[self._mctree][1]
-                # For some strange reason the second entry is identical in
-                # all variables and has no nans (always muon)
+            if self._ice_top:
+                MCInIcePrimary = frame["MCPrimary"]  # IceTop CORSIKA primary
+            else:
+                try:
+                    MCInIcePrimary = frame["MCInIcePrimary"]
+                except KeyError:
+                    MCInIcePrimary = frame[self._mctree][0]
+                if (
+                    MCInIcePrimary.energy != MCInIcePrimary.energy
+                ):  # This is a nan check. Only happens for some muons
+                    # where second item in MCTree is primary. Weird!
+                    MCInIcePrimary = frame[self._mctree][1]
+                    # For some strange reason the second entry is identical in
+                    # all variables and has no nans (always muon)
         else:
             MCInIcePrimary = None
 
@@ -472,6 +483,11 @@ def _get_primary_track_energy_and_inelasticity(
             Tuple containing the energy of tracks from primary, and the
                 corresponding inelasticity.
         """
+        if self._ice_top and self._mctree not in frame:
+            raise RuntimeError(
+                f"Expected MCTree key '{self._mctree}' not found in frame for IceTop data."
+            )
+
         mc_tree = frame[self._mctree]
         primary = mc_tree.primaries[0]
         daughters = mc_tree.get_daughters(primary)
@@ -509,7 +525,7 @@ def _find_data_type(
             sim_type = "data"
         elif "muon" in input_file:
             sim_type = "muongun"
-        elif "corsika" in input_file:
+        elif "corsika" in input_file or self._ice_top:
             sim_type = "corsika"
         elif "genie" in input_file or "nu" in input_file.lower():
             sim_type = "genie"

src/graphnet/data/extractors/icecube/utilities/frames.py

@@ -12,16 +12,29 @@
 
 
 def frame_is_montecarlo(
-    frame: "icetray.I3Frame", mctree: Optional[str] = "I3MCTree"
+    frame: "icetray.I3Frame",
+    mctree: Optional[str] = "I3MCTree",
+    ice_top: bool = False,
 ) -> bool:
     """Check whether `frame` is from Monte Carlo simulation."""
+    if ice_top:
+        return "MCPrimary" in frame
     return ("MCInIcePrimary" in frame) or (mctree in frame)
 
 
 def frame_is_noise(
-    frame: "icetray.I3Frame", mctree: Optional[str] = "I3MCTree"
+    frame: "icetray.I3Frame",
+    mctree: Optional[str] = "I3MCTree",
+    ice_top: bool = False,
 ) -> bool:
     """Check whether `frame` is from noise."""
+    if ice_top:
+        try:
+            frame["MCPrimary"].energy
+            return False
+        except:  # noqa: E722
+            return True
+
     try:
         frame[mctree][0].energy
         return False