From 6b421d3c7e90b3367421798c25f1bc1be8d9f964 Mon Sep 17 00:00:00 2001 From: "askerosted@gmail.com" Date: Tue, 13 Jan 2026 01:05:37 -0600 Subject: [PATCH 1/9] 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" Date: Tue, 13 Jan 2026 01:17:15 -0600 Subject: [PATCH 2/9] 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" Date: Tue, 13 Jan 2026 18:36:27 -0600 Subject: [PATCH 3/9] 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" Date: Tue, 13 Jan 2026 18:59:41 -0600 Subject: [PATCH 4/9] 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" Date: Mon, 26 Jan 2026 00:37:35 -0600 Subject: [PATCH 5/9] 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" Date: Mon, 26 Jan 2026 03:42:38 -0600 Subject: [PATCH 6/9] 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" Date: Tue, 27 Jan 2026 00:15:49 -0600 Subject: [PATCH 7/9] 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" Date: Sun, 8 Mar 2026 21:12:16 -0500 Subject: [PATCH 8/9] 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" Date: Wed, 8 Apr 2026 20:27:08 -0500 Subject: [PATCH 9/9] 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,