From 6b421d3c7e90b3367421798c25f1bc1be8d9f964 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 13 Jan 2026 01:05:37 -0600
Subject: [PATCH 01/19] revert calor fix

---
 .../data/extractors/icecube/i3calorimetry.py  | 382 +++++++-----------
 1 file changed, 147 insertions(+), 235 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 09f98e73b..f114b05b9 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -1,6 +1,6 @@
 """Extract all the visible particles entering the volume."""
 
-from typing import Dict, Any, TYPE_CHECKING, Tuple, List
+from typing import Dict, Any, TYPE_CHECKING, Tuple
 
 from .utilities.gcd_hull import GCD_hull
 from .i3extractor import I3Extractor
@@ -14,6 +14,7 @@
         icetray,
         dataclasses,
         MuonGun,
+        simclasses,
     )  # pyright: reportMissingImports=false
 
 
@@ -31,8 +32,6 @@ def __init__(
         mmctracklist: str = "MMCTrackList",
         extractor_name: str = "I3Calorimetry",
         daughters: bool = False,
-        highest_energy_primary: bool = False,
-        cascade_deposited_only: bool = True,
         **kwargs: Any,
     ) -> None:
         """Create a ConvexHull object from the GCD file.
@@ -42,37 +41,14 @@ def __init__(
         mctree: Name of the I3MCTree in the frame.
         mmctracklist: Name of the MMCTrackList in the frame.
         extractor_name: Name of the extractor.
-        daughters: If True, only calculate energies for particles
-            that are daughters of the primary.
-        highest_energy_primary: If True, takes into account only the
-            primary with the highest energy.
-            NOTE: Only makes a difference if daughters is False
-                and the event is not a Corsika event.
-        cascade_deposited_only: If True, consider only energies from
-            cascades that are marked as visible. If False the total
-            energy of a cascade is counted.
-
-        Variable explanation:
-        - e_entrance_track: Total energy of tracks entering the hull.
-        - e_deposited_track: Total energy deposited by tracks in the hull.
-        - e_cascade: Total energy of cascade particles in the hull.
-        - e_visible: Total energy of particles entering the hull.
-            NOTE: if daughters is True, this is the total visible energy
-            of daughter particles of the primary particles. If this is 0
-            that means that all the light in the detector comes from
-            particles that are daughters of coincident primaries.
-        - fraction_primary: Fraction of `e_visible` compared to
-            the primary energy.
-        - fraction_cascade: Fraction of the total energy that is
-            deposited by cascade particles compared to the total energy.
+        daughters: If True, only consider particles that are
+            daughters of the primary.
         """
         # Member variable(s)
         self.hull = hull
         self.mctree = mctree
         self.mmctracklist = mmctracklist
         self.daughters = daughters
-        self.highest_energy_primary = highest_energy_primary
-        self.cascade_deposited_only = cascade_deposited_only
         # Base class constructor
         super().__init__(extractor_name=extractor_name, **kwargs)
 
@@ -81,71 +57,23 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         output = {}
         if self.frame_contains_info(frame):
 
-            primaries = self.get_primaries(
-                frame,
-                self.daughters,
-                self.highest_energy_primary,
+            e_entrance_track, e_deposited_track = self.total_track_energy(
+                frame
             )
 
-            if not len(primaries) == 0:
-
-                MMCTrackList = frame[self.mmctracklist]
-                # Filter tracks that are not daughters of the desired
-                if self.daughters:
-                    temp_MMCTrackList = []
-                    for track in MMCTrackList:
-                        for p in primaries:
-                            if frame[self.mctree].is_in_subtree(
-                                p.id, track.GetI3Particle().id
-                            ):
-                                temp_MMCTrackList.append(track)
-                                break
-                    MMCTrackList = temp_MMCTrackList
-
-                # Create a lookup dict for the tracks
-                track_lookup = {}
-                for track in MuonGun.Track.harvest(
-                    frame[self.mctree], MMCTrackList
-                ):
-                    track_lookup[track.id] = track
-
-                e_cascade, e_dep_track, e_ent_track = self.get_energies(
-                    frame, primaries, track_lookup
-                )
-
-                primary_energy = sum([p.energy for p in primaries])
-            else:
-                e_ent_track = np.nan
-                e_dep_track = np.nan
-                e_cascade = np.nan
-                primary_energy = np.nan
-
-            e_total = e_ent_track + e_cascade
-
-            # In case all particles are considered and
-            # there is no energy deposited in the hull,
-            # we warn the user.
-            if all(
-                (
-                    not self.daughters,
-                    not self.highest_energy_primary,
-                    e_total == 0,
-                )
-            ):
-                self.warning(
-                    "No energy deposited in the hull, "
-                    "Think about increasing the padding."
-                    f"\nCurrent padding: {self.hull.padding}"
-                    f"\nTotal energy: {e_total}"
-                    f"\nTrack energy: {e_ent_track}"
-                    f"\nCascade energy: {e_cascade}"
-                    f"\nEvent header: {frame['I3EventHeader']}"
-                )
+            e_deposited_cascade = self.total_cascade_energy(frame)
 
-            # Check only in the case that there were primaries
-            if not len(primaries) == 0 and (not np.isnan(e_total)):
+            primary_energy = sum(
+                [
+                    p.energy
+                    for p in self.check_primary_energy(
+                        frame, self.get_primaries(frame, self.daughters)
+                    )
+                ]
+            )
+            e_total = e_entrance_track + e_deposited_cascade
 
-                # total energy should always be less than the primary energy
+            if self.daughters:
                 assert e_total <= (
                     primary_energy * (1 + 1e-6)
                 ), "Total energy on entrance is greater than primary energy\
@@ -156,34 +84,26 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     {}".format(
                     e_total,
                     primary_energy,
-                    e_ent_track,
-                    e_cascade,
+                    e_entrance_track,
+                    e_deposited_cascade,
                     frame["I3EventHeader"],
                 )
 
-                assert (
-                    primary_energy > 0
-                ), "Primary energy is 0, this should not happen.\
-                    \nTotal energy: {}\
-                    \nTrack energy: {}\
-                    \nCascade energy: {}\
-                    {}".format(
-                    e_total,
-                    e_ent_track,
-                    e_cascade,
-                    frame["I3EventHeader"],
-                )
-            fraction_primary = e_total / primary_energy
-
             cascade_fraction = None
             if e_total > 0:
-                cascade_fraction = e_cascade / e_total
+                cascade_fraction = e_deposited_cascade / e_total
 
+            if primary_energy > 0:
+                fraction_primary = e_total / primary_energy
+            else:
+                fraction_primary = None
             output.update(
                 {
-                    "e_entrance_track_" + self._extractor_name: e_ent_track,
-                    "e_deposited_track_" + self._extractor_name: e_dep_track,
-                    "e_cascade_" + self._extractor_name: e_cascade,
+                    "e_entrance_track_"
+                    + self._extractor_name: e_entrance_track,
+                    "e_deposited_track_"
+                    + self._extractor_name: e_deposited_track,
+                    "e_cascade_" + self._extractor_name: e_deposited_cascade,
                     "e_visible_" + self._extractor_name: e_total,
                     "fraction_primary_"
                     + self._extractor_name: fraction_primary,
@@ -195,144 +115,136 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         output = {k: v for k, v in output.items() if k not in self._exclude}
         return output
 
-    def get_energies(
+    def frame_contains_info(self, frame: "icetray.I3Frame") -> bool:
+        """Check if the frame contains the necessary information."""
+        return self.mctree in frame and self.mmctracklist in frame
+
+    def total_track_energy(
+        self, frame: "icetray.I3Frame"
+    ) -> Tuple[float, float]:
+        """Get the total energy of track particles on entrance."""
+        e_entrance = 0
+        e_deposited = 0
+        primaries = self.get_primaries(frame, self.daughters)
+        primaries = self.check_primary_energy(frame, primaries)
+
+        MMCTrackList = frame[self.mmctracklist]
+        if self.daughters:
+            MMCTrackList = [
+                track
+                for track in MMCTrackList
+                if frame[self.mctree].get_primary(track.GetI3Particle())
+                in primaries
+            ]
+            MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList)
+
+        for track in MuonGun.Track.harvest(frame[self.mctree], MMCTrackList):
+            assert track.is_track, "Track is not a track"
+
+            # Find distance to entrance and exit from sampling volume
+            intersections = self.hull.surface.intersection(
+                track.pos, track.dir
+            )
+            # Get the corresponding energies
+
+            e0 = track.get_energy(intersections.first)
+            e1 = track.get_energy(intersections.second)
+
+            # Accumulate
+            e_deposited += e0 - e1
+            e_entrance += e0
+            if self.daughters:
+                assert e_entrance <= sum(
+                    [p.energy for p in primaries]
+                ), "Energy on entrance is greater than primary energy"
+                assert e_deposited <= sum(
+                    [p.energy for p in primaries]
+                ), "Energy deposited is greater than primary energy"
+        return e_entrance, e_deposited
+
+    def total_cascade_energy(
         self,
         frame: "icetray.I3Frame",
-        particles: List["dataclasses.I3Particle"],
-        track_lookup: Dict["icetray.I3ParticleID", "icetray.I3Particle"],
-    ) -> Tuple[float, float, float]:
+    ) -> float:
         """Get the total energy of cascade particles on entrance."""
-        e_cascade = 0
-        e_dep_track = 0
-        e_ent_track = 0
+        e_deposited = 0
+
+        particles = self.get_primaries(frame, self.daughters)
+
+        particles = np.array([p for p in particles if (not p.is_track)])
 
         if len(particles) == 0:
-            return e_cascade, e_dep_track, e_ent_track
+            return e_deposited
+
+        pos, direc, length = np.asarray(
+            [
+                [
+                    np.array(p.pos),
+                    np.array([p.dir.x, p.dir.y, p.dir.z]),
+                    p.length,
+                ]
+                for p in particles
+            ],
+            dtype=object,
+        ).T
+
+        length = length.astype(float)
+
+        # replace length nan with 0
+        length[np.isnan(length)] = 0
+        pos = pos + direc * length
+        pos = np.stack(pos)
+
+        in_volume = self.hull.point_in_hull(pos)
+        particles = particles[in_volume]
 
         for particle in particles:
-            length = particle.length
-            if length != length:
-                length = 0
-            # If the particle is a track in the MMCTrackList take the
-            # energy at the entrance and exit of the hull.
-            # NOTE: We do not consider daughters of tracks,
-            # because they are already included in the track energy.
-            if particle.is_track & (particle.id in track_lookup):
-                track = track_lookup[particle.id]
-
-                # Find distance to entrance and exit from sampling volume
-                intersections = self.hull.surface.intersection(
-                    track.pos, track.dir
-                )
-                # Get the corresponding energies
-                try:
-                    e0 = track.get_energy(intersections.first)
-                    e1 = track.get_energy(intersections.second)
-
-                # Catch MuonGun errors
-                except RuntimeError as e:
-                    if (
-                        "sum of losses is smaller than "
-                        "energy at last checkpoint" in str(e)
-                    ):
-                        hdr = frame["I3EventHeader"]
-                        e.add_note(f"Error in MuonGun track in event {hdr}")
-                        self.warning(f"Skipping bad event {hdr}: {e}")
-                        e0 = np.nan
-                        e1 = np.nan
-                        e_cascade = np.nan
-                        continue  # skip this frame
-                    else:
-                        raise  # re-raise unexpected errors
-
-                e_dep_track += e0 - e1
-                e_ent_track += e0
-            # if the particle is not in the hull, but has daughters,
-            # we add the energies of the daughters.
-            elif not self.is_in_hull(particle):
+            if particle.is_cascade:
+                e_deposited += particle.energy
+            else:
                 daughters = dataclasses.I3MCTree.get_daughters(
                     frame[self.mctree], particle
                 )
-                if len(daughters) == 0:
+
+                pos, direc, length = np.asarray(
+                    [
+                        [
+                            np.array(p.pos),
+                            np.array([p.dir.x, p.dir.y, p.dir.z]),
+                            p.length,
+                        ]
+                        for p in daughters
+                    ],
+                    dtype=object,
+                ).T
+
+                length = length.astype(float)
+
+                # replace length nan with 0
+                length[np.isnan(length)] = 0
+                pos = pos + direc * length
+                pos = np.stack(pos)
+
+                in_volume = self.hull.point_in_hull(pos)
+                daughters = np.array(daughters)[in_volume]
+
+            while len(daughters) > 0:
+                daughter = daughters[0]
+                daughters = daughters[1:]
+                length = daughter.length
+                if daughter.is_track:
                     continue
-                (
-                    e_cascade,
-                    e_dep_track,
-                    e_ent_track,
-                ) = tuple(
-                    np.add(
-                        (e_cascade, e_dep_track, e_ent_track),
-                        self.get_energies(
-                            frame,
-                            daughters,
-                            track_lookup,
-                        ),
-                    )
-                )
-            # If the particle is a cascade in the hull, we add its energy.
-            elif particle.is_cascade:
-                if self.cascade_deposited_only:
-                    # Check wether the cascade is made up of smaller segments
-                    # in this case the shape is dark and we want to count
-                    # the energy of its daughters.
-                    if (
-                        particle.shape
-                        != dataclasses.I3Particle.ParticleShape.Dark
-                    ):
-                        e_cascade += particle.energy
-                    else:
-                        (
-                            e_cascade,
-                            e_dep_track,
-                            e_ent_track,
-                        ) = tuple(
-                            np.add(
-                                (e_cascade, e_dep_track, e_ent_track),
-                                self.get_energies(
-                                    frame,
-                                    dataclasses.I3MCTree.get_daughters(
-                                        frame[self.mctree], particle
-                                    ),
-                                    track_lookup,
-                                ),
-                            )
-                        )
+                if length == np.nan:
+                    length = 0
+                if daughter.is_cascade and daughter.shape != "Dark":
+                    e_deposited += daughter.energy
                 else:
-                    # In this case we consider the total cascade
-                    # energy and therefore do not look at the daughters
-                    e_cascade += particle.energy
-            # The particle is in the hull and not a track in the MMCTrackList,
-            # or a cascade, so we look at its daughters.
-            # Could be a NuMu interacting within the hull.
-            else:
-                (
-                    e_cascade,
-                    e_dep_track,
-                    e_ent_track,
-                ) = tuple(
-                    np.add(
-                        (e_cascade, e_dep_track, e_ent_track),
-                        self.get_energies(
-                            frame,
+                    daughters = np.concatenate(
+                        [
+                            daughters,
                             dataclasses.I3MCTree.get_daughters(
-                                frame[self.mctree], particle
+                                frame[self.mctree], daughter
                             ),
-                            track_lookup,
-                        ),
+                        ]
                     )
-                )
-
-        return e_cascade, e_dep_track, e_ent_track
-
-    def frame_contains_info(self, frame: "icetray.I3Frame") -> bool:
-        """Check if the frame contains the necessary information."""
-        return self.mctree in frame and self.mmctracklist in frame
-
-    def is_in_hull(self, particle: "dataclasses.I3Particle") -> bool:
-        """Check if a particle is in the hull."""
-        pos = np.array(particle.pos)
-        direc = np.array([particle.dir.x, particle.dir.y, particle.dir.z])
-        length = particle.length if particle.length is not None else 0
-        pos = pos + direc * length
-
-        return self.hull.point_in_hull(pos)
+        return e_deposited

From b9b565ccdd25d0e8f09a4d3b6786c7b9eb4254a6 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 13 Jan 2026 01:17:15 -0600
Subject: [PATCH 02/19] boolean logic fix

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index f114b05b9..ced1afc39 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -236,7 +236,9 @@ def total_cascade_energy(
                     continue
                 if length == np.nan:
                     length = 0
-                if daughter.is_cascade and daughter.shape != "Dark":
+                if daughter.is_cascade and (
+                    daughter.shape != dataclasses.I3Particle.ParticleShape.Dark
+                ):
                     e_deposited += daughter.energy
                 else:
                     daughters = np.concatenate(

From 858c1809ff85f28c5a258b515300d9135b54955a Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 13 Jan 2026 18:36:27 -0600
Subject: [PATCH 03/19] alternative fix

---
 .../data/extractors/icecube/i3calorimetry.py  | 36 +++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index ced1afc39..88de3d325 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -8,6 +8,7 @@
 import numpy as np
 
 from graphnet.utilities.imports import has_icecube_package
+from copy import deepcopy
 
 if has_icecube_package() or TYPE_CHECKING:
     from icecube import (
@@ -32,6 +33,7 @@ def __init__(
         mmctracklist: str = "MMCTrackList",
         extractor_name: str = "I3Calorimetry",
         daughters: bool = False,
+        highest_energy_primary: bool = False,
         **kwargs: Any,
     ) -> None:
         """Create a ConvexHull object from the GCD file.
@@ -49,12 +51,15 @@ def __init__(
         self.mctree = mctree
         self.mmctracklist = mmctracklist
         self.daughters = daughters
+        self.highest_energy_primary = highest_energy_primary
         # Base class constructor
         super().__init__(extractor_name=extractor_name, **kwargs)
 
     def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         """Extract all the visible particles entering the volume."""
         output = {}
+        # copy the original mctree because we will be modifying it
+        tree_copy = deepcopy(frame[self.mctree])
         if self.frame_contains_info(frame):
 
             e_entrance_track, e_deposited_track = self.total_track_energy(
@@ -67,7 +72,12 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 [
                     p.energy
                     for p in self.check_primary_energy(
-                        frame, self.get_primaries(frame, self.daughters)
+                        frame,
+                        self.get_primaries(
+                            frame,
+                            self.daughters,
+                            highest_energy_primary=self.highest_energy_primary,
+                        ),
                     )
                 ]
             )
@@ -113,6 +123,9 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             )
 
         output = {k: v for k, v in output.items() if k not in self._exclude}
+        # restore original mctree
+        frame.Delete(self.mctree)
+        frame[self.mctree] = tree_copy
         return output
 
     def frame_contains_info(self, frame: "icetray.I3Frame") -> bool:
@@ -138,8 +151,12 @@ def total_track_energy(
             ]
             MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList)
 
-        for track in MuonGun.Track.harvest(frame[self.mctree], MMCTrackList):
-            assert track.is_track, "Track is not a track"
+        track_list = MuonGun.Track.harvest(frame[self.mctree], MMCTrackList)
+        track_ids = np.array([track.id for track in track_list])
+
+        total_tracks = len(track_list)
+        while len(track_list) > 0:
+            track = track_list[0]
 
             # Find distance to entrance and exit from sampling volume
             intersections = self.hull.surface.intersection(
@@ -160,6 +177,19 @@ def total_track_energy(
                 assert e_deposited <= sum(
                     [p.energy for p in primaries]
                 ), "Energy deposited is greater than primary energy"
+            # erase particle and children
+            exclude_ids = set(
+                [c.id for c in frame[self.mctree].children(track.id)]
+            )
+            exclude_ids.add(track.id)
+            frame[self.mctree].erase_children(track.id)
+            frame[self.mctree].erase(track.id)
+            track_mask = [tid not in exclude_ids for tid in track_ids]
+            track_list = list(np.array(track_list)[track_mask])
+            track_ids = track_ids[track_mask]
+            self._logger.debug(
+                f"Remaining tracks: {len(track_list)}/{total_tracks}"
+            )
         return e_entrance, e_deposited
 
     def total_cascade_energy(

From 79adb2bcf531a852b727913ed1e3303db2ae3eef Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 13 Jan 2026 18:59:41 -0600
Subject: [PATCH 04/19] move assert and check intersection

---
 .../data/extractors/icecube/i3calorimetry.py  | 36 ++++++++++++++-----
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 88de3d325..ee3c9409d 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -155,6 +155,7 @@ def total_track_energy(
         track_ids = np.array([track.id for track in track_list])
 
         total_tracks = len(track_list)
+        exclude_ids = set()
         while len(track_list) > 0:
             track = track_list[0]
 
@@ -162,23 +163,32 @@ def total_track_energy(
             intersections = self.hull.surface.intersection(
                 track.pos, track.dir
             )
-            # Get the corresponding energies
 
+            # Check if the track actually enters the volume
+            if not (
+                np.isfinite(intersections.first)
+                and (
+                    intersections.first
+                    < frame[self.mctree].get_particle(track.id).length
+                )
+            ):
+                track_list = track_list[1:]
+                track_ids = track_ids[1:]
+                frame[self.mctree].erase(track.id)
+                self._logger.debug(
+                    f"Remaining tracks: {len(track_list)}/{total_tracks}"
+                )
+                continue
+
+            # Get the corresponding energies
             e0 = track.get_energy(intersections.first)
             e1 = track.get_energy(intersections.second)
 
             # Accumulate
             e_deposited += e0 - e1
             e_entrance += e0
-            if self.daughters:
-                assert e_entrance <= sum(
-                    [p.energy for p in primaries]
-                ), "Energy on entrance is greater than primary energy"
-                assert e_deposited <= sum(
-                    [p.energy for p in primaries]
-                ), "Energy deposited is greater than primary energy"
             # erase particle and children
-            exclude_ids = set(
+            exclude_ids.update(
                 [c.id for c in frame[self.mctree].children(track.id)]
             )
             exclude_ids.add(track.id)
@@ -190,6 +200,14 @@ def total_track_energy(
             self._logger.debug(
                 f"Remaining tracks: {len(track_list)}/{total_tracks}"
             )
+        # Sanity check ensuring no double counting
+        if self.daughters:
+            assert e_entrance <= sum(
+                [p.energy for p in primaries]
+            ), "Energy on entrance is greater than primary energy"
+            assert e_deposited <= sum(
+                [p.energy for p in primaries]
+            ), "Energy deposited is greater than primary energy"
         return e_entrance, e_deposited
 
     def total_cascade_energy(

From 4fdd013b9098e6414f6c162579a2211b39bd3e6f Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Mon, 26 Jan 2026 00:37:35 -0600
Subject: [PATCH 05/19] big remake

---
 .../data/extractors/icecube/i3calorimetry.py  | 197 ++++++++----------
 1 file changed, 85 insertions(+), 112 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index ee3c9409d..6bb9740b7 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -1,6 +1,6 @@
 """Extract all the visible particles entering the volume."""
 
-from typing import Dict, Any, TYPE_CHECKING, Tuple
+from typing import Dict, Any, TYPE_CHECKING, Tuple, Union, List
 
 from .utilities.gcd_hull import GCD_hull
 from .i3extractor import I3Extractor
@@ -45,6 +45,8 @@ def __init__(
         extractor_name: Name of the extractor.
         daughters: If True, only consider particles that are
             daughters of the primary.
+        highest_energy_primary: If True, only consider particles that are
+            daughters of the highest energy primary.
         """
         # Member variable(s)
         self.hull = hull
@@ -62,12 +64,6 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         tree_copy = deepcopy(frame[self.mctree])
         if self.frame_contains_info(frame):
 
-            e_entrance_track, e_deposited_track = self.total_track_energy(
-                frame
-            )
-
-            e_deposited_cascade = self.total_cascade_energy(frame)
-
             primary_energy = sum(
                 [
                     p.energy
@@ -81,8 +77,32 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     )
                 ]
             )
+
+            e_entrance_track, e_deposited_track = self.total_track_energy(
+                frame
+            )
+
+            e_deposited_cascade = self.total_cascade_energy(frame)
+
             e_total = e_entrance_track + e_deposited_cascade
 
+            if all(
+                (
+                    not self.daughters,
+                    not self.highest_energy_primary,
+                    e_total == 0,
+                )
+            ):
+                self.warning(
+                    "No energy deposited in the hull, "
+                    "Think about increasing the padding."
+                    f"\nCurrent padding: {self.hull.padding}"
+                    f"\nTotal energy: {e_total}"
+                    f"\nTrack energy: {e_entrance_track}"
+                    f"\nCascade energy: {e_deposited_cascade}"
+                    f"\nEvent header: {frame['I3EventHeader']}"
+                )
+
             if self.daughters:
                 assert e_total <= (
                     primary_energy * (1 + 1e-6)
@@ -138,7 +158,9 @@ def total_track_energy(
         """Get the total energy of track particles on entrance."""
         e_entrance = 0
         e_deposited = 0
-        primaries = self.get_primaries(frame, self.daughters)
+        primaries = self.get_primaries(
+            frame, self.daughters, self.highest_energy_primary
+        )
         primaries = self.check_primary_energy(frame, primaries)
 
         MMCTrackList = frame[self.mmctracklist]
@@ -151,13 +173,18 @@ def total_track_energy(
             ]
             MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList)
 
-        track_list = MuonGun.Track.harvest(frame[self.mctree], MMCTrackList)
-        track_ids = np.array([track.id for track in track_list])
+        track_list = np.array(
+            MuonGun.Track.harvest(frame[self.mctree], MMCTrackList)
+        )
 
-        total_tracks = len(track_list)
-        exclude_ids = set()
         while len(track_list) > 0:
             track = track_list[0]
+            track_list = track_list[1:]
+            try:
+                particle = frame[self.mctree].get_particle(track.id)
+            except RuntimeError:
+                # If the particle does not exist in the mctree, that means a partcle further up was processed and therefore it should not be counted
+                continue
 
             # Find distance to entrance and exit from sampling volume
             intersections = self.hull.surface.intersection(
@@ -167,17 +194,8 @@ def total_track_energy(
             # Check if the track actually enters the volume
             if not (
                 np.isfinite(intersections.first)
-                and (
-                    intersections.first
-                    < frame[self.mctree].get_particle(track.id).length
-                )
+                and (intersections.first < particle.length)
             ):
-                track_list = track_list[1:]
-                track_ids = track_ids[1:]
-                frame[self.mctree].erase(track.id)
-                self._logger.debug(
-                    f"Remaining tracks: {len(track_list)}/{total_tracks}"
-                )
                 continue
 
             # Get the corresponding energies
@@ -187,19 +205,10 @@ def total_track_energy(
             # Accumulate
             e_deposited += e0 - e1
             e_entrance += e0
-            # erase particle and children
-            exclude_ids.update(
-                [c.id for c in frame[self.mctree].children(track.id)]
-            )
-            exclude_ids.add(track.id)
-            frame[self.mctree].erase_children(track.id)
+            # get descendant ids
+            # erase particle and children from mctree
             frame[self.mctree].erase(track.id)
-            track_mask = [tid not in exclude_ids for tid in track_ids]
-            track_list = list(np.array(track_list)[track_mask])
-            track_ids = track_ids[track_mask]
-            self._logger.debug(
-                f"Remaining tracks: {len(track_list)}/{total_tracks}"
-            )
+
         # Sanity check ensuring no double counting
         if self.daughters:
             assert e_entrance <= sum(
@@ -215,86 +224,50 @@ def total_cascade_energy(
         frame: "icetray.I3Frame",
     ) -> float:
         """Get the total energy of cascade particles on entrance."""
-        e_deposited = 0
-
-        particles = self.get_primaries(frame, self.daughters)
-
-        particles = np.array([p for p in particles if (not p.is_track)])
+        particles = np.array(
+            self.get_primaries(
+                frame, self.daughters, self.highest_energy_primary
+            )
+        )
 
         if len(particles) == 0:
-            return e_deposited
-
-        pos, direc, length = np.asarray(
-            [
-                [
-                    np.array(p.pos),
-                    np.array([p.dir.x, p.dir.y, p.dir.z]),
-                    p.length,
-                ]
-                for p in particles
-            ],
-            dtype=object,
-        ).T
+            return 0.0
+
+        pos_list, direc_list, length_list, cascade_bool, energies = (
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+        while len(particles) > 0:
+            p = particles[0]
+            particles = particles[1:]
+            p_children = dataclasses.I3MCTree.get_daughters(
+                frame[self.mctree], p
+            )
+            if len(p_children) > 0:
+                particles = np.concatenate((particles, p_children))
+                continue
+            elif (
+                p.is_track
+                or p.shape == dataclasses.I3Particle.ParticleShape.Dark
+            ):
+                continue
 
-        length = length.astype(float)
+            pos_list.append(np.array(p.pos))
+            direc_list.append(np.array([p.dir.x, p.dir.y, p.dir.z]))
+            length_list.append(p.length)
+            cascade_bool.append(p.is_cascade)
+            energies.append(p.energy)
 
-        # replace length nan with 0
+        length = np.array(length_list).astype(float)
         length[np.isnan(length)] = 0
-        pos = pos + direc * length
-        pos = np.stack(pos)
-
-        in_volume = self.hull.point_in_hull(pos)
-        particles = particles[in_volume]
-
-        for particle in particles:
-            if particle.is_cascade:
-                e_deposited += particle.energy
-            else:
-                daughters = dataclasses.I3MCTree.get_daughters(
-                    frame[self.mctree], particle
-                )
-
-                pos, direc, length = np.asarray(
-                    [
-                        [
-                            np.array(p.pos),
-                            np.array([p.dir.x, p.dir.y, p.dir.z]),
-                            p.length,
-                        ]
-                        for p in daughters
-                    ],
-                    dtype=object,
-                ).T
-
-                length = length.astype(float)
-
-                # replace length nan with 0
-                length[np.isnan(length)] = 0
-                pos = pos + direc * length
-                pos = np.stack(pos)
-
-                in_volume = self.hull.point_in_hull(pos)
-                daughters = np.array(daughters)[in_volume]
-
-            while len(daughters) > 0:
-                daughter = daughters[0]
-                daughters = daughters[1:]
-                length = daughter.length
-                if daughter.is_track:
-                    continue
-                if length == np.nan:
-                    length = 0
-                if daughter.is_cascade and (
-                    daughter.shape != dataclasses.I3Particle.ParticleShape.Dark
-                ):
-                    e_deposited += daughter.energy
-                else:
-                    daughters = np.concatenate(
-                        [
-                            daughters,
-                            dataclasses.I3MCTree.get_daughters(
-                                frame[self.mctree], daughter
-                            ),
-                        ]
-                    )
-        return e_deposited
+        pos = np.array(pos_list)
+        direc = np.array(direc_list)
+        cascade_bool = np.array(cascade_bool)
+        energies = np.array(energies)
+        pos = (pos.T + direc.T * length).T
+        in_hull = self.hull.point_in_hull(pos)
+
+        return np.sum(energies[cascade_bool & in_hull])

From 8096569994827eaeee3854a87abc71dc2ff62cb7 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Mon, 26 Jan 2026 03:42:38 -0600
Subject: [PATCH 06/19] speed improvements

---
 .../data/extractors/icecube/i3calorimetry.py  | 29 +++++++++----------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 6bb9740b7..2417a6256 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -9,6 +9,7 @@
 
 from graphnet.utilities.imports import has_icecube_package
 from copy import deepcopy
+from collections import deque
 
 if has_icecube_package() or TYPE_CHECKING:
     from icecube import (
@@ -18,6 +19,8 @@
         simclasses,
     )  # pyright: reportMissingImports=false
 
+DARK = dataclasses.I3Particle.ParticleShape.Dark
+
 
 class I3Calorimetry(I3Extractor):
     """Event level energy labeling for IceCube data.
@@ -224,7 +227,7 @@ def total_cascade_energy(
         frame: "icetray.I3Frame",
     ) -> float:
         """Get the total energy of cascade particles on entrance."""
-        particles = np.array(
+        particles = deque(
             self.get_primaries(
                 frame, self.daughters, self.highest_energy_primary
             )
@@ -240,31 +243,27 @@ def total_cascade_energy(
             [],
             [],
         )
+
+        mctree = frame[self.mctree]
         while len(particles) > 0:
-            p = particles[0]
-            particles = particles[1:]
-            p_children = dataclasses.I3MCTree.get_daughters(
-                frame[self.mctree], p
-            )
+            p = particles.popleft()
+            p_children = mctree.get_daughters(p)
             if len(p_children) > 0:
-                particles = np.concatenate((particles, p_children))
+                particles.extend(p_children)
                 continue
-            elif (
-                p.is_track
-                or p.shape == dataclasses.I3Particle.ParticleShape.Dark
-            ):
+            if p.is_track or p.shape == DARK:
                 continue
 
-            pos_list.append(np.array(p.pos))
-            direc_list.append(np.array([p.dir.x, p.dir.y, p.dir.z]))
+            pos_list.append([p.pos.x, p.pos.y, p.pos.z])
+            direc_list.append([p.dir.x, p.dir.y, p.dir.z])
             length_list.append(p.length)
             cascade_bool.append(p.is_cascade)
             energies.append(p.energy)
 
         length = np.array(length_list).astype(float)
         length[np.isnan(length)] = 0
-        pos = np.array(pos_list)
-        direc = np.array(direc_list)
+        pos = np.asarray(pos_list)
+        direc = np.asarray(direc_list)
         cascade_bool = np.array(cascade_bool)
         energies = np.array(energies)
         pos = (pos.T + direc.T * length).T

From 14831e8426b4909ac64ac27a8e7358e707b52751 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 27 Jan 2026 00:15:49 -0600
Subject: [PATCH 07/19] move DARK instantiation

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 2417a6256..943a22d49 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -19,7 +19,7 @@
         simclasses,
     )  # pyright: reportMissingImports=false
 
-DARK = dataclasses.I3Particle.ParticleShape.Dark
+    DARK = dataclasses.I3Particle.ParticleShape.Dark
 
 
 class I3Calorimetry(I3Extractor):

From f24b6ec2cbe97edb5f9c84c35622047b9044e654 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Sun, 8 Mar 2026 21:12:16 -0500
Subject: [PATCH 08/19] catch runtime errors

---
 .../data/extractors/icecube/i3calorimetry.py  | 50 +++++++++++++++----
 1 file changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 943a22d49..30735821a 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -168,13 +168,26 @@ def total_track_energy(
 
         MMCTrackList = frame[self.mmctracklist]
         if self.daughters:
-            MMCTrackList = [
-                track
-                for track in MMCTrackList
-                if frame[self.mctree].get_primary(track.GetI3Particle())
-                in primaries
-            ]
-            MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList)
+            MMCTrackList_filtered = []
+            for track in MMCTrackList:
+                try:
+                    if (
+                        frame[self.mctree].get_primary(track.GetI3Particle())
+                        in primaries
+                    ):
+                        MMCTrackList_filtered.append(track)
+                except RuntimeError as e:
+                    if "particle not found" in str(e):
+                        # log warning with event header
+                        self.warning(
+                            f"Could not find primary for track {track.GetI3Particle()}"
+                            f" in event {frame['I3EventHeader']}: {e}"
+                        )
+                        # continue to next track
+                    else:
+                        raise e
+
+            MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList_filtered)
 
         track_list = np.array(
             MuonGun.Track.harvest(frame[self.mctree], MMCTrackList)
@@ -202,8 +215,24 @@ def total_track_energy(
                 continue
 
             # Get the corresponding energies
-            e0 = track.get_energy(intersections.first)
-            e1 = track.get_energy(intersections.second)
+            try:
+                e0 = track.get_energy(intersections.first)
+                e1 = track.get_energy(intersections.second)
+
+            except RuntimeError as e:
+                if (
+                    "sum of losses is smaller than "
+                    "energy at last checkpoint" in str(e)
+                ):
+                    hdr = frame["I3EventHeader"]
+                    e.add_note(f"Error in MuonGun track in event {hdr}")
+                    self.warning(
+                        f"Skipping bad track {hdr}: {e}"
+                        f"\nTotal energy of offending particle: {particle.energy}"
+                    )
+                    continue
+                else:
+                    raise
 
             # Accumulate
             e_deposited += e0 - e1
@@ -260,6 +289,9 @@ def total_cascade_energy(
             cascade_bool.append(p.is_cascade)
             energies.append(p.energy)
 
+        if len(energies) == 0:
+            return 0.0
+
         length = np.array(length_list).astype(float)
         length[np.isnan(length)] = 0
         pos = np.asarray(pos_list)

From fe26e91d8e11abe7b46827c37a6e1fbc51281681 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Wed, 8 Apr 2026 20:27:08 -0500
Subject: [PATCH 09/19] allow for mass to kinetic conversion

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 30735821a..f0af73aab 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -107,14 +107,14 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 )
 
             if self.daughters:
-                assert e_total <= (
-                    primary_energy * (1 + 1e-6)
+                assert e_total <= (primary_energy * (1 + 1e-6)) or (
+                    e_total - primary_energy < 0.5
                 ), "Total energy on entrance is greater than primary energy\
                     \nTotal energy: {}\
                     \nPrimary energy: {}\
                     \nTrack energy: {}\
                     \nCascade energy: {}\
-                    {}".format(
+                    {}".format(  # allow for differences due to mass -> kinetic energy conversion and numerical precision
                     e_total,
                     primary_energy,
                     e_entrance_track,

From 0aad12727c54f2ea9496f870b73d31749915d194 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Mon, 1 Jun 2026 23:29:13 -0500
Subject: [PATCH 10/19] Major rework

---
 .../data/extractors/icecube/i3calorimetry.py  | 288 ++++++++++--------
 .../data/extractors/icecube/i3extractor.py    | 103 +++++--
 .../icecube/i3highesteparticleextractor.py    |  15 +-
 3 files changed, 262 insertions(+), 144 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index f0af73aab..5f1cdcc14 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -35,8 +35,8 @@ def __init__(
         mctree: str = "I3MCTree",
         mmctracklist: str = "MMCTrackList",
         extractor_name: str = "I3Calorimetry",
-        daughters: bool = False,
         highest_energy_primary: bool = False,
+        entrance_energy: bool = False,
         **kwargs: Any,
     ) -> None:
         """Create a ConvexHull object from the GCD file.
@@ -46,17 +46,16 @@ def __init__(
         mctree: Name of the I3MCTree in the frame.
         mmctracklist: Name of the MMCTrackList in the frame.
         extractor_name: Name of the extractor.
-        daughters: If True, only consider particles that are
-            daughters of the primary.
         highest_energy_primary: If True, only consider particles that are
             daughters of the highest energy primary.
+        entrance_energy: If True, consider entrance energy.
         """
         # Member variable(s)
         self.hull = hull
         self.mctree = mctree
         self.mmctracklist = mmctracklist
-        self.daughters = daughters
         self.highest_energy_primary = highest_energy_primary
+        self.entrance_energy = entrance_energy
         # Base class constructor
         super().__init__(extractor_name=extractor_name, **kwargs)
 
@@ -64,91 +63,145 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         """Extract all the visible particles entering the volume."""
         output = {}
         # copy the original mctree because we will be modifying it
-        tree_copy = deepcopy(frame[self.mctree])
         if self.frame_contains_info(frame):
-
-            primary_energy = sum(
-                [
-                    p.energy
-                    for p in self.check_primary_energy(
-                        frame,
-                        self.get_primaries(
-                            frame,
-                            self.daughters,
-                            highest_energy_primary=self.highest_energy_primary,
-                        ),
-                    )
-                ]
+            target_tree, bkg_tree = self.split_mc_tree(
+                frame, highest_energy_primary=self.highest_energy_primary
             )
+            if len(target_tree) + len(bkg_tree) != len(frame[self.mctree]):
+                raise ValueError(
+                    "Split mctree has different number of particles than original mctree"
+                )
 
-            e_entrance_track, e_deposited_track = self.total_track_energy(
-                frame
+            # For the target we consider either all neutrino primary products or only the highest energy primary of the neutrino depending on the flag.
+            target_primaries = self.get_primaries(
+                target_tree,
+                daughters=True,
+                highest_energy_primary=self.highest_energy_primary,
             )
+            target_primaries = self.check_primary_energy(
+                target_tree, target_primaries
+            )
+            # For background we consider everything in the background tree.
+            bkg_primaries = self.get_primaries(
+                bkg_tree, daughters=False, highest_energy_primary=False
+            )
+            bkg_primaries = self.check_primary_energy(bkg_tree, bkg_primaries)
 
-            e_deposited_cascade = self.total_cascade_energy(frame)
+            target_primaries_energy = sum([p.energy for p in target_primaries])
+            bkg_primaries_energy = sum([p.energy for p in bkg_primaries])
 
-            e_total = e_entrance_track + e_deposited_cascade
+            e_track_target = self.total_track_energy(
+                frame, target_tree, entrance_energy=self.entrance_energy
+            )
+            # Sanity check ensuring no double counting
 
-            if all(
-                (
-                    not self.daughters,
-                    not self.highest_energy_primary,
-                    e_total == 0,
+            if e_track_target > target_primaries_energy:
+                raise ValueError(
+                    f"Energy deposited in target is greater than primary energy: {e_track_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                 )
-            ):
+            if len(bkg_primaries) > 0:
+                e_track_bkg = self.total_track_energy(
+                    frame, bkg_tree, entrance_energy=self.entrance_energy
+                )
+                if e_track_bkg > bkg_primaries_energy:
+                    raise ValueError(
+                        f"Energy deposited in background is greater than primary energy: {e_track_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
+                    )
+            else:
+                e_track_bkg = 0.0
+
+            e_cascade_target = self.total_cascade_energy(
+                target_tree, target_primaries
+            )
+            if e_cascade_target > target_primaries_energy:
+                raise ValueError(
+                    f"Energy deposited in cascades is greater than primary energy: {e_cascade_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
+                )
+            if len(bkg_primaries) > 0:
+                e_cascade_bkg = self.total_cascade_energy(
+                    bkg_tree, bkg_primaries
+                )
+                if e_cascade_bkg > bkg_primaries_energy:
+                    raise ValueError(
+                        f"Energy deposited in cascades is greater than primary energy: {e_cascade_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
+                    )
+            else:
+                e_cascade_bkg = 0.0
+
+            e_total_target = e_track_target + e_cascade_target
+            e_total_bkg = e_track_bkg + e_cascade_bkg
+
+            e_total = e_total_target + e_total_bkg
+
+            if e_total == 0.0:
                 self.warning(
                     "No energy deposited in the hull, "
                     "Think about increasing the padding."
                     f"\nCurrent padding: {self.hull.padding}"
-                    f"\nTotal energy: {e_total}"
-                    f"\nTrack energy: {e_entrance_track}"
-                    f"\nCascade energy: {e_deposited_cascade}"
                     f"\nEvent header: {frame['I3EventHeader']}"
                 )
 
-            if self.daughters:
-                assert e_total <= (primary_energy * (1 + 1e-6)) or (
-                    e_total - primary_energy < 0.5
-                ), "Total energy on entrance is greater than primary energy\
-                    \nTotal energy: {}\
-                    \nPrimary energy: {}\
-                    \nTrack energy: {}\
-                    \nCascade energy: {}\
-                    {}".format(  # allow for differences due to mass -> kinetic energy conversion and numerical precision
-                    e_total,
-                    primary_energy,
-                    e_entrance_track,
-                    e_deposited_cascade,
-                    frame["I3EventHeader"],
+            if not (
+                e_total_target <= (target_primaries_energy * (1 + 1e-6))
+                or (e_total_target - target_primaries_energy < 0.5)
+            ):
+                raise ValueError(
+                    "Total energy on entrance is greater than primary energy\n"
+                    f"Total energy: {e_total_target}\n"
+                    f"Primary energy: {target_primaries_energy}\n"
+                    f"Track deposited energy: {e_track_target}\n"
+                    f"Cascade deposited energy: {e_cascade_target}\n"
+                    f"{frame['I3EventHeader']}"
                 )
 
-            cascade_fraction = None
+            e_target_fraction = (
+                e_total_target / e_total if e_total > 0 else 0.0
+            )
+            target_cascade_fraction = (
+                e_cascade_target / e_total_target
+                if e_total_target > 0
+                else 0.0
+            )
+
             if e_total > 0:
-                cascade_fraction = e_deposited_cascade / e_total
+                cascade_fraction_tot = (
+                    e_cascade_target + e_cascade_bkg
+                ) / e_total
 
-            if primary_energy > 0:
-                fraction_primary = e_total / primary_energy
+            if target_primaries_energy > 0:
+                fraction_primary = e_total_target / target_primaries_energy
             else:
                 fraction_primary = None
             output.update(
                 {
-                    "e_entrance_track_"
-                    + self._extractor_name: e_entrance_track,
-                    "e_deposited_track_"
-                    + self._extractor_name: e_deposited_track,
-                    "e_cascade_" + self._extractor_name: e_deposited_cascade,
-                    "e_visible_" + self._extractor_name: e_total,
-                    "fraction_primary_"
+                    "e_track_target_" + self._extractor_name: e_track_target,
+                    "e_cascade_target_"
+                    + self._extractor_name: e_cascade_target,
+                    "e_target_" + self._extractor_name: e_total_target,
+                    "e_track_bkg_" + self._extractor_name: e_track_bkg,
+                    "e_cascade_bkg_" + self._extractor_name: e_cascade_bkg,
+                    "e_bkg_" + self._extractor_name: e_total_bkg,
+                    "e_target_fraction_"
+                    + self._extractor_name: e_target_fraction,
+                    "fraction_target_primary_"
                     + self._extractor_name: fraction_primary,
-                    "fraction_cascade_"
-                    + self._extractor_name: cascade_fraction,
+                    "fraction_cascade_target_"
+                    + self._extractor_name: target_cascade_fraction,
+                    "e_track_total_"
+                    + self._extractor_name: e_track_target
+                    + e_track_bkg,
+                    "e_cascade_total_"
+                    + self._extractor_name: e_cascade_target
+                    + e_cascade_bkg,
+                    "e_dep_total_" + self._extractor_name: e_total,
+                    "fraction_cascade_total_"
+                    + self._extractor_name: (
+                        cascade_fraction_tot if e_total > 0 else 0.0
+                    ),
                 }
             )
 
         output = {k: v for k, v in output.items() if k not in self._exclude}
-        # restore original mctree
-        frame.Delete(self.mctree)
-        frame[self.mctree] = tree_copy
         return output
 
     def frame_contains_info(self, frame: "icetray.I3Frame") -> bool:
@@ -156,50 +209,31 @@ def frame_contains_info(self, frame: "icetray.I3Frame") -> bool:
         return self.mctree in frame and self.mmctracklist in frame
 
     def total_track_energy(
-        self, frame: "icetray.I3Frame"
-    ) -> Tuple[float, float]:
-        """Get the total energy of track particles on entrance."""
-        e_entrance = 0
-        e_deposited = 0
-        primaries = self.get_primaries(
-            frame, self.daughters, self.highest_energy_primary
-        )
-        primaries = self.check_primary_energy(frame, primaries)
-
-        MMCTrackList = frame[self.mmctracklist]
-        if self.daughters:
-            MMCTrackList_filtered = []
-            for track in MMCTrackList:
-                try:
-                    if (
-                        frame[self.mctree].get_primary(track.GetI3Particle())
-                        in primaries
-                    ):
-                        MMCTrackList_filtered.append(track)
-                except RuntimeError as e:
-                    if "particle not found" in str(e):
-                        # log warning with event header
-                        self.warning(
-                            f"Could not find primary for track {track.GetI3Particle()}"
-                            f" in event {frame['I3EventHeader']}: {e}"
-                        )
-                        # continue to next track
-                    else:
-                        raise e
-
-            MMCTrackList = simclasses.I3MMCTrackList(MMCTrackList_filtered)
-
-        track_list = np.array(
-            MuonGun.Track.harvest(frame[self.mctree], MMCTrackList)
+        self,
+        frame: "icetray.I3Frame",
+        mctree: "dataclasses.I3MCTree",
+        entrance_energy: bool = False,
+    ) -> float:
+        """Get the total energy deposited by tracks entering the volume.
+
+        If entrance_energy is True, return the total energy entering the
+        volume as tracks instead of the energy deposited.
+        """
+        energy = 0
+
+        mmc_track_list = self.filter_track_list(
+            mctree, frame[self.mmctracklist]
         )
 
+        track_list = np.array(MuonGun.Track.harvest(mctree, mmc_track_list))
+
         while len(track_list) > 0:
             track = track_list[0]
             track_list = track_list[1:]
             try:
-                particle = frame[self.mctree].get_particle(track.id)
+                particle = mctree.get_particle(track.id)
             except RuntimeError:
-                # If the particle does not exist in the mctree, that means a partcle further up was processed and therefore it should not be counted
+                # If the particle does not exist in the mctree, that means a particle further up was processed and therefore it should not be counted
                 continue
 
             # Find distance to entrance and exit from sampling volume
@@ -235,32 +269,27 @@ def total_track_energy(
                     raise
 
             # Accumulate
-            e_deposited += e0 - e1
-            e_entrance += e0
+            if entrance_energy:
+                energy += e0
+            else:
+                energy += e0 - e1
             # get descendant ids
-            # erase particle and children from mctree
-            frame[self.mctree].erase(track.id)
-
-        # Sanity check ensuring no double counting
-        if self.daughters:
-            assert e_entrance <= sum(
-                [p.energy for p in primaries]
-            ), "Energy on entrance is greater than primary energy"
-            assert e_deposited <= sum(
-                [p.energy for p in primaries]
-            ), "Energy deposited is greater than primary energy"
-        return e_entrance, e_deposited
+            if entrance_energy:
+                # if we are looking at the entrance energy then all energy entering the volume as a track is considered "track energy" even if it is later deposited in a cascade, so we remove all descendants of the track from the mctree to avoid double counting
+                mctree.erase(track.id)
+            else:
+                # if we are looking at the deposited energy then we only want to remove the tracks that have either deposited all their energy in the volume or left the volume again thus descendants cannot produce cascades in the volume.
+                if (e1 == 0) or (intersections.second < particle.length):
+                    mctree.erase(track.id)
+        return energy
 
     def total_cascade_energy(
         self,
-        frame: "icetray.I3Frame",
+        mctree: "dataclasses.I3MCTree",
+        primaries: "dataclasses.ListI3Particle",
     ) -> float:
         """Get the total energy of cascade particles on entrance."""
-        particles = deque(
-            self.get_primaries(
-                frame, self.daughters, self.highest_energy_primary
-            )
-        )
+        particles = deque(primaries)
 
         if len(particles) == 0:
             return 0.0
@@ -273,7 +302,6 @@ def total_cascade_energy(
             [],
         )
 
-        mctree = frame[self.mctree]
         while len(particles) > 0:
             p = particles.popleft()
             p_children = mctree.get_daughters(p)
@@ -302,3 +330,25 @@ def total_cascade_energy(
         in_hull = self.hull.point_in_hull(pos)
 
         return np.sum(energies[cascade_bool & in_hull])
+
+    def filter_track_list(
+        self,
+        mctree: "dataclasses.I3MCTree",
+        track_list: "simclasses.I3MMCTrackList",
+    ) -> "simclasses.I3MMCTrackList":
+        """Filter the track list based on the mctree provided.
+
+        (This function is only meant to run on target/bkg split trees)
+        """
+        filtered_track_list = []
+        for track in track_list:
+            try:
+                mctree.get_particle(track.particle.id)
+                filtered_track_list.append(track)
+            except RuntimeError as e:
+                if "particleID not found" in str(e):
+                    # if particle is not found in the mctree then it should not be included as it in the other tree.
+                    continue
+                else:
+                    raise e
+        return simclasses.I3MMCTrackList(filtered_track_list)
diff --git a/src/graphnet/data/extractors/icecube/i3extractor.py b/src/graphnet/data/extractors/icecube/i3extractor.py
index 9b8237fe5..57df7ca9c 100644
--- a/src/graphnet/data/extractors/icecube/i3extractor.py
+++ b/src/graphnet/data/extractors/icecube/i3extractor.py
@@ -14,6 +14,8 @@
         dataclasses,
     )  # pyright: reportMissingImports=false
 
+from copy import deepcopy
+
 
 class I3Extractor(Extractor):
     """Base class for extracting information from physics I3-frames.
@@ -111,7 +113,7 @@ def __call__(self, frame: "icetray.I3Frame") -> dict:
 
     def check_primary_energy(
         self,
-        frame: "icetray.I3Frame",
+        mctree: "dataclasses.I3MCTree",
         primaries: Union[
             "dataclasses.ListI3Particle", "dataclasses.I3Particle"
         ],
@@ -122,7 +124,7 @@ def check_primary_energy(
         primary particle(s) are returned instead.
 
         Args:
-            frame: I3Frame object.
+            mctree: I3MCTree object.
             primaries: Primary particle or a list of primary particles.
         """
         assert hasattr(
@@ -132,7 +134,7 @@ def check_primary_energy(
         if isinstance(primaries, dataclasses.ListI3Particle):
             new_primaries = dataclasses.ListI3Particle()
             for primary in primaries:
-                primary = self.check_primary_energy(frame, primary)
+                primary = self.check_primary_energy(mctree, primary)
                 if isinstance(primary, dataclasses.ListI3Particle):
                     new_primaries.extend(primary)
                 elif isinstance(primary, dataclasses.I3Particle):
@@ -151,9 +153,7 @@ def check_primary_energy(
 
         if primary.energy != primary.energy:
             self.warning_once("Primary energy is nan checking daughters")
-            daughters = dataclasses.I3MCTree.get_daughters(
-                frame[self.mctree], primary
-            )
+            daughters = dataclasses.I3MCTree.get_daughters(mctree, primary)
             if len(daughters) == 0:
                 raise ValueError(
                     "Primary energy is nan and no daughters found"
@@ -165,7 +165,7 @@ def check_primary_energy(
 
     def get_primaries(
         self,
-        frame: "icetray.I3Frame",
+        mctree: "dataclasses.I3MCTree",
         daughters: bool = False,
         highest_energy_primary: bool = True,
     ) -> "dataclasses.ListI3Particle":
@@ -183,7 +183,7 @@ def get_primaries(
             Corsika case.
 
         Input:
-        frame: I3Frame object
+        mctree: I3MCTree object
         daughters: If True only daughters of the primary neutrino are returned
         highest_energy_primary: If True, return the primary with the highest
             energy. If False, return all primaries.
@@ -195,7 +195,7 @@ def get_primaries(
         ), "mctree should be instantiated by subclass"
 
         if not self._is_corsika:
-            primaries = frame[self.mctree].get_primaries()
+            primaries = mctree.get_primaries()
             if daughters:
                 primaries = [
                     p
@@ -217,16 +217,13 @@ def get_primaries(
 
                     # get the original primary neutrino(s)
                     primary_nus = [
-                        p
-                        for p in frame[self.mctree].get_primaries()
-                        if p.is_neutrino
+                        p for p in mctree.get_primaries() if p.is_neutrino
                     ]
 
                     # recursively search for in-ice neutrino daughters
                     primaries = self.find_in_ice_daughters(
-                        frame,
+                        mctree,
                         primary_nus,
-                        self.mctree,
                     )
 
                 # This is not expected to happen
@@ -246,14 +243,13 @@ def get_primaries(
             primaries = dataclasses.ListI3Particle(primaries)
 
         if self._is_corsika:
-            primaries = frame[self.mctree].get_primaries()
+            primaries = mctree.get_primaries()
         return primaries
 
     def find_in_ice_daughters(
         self,
-        frame: "icetray.I3Frame",
+        mctree: "dataclasses.I3MCTree",
         particles: "dataclasses.ListI3Particle",
-        mctree: str,
     ) -> "dataclasses.ListI3Particle":
         """Find in-ice particles in the frame."""
         if particles == []:
@@ -268,9 +264,76 @@ def find_in_ice_daughters(
             else:
                 ret.extend(
                     self.find_in_ice_daughters(
-                        frame,
-                        frame[mctree].get_daughters(p.id),
-                        mctree=mctree,
+                        mctree,
+                        mctree.get_daughters(p.id),
                     )
                 )
         return ret
+
+    def split_mc_tree(
+        self, frame: "icetray.I3Frame", highest_energy_primary: bool = True
+    ) -> "dataclasses.I3MCTree":
+        """Split the mctree in subtrees corresponding to each primary particle.
+
+        Into a subtree containing only the daughters of the primary
+        particle, and a subtree containing the rest of the particles in
+        the event.
+        """
+        assert hasattr(
+            self, "mctree"
+        ), "mctree should be instantiated by subclass"
+
+        main_tree = deepcopy(frame[self.mctree])
+        bkg_tree = deepcopy(frame[self.mctree])
+
+        if self._is_corsika:
+            # create empty main tree and return bkg tree.
+            return dataclasses.I3MCTree(), bkg_tree
+
+        all_primaries = main_tree.get_primaries()
+        if highest_energy_primary:
+            # grab the id of the highest energy primary
+
+            energies = np.array(
+                [
+                    p.energy
+                    for p in self.find_in_ice_daughters(
+                        main_tree, all_primaries
+                    )
+                ]
+            )
+            p_highest = np.array(all_primaries)[np.argmax(energies)]
+            parent_ids = [
+                p.id for p in self.get_all_parents(main_tree, p_highest)
+            ]
+            parent_ids.append(p_highest.id)
+
+        for primary in all_primaries:
+            if primary.is_neutrino:
+                if highest_energy_primary:
+                    if primary.id not in parent_ids:
+                        main_tree.erase(primary.id)
+                    else:
+                        bkg_tree.erase(primary.id)
+                else:
+                    bkg_tree.erase(primary.id)
+            else:
+                main_tree.erase(primary.id)
+        return main_tree, bkg_tree
+
+    def get_all_parents(
+        self,
+        mctree: "dataclasses.I3MCTree",
+        particle: "dataclasses.I3Particle",
+    ) -> list:
+        """Get all parents of a particle."""
+        assert hasattr(
+            self, "mctree"
+        ), "mctree should be instantiated by subclass"
+
+        parents = []
+        while mctree.has_parent(particle.id):
+            parent = mctree.get_parent(particle.id)
+            parents.append(parent)
+            particle = parent
+        return parents
diff --git a/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py b/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py
index 19228d21c..b303ae17a 100644
--- a/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py
+++ b/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py
@@ -66,7 +66,9 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             HEParticle.energy = 0
             primary_energy = sum(
                 prim.energy
-                for prim in self.get_primaries(frame, self.daughters)
+                for prim in self.get_primaries(
+                    frame[self.mctree], self.daughters
+                )
             )
             distance = -1.0
             EonEntrance = 0.0
@@ -178,8 +180,10 @@ def get_tracks(
         Args:
         frame: I3Frame object
         """
-        primaries = self.get_primaries(frame, self.daughters)
-        primaries = [self.check_primary_energy(frame, p) for p in primaries]
+        primaries = self.get_primaries(frame[self.mctree], self.daughters)
+        primaries = [
+            self.check_primary_energy(frame[self.mctree], p) for p in primaries
+        ]
 
         MMCTrackList = frame[self.mmctracklist]
         if self.daughters:
@@ -419,9 +423,10 @@ def highest_energy_starting(  # noqa: C901
         containment = GN_containment_types.no_intersect.value
         visible_length = 0.0
         if self.daughters:
-            primaries = self.get_primaries(frame, self.daughters)
+            primaries = self.get_primaries(frame[self.mctree], self.daughters)
             primaries = [
-                self.check_primary_energy(frame, p) for p in primaries
+                self.check_primary_energy(frame[self.mctree], p)
+                for p in primaries
             ]
 
             particles = self.get_descendants(frame, primaries)

From 34d2743d5b159f5125dbcc8d623786add64fd39c Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Mon, 1 Jun 2026 23:49:35 -0500
Subject: [PATCH 11/19] target/total fraction rename

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 5f1cdcc14..d932e9546 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -154,7 +154,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     f"{frame['I3EventHeader']}"
                 )
 
-            e_target_fraction = (
+            fraction_target_total = (
                 e_total_target / e_total if e_total > 0 else 0.0
             )
             target_cascade_fraction = (
@@ -181,8 +181,8 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     "e_track_bkg_" + self._extractor_name: e_track_bkg,
                     "e_cascade_bkg_" + self._extractor_name: e_cascade_bkg,
                     "e_bkg_" + self._extractor_name: e_total_bkg,
-                    "e_target_fraction_"
-                    + self._extractor_name: e_target_fraction,
+                    "fraction_target_total_"
+                    + self._extractor_name: fraction_target_total,
                     "fraction_target_primary_"
                     + self._extractor_name: fraction_primary,
                     "fraction_cascade_target_"

From 858bdb8fef8742c571c7dbdf8294701bb8aa7a9d Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 2 Jun 2026 00:30:05 -0500
Subject: [PATCH 12/19] remove daughters dependence in example

---
 examples/01_icetray/05_convert_i3_files_advanced.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/01_icetray/05_convert_i3_files_advanced.py b/examples/01_icetray/05_convert_i3_files_advanced.py
index 4d6f3b6f6..c83356466 100644
--- a/examples/01_icetray/05_convert_i3_files_advanced.py
+++ b/examples/01_icetray/05_convert_i3_files_advanced.py
@@ -91,7 +91,6 @@ def main(
         mctree="I3MCTree",
         mmctracklist="MMCTrackList",
         extractor_name=f"calorimetry_pad_{str(padding)}",
-        daughters=False,
         is_corsika=False,
     )
 

From 3da19af1a635bd956fdc0581dc0264e4e17d3a37 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 2 Jun 2026 02:01:21 -0500
Subject: [PATCH 13/19] a little leniency

---
 .../data/extractors/icecube/i3calorimetry.py   | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index d932e9546..6d52a4a16 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -69,7 +69,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             )
             if len(target_tree) + len(bkg_tree) != len(frame[self.mctree]):
                 raise ValueError(
-                    "Split mctree has different number of particles than original mctree"
+                    f"Split mctree has different number of particles than original mctree\nOriginal mctree: {len(frame[self.mctree])}\nTarget tree: {len(target_tree)}\nBkg tree: {len(bkg_tree)}\nHighest energy primary flag: {self.highest_energy_primary}\nEvent header: {frame['I3EventHeader']}"
                 )
 
             # For the target we consider either all neutrino primary products or only the highest energy primary of the neutrino depending on the flag.
@@ -95,7 +95,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             )
             # Sanity check ensuring no double counting
 
-            if e_track_target > target_primaries_energy:
+            if not (e_track_target <= target_primaries_energy * (1 + 1e-6)):
                 raise ValueError(
                     f"Energy deposited in target is greater than primary energy: {e_track_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                 )
@@ -103,7 +103,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 e_track_bkg = self.total_track_energy(
                     frame, bkg_tree, entrance_energy=self.entrance_energy
                 )
-                if e_track_bkg > bkg_primaries_energy:
+                if not (e_track_bkg <= bkg_primaries_energy * (1 + 1e-6)):
                     raise ValueError(
                         f"Energy deposited in background is greater than primary energy: {e_track_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                     )
@@ -113,7 +113,10 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             e_cascade_target = self.total_cascade_energy(
                 target_tree, target_primaries
             )
-            if e_cascade_target > target_primaries_energy:
+            if not (
+                e_cascade_target <= target_primaries_energy * (1 + 1e-6)
+                or (e_cascade_target - target_primaries_energy < 0.5)
+            ):
                 raise ValueError(
                     f"Energy deposited in cascades is greater than primary energy: {e_cascade_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                 )
@@ -121,7 +124,10 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 e_cascade_bkg = self.total_cascade_energy(
                     bkg_tree, bkg_primaries
                 )
-                if e_cascade_bkg > bkg_primaries_energy:
+                if not (
+                    e_cascade_bkg <= bkg_primaries_energy * (1 + 1e-6)
+                    or (e_cascade_bkg - bkg_primaries_energy < 0.5)
+                ):
                     raise ValueError(
                         f"Energy deposited in cascades is greater than primary energy: {e_cascade_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                     )
@@ -146,7 +152,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 or (e_total_target - target_primaries_energy < 0.5)
             ):
                 raise ValueError(
-                    "Total energy on entrance is greater than primary energy\n"
+                    "Total energy is greater than primary energy\n"
                     f"Total energy: {e_total_target}\n"
                     f"Primary energy: {target_primaries_energy}\n"
                     f"Track deposited energy: {e_track_target}\n"

From 08aefcec4885569265211c4e426010b9e5cf9a80 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 2 Jun 2026 02:02:21 -0500
Subject: [PATCH 14/19] fix highest_energy_particle True imlpementation

---
 .../data/extractors/icecube/i3extractor.py        | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3extractor.py b/src/graphnet/data/extractors/icecube/i3extractor.py
index 57df7ca9c..1573d6e33 100644
--- a/src/graphnet/data/extractors/icecube/i3extractor.py
+++ b/src/graphnet/data/extractors/icecube/i3extractor.py
@@ -293,16 +293,11 @@ def split_mc_tree(
         all_primaries = main_tree.get_primaries()
         if highest_energy_primary:
             # grab the id of the highest energy primary
-
-            energies = np.array(
-                [
-                    p.energy
-                    for p in self.find_in_ice_daughters(
-                        main_tree, all_primaries
-                    )
-                ]
+            in_ice_daughters = self.find_in_ice_daughters(
+                main_tree, [p for p in all_primaries if p.is_neutrino]
             )
-            p_highest = np.array(all_primaries)[np.argmax(energies)]
+            energies = np.array([p.energy for p in in_ice_daughters])
+            p_highest = np.array(in_ice_daughters)[np.argmax(energies)]
             parent_ids = [
                 p.id for p in self.get_all_parents(main_tree, p_highest)
             ]
@@ -333,7 +328,7 @@ def get_all_parents(
 
         parents = []
         while mctree.has_parent(particle.id):
-            parent = mctree.get_parent(particle.id)
+            parent = mctree.parent(particle.id)
             parents.append(parent)
             particle = parent
         return parents

From 533300bcc347e2fe1cac70df1d1a7a053ae31c26 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Tue, 2 Jun 2026 06:02:19 -0500
Subject: [PATCH 15/19] check bkg + remove individual checks + optimizations

---
 .../data/extractors/icecube/i3calorimetry.py  | 85 +++++++++++--------
 1 file changed, 49 insertions(+), 36 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 6d52a4a16..e4510e03c 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -90,47 +90,39 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
             target_primaries_energy = sum([p.energy for p in target_primaries])
             bkg_primaries_energy = sum([p.energy for p in bkg_primaries])
 
-            e_track_target = self.total_track_energy(
-                frame, target_tree, entrance_energy=self.entrance_energy
-            )
+            if len(target_tree) > 0:
+                e_track_target = self.total_track_energy(
+                    frame,
+                    target_tree,
+                    entrance_energy=self.entrance_energy,
+                )
+            else:
+                e_track_target = 0.0
             # Sanity check ensuring no double counting
 
             if not (e_track_target <= target_primaries_energy * (1 + 1e-6)):
                 raise ValueError(
                     f"Energy deposited in target is greater than primary energy: {e_track_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
                 )
-            if len(bkg_primaries) > 0:
+            if len(bkg_tree) > 0:
                 e_track_bkg = self.total_track_energy(
-                    frame, bkg_tree, entrance_energy=self.entrance_energy
+                    frame,
+                    bkg_tree,
+                    entrance_energy=self.entrance_energy,
                 )
-                if not (e_track_bkg <= bkg_primaries_energy * (1 + 1e-6)):
-                    raise ValueError(
-                        f"Energy deposited in background is greater than primary energy: {e_track_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
-                    )
             else:
                 e_track_bkg = 0.0
 
-            e_cascade_target = self.total_cascade_energy(
-                target_tree, target_primaries
-            )
-            if not (
-                e_cascade_target <= target_primaries_energy * (1 + 1e-6)
-                or (e_cascade_target - target_primaries_energy < 0.5)
-            ):
-                raise ValueError(
-                    f"Energy deposited in cascades is greater than primary energy: {e_cascade_target} > {target_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
+            if len(target_tree) > 0:
+                e_cascade_target = self.total_cascade_energy(
+                    target_tree, target_primaries
                 )
-            if len(bkg_primaries) > 0:
+            else:
+                e_cascade_target = 0.0
+            if len(bkg_tree) > 0:
                 e_cascade_bkg = self.total_cascade_energy(
                     bkg_tree, bkg_primaries
                 )
-                if not (
-                    e_cascade_bkg <= bkg_primaries_energy * (1 + 1e-6)
-                    or (e_cascade_bkg - bkg_primaries_energy < 0.5)
-                ):
-                    raise ValueError(
-                        f"Energy deposited in cascades is greater than primary energy: {e_cascade_bkg} > {bkg_primaries_energy}\nEvent header: {frame['I3EventHeader']}"
-                    )
             else:
                 e_cascade_bkg = 0.0
 
@@ -160,6 +152,19 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     f"{frame['I3EventHeader']}"
                 )
 
+            if not (
+                e_total_bkg <= (bkg_primaries_energy * (1 + 1e-6))
+                or (e_total_bkg - bkg_primaries_energy < 0.5)
+            ):
+                raise ValueError(
+                    "Total background energy is greater than primary energy\n"
+                    f"Total background energy: {e_total_bkg}\n"
+                    f"Background primary energy: {bkg_primaries_energy}\n"
+                    f"Track deposited background energy: {e_track_bkg}\n"
+                    f"Cascade deposited background energy: {e_cascade_bkg}\n"
+                    f"{frame['I3EventHeader']}"
+                )
+
             fraction_target_total = (
                 e_total_target / e_total if e_total > 0 else 0.0
             )
@@ -227,15 +232,18 @@ def total_track_energy(
         """
         energy = 0
 
-        mmc_track_list = self.filter_track_list(
-            mctree, frame[self.mmctracklist]
-        )
+        if self._is_corsika:
+            mmc_track_list = frame[self.mmctracklist]
+        else:
+            mmc_track_list = self.filter_track_list(
+                mctree, frame[self.mmctracklist]
+            )
 
-        track_list = np.array(MuonGun.Track.harvest(mctree, mmc_track_list))
+        track_list = deque(MuonGun.Track.harvest(mctree, mmc_track_list))
 
         while len(track_list) > 0:
-            track = track_list[0]
-            track_list = track_list[1:]
+            track = track_list.popleft()
+
             try:
                 particle = mctree.get_particle(track.id)
             except RuntimeError:
@@ -247,10 +255,16 @@ def total_track_energy(
                 track.pos, track.dir
             )
 
-            # Check if the track actually enters the volume
+            # Check if the track actually enters the volume. Values are NAN if the ray does not intersect the hull, negative if the intersection is behind the "origin". Uncertain if we can have negative intersections for both first and second intersection without it being converted to NAN (no intersection) but we check for both to be sure.
             if not (
-                np.isfinite(intersections.first)
-                and (intersections.first < particle.length)
+                (
+                    np.isfinite(intersections.first)
+                    and (intersections.first < particle.length)
+                )
+                and (
+                    np.isfinite(intersections.second)
+                    and intersections.second > 0
+                )
             ):
                 continue
 
@@ -279,7 +293,6 @@ def total_track_energy(
                 energy += e0
             else:
                 energy += e0 - e1
-            # get descendant ids
             if entrance_energy:
                 # if we are looking at the entrance energy then all energy entering the volume as a track is considered "track energy" even if it is later deposited in a cascade, so we remove all descendants of the track from the mctree to avoid double counting
                 mctree.erase(track.id)

From af011f233a1177af379a65fc8216dec9c7e2e011 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Fri, 12 Jun 2026 00:03:24 -0500
Subject: [PATCH 16/19] combine if check

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index e4510e03c..1cfd11773 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -291,12 +291,10 @@ def total_track_energy(
             # Accumulate
             if entrance_energy:
                 energy += e0
-            else:
-                energy += e0 - e1
-            if entrance_energy:
                 # if we are looking at the entrance energy then all energy entering the volume as a track is considered "track energy" even if it is later deposited in a cascade, so we remove all descendants of the track from the mctree to avoid double counting
                 mctree.erase(track.id)
             else:
+                energy += e0 - e1
                 # if we are looking at the deposited energy then we only want to remove the tracks that have either deposited all their energy in the volume or left the volume again thus descendants cannot produce cascades in the volume.
                 if (e1 == 0) or (intersections.second < particle.length):
                     mctree.erase(track.id)

From 2935c08038781fea4c196468073e987257e74a38 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Fri, 12 Jun 2026 00:03:51 -0500
Subject: [PATCH 17/19] remove outdated comment

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 1cfd11773..99f4d805e 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -62,7 +62,6 @@ def __init__(
     def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
         """Extract all the visible particles entering the volume."""
         output = {}
-        # copy the original mctree because we will be modifying it
         if self.frame_contains_info(frame):
             target_tree, bkg_tree = self.split_mc_tree(
                 frame, highest_energy_primary=self.highest_energy_primary

From 54229caa5da018291ce3b40b80fc8f9e011baa93 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Fri, 12 Jun 2026 00:07:59 -0500
Subject: [PATCH 18/19] output key name change

---
 src/graphnet/data/extractors/icecube/i3calorimetry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 99f4d805e..0f07d5883 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -203,7 +203,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     "e_cascade_total_"
                     + self._extractor_name: e_cascade_target
                     + e_cascade_bkg,
-                    "e_dep_total_" + self._extractor_name: e_total,
+                    "e_total_" + self._extractor_name: e_total,
                     "fraction_cascade_total_"
                     + self._extractor_name: (
                         cascade_fraction_tot if e_total > 0 else 0.0

From 370a6db29e276581fe58ea467711772aa579f131 Mon Sep 17 00:00:00 2001
From: "askerosted@gmail.com" <askerosted@gmail.com>
Date: Fri, 12 Jun 2026 00:27:40 -0500
Subject: [PATCH 19/19] update docstring to be more descriptive

---
 .../data/extractors/icecube/i3calorimetry.py    | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/graphnet/data/extractors/icecube/i3calorimetry.py b/src/graphnet/data/extractors/icecube/i3calorimetry.py
index 0f07d5883..90023e7f6 100644
--- a/src/graphnet/data/extractors/icecube/i3calorimetry.py
+++ b/src/graphnet/data/extractors/icecube/i3calorimetry.py
@@ -26,7 +26,22 @@ class I3Calorimetry(I3Extractor):
     """Event level energy labeling for IceCube data.
 
     This class extracts cumulative energy information from all visible
-    particles entering the detector volume, during the event.
+    particles entering the detector volume, during the event. The recorded energy is split into a "target" and "background" contribution, where the target contribution consists of all particles that downstream of neutrino primaries (or only the highest energy neutrino primary if the corresponding flag is set) and the background contribution consists of all other particles. The recorded energy is further split into a "track" and "cascade" contribution, where the track contribution consists of all energy deposited by particles that are classified as tracks (i.e. if a track is recorded in the MMCTrackList) and the cascade contribution consists of all energy deposited by particles that are not classified as tracks. The recorded energy varies depending on whether or not the entrance_energy flag is set. If the entrance_energy flag is set, the energy recorded is the energy of all visible particles entering the volume as they enter the volume. If the entrance_energy flag is not set, then the recorded energy is only the energy deposited inside the volume i.e. if a muon enters the volume with 100 GeV and leaves with 80 GeV, then the track energy recorded would be 20 GeV.
+
+    Returns a dictionary with the following keys
+    - e_track_target: Energy entering/deposited by target tracks.
+    - e_cascade_target: Energy entering/deposited by target cascades.
+    - e_target: Total energy entering/deposited by target particles.
+    - e_track_bkg: Energy entering/deposited by background tracks.
+    - e_cascade_bkg: Energy entering/deposited by background cascades.
+    - e_bkg: Total energy entering/deposited by background particles.
+    - fraction_target_total: Fraction of total recorded energy that is from target particles.
+    - fraction_target_primary: Fraction of primar(y/ies) energy that is recorded as entering/deposited by target particles.
+    - fraction_cascade_target: Fraction of target energy that is from cascades.
+    - e_track_total: Total energy entering/deposited by tracks.
+    - e_cascade_total: Total energy entering/deposited by cascades.
+    - e_total: Total energy entering/deposited by all particles.
+    - fraction_cascade_total: Fraction of total recorded energy that is from cascades.
     """
 
     def __init__(