changes to bundle pos/timing determination

Author: Aske-Rosted

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

@@ -377,14 +377,6 @@ def highest_energy_track(
                     if tmp_EonEntrance > EonEntrance:
                         particle = track_particle
 
-                        closest_pos = np.array(
-                            [
-                                track.GetXc(),
-                                track.GetYc(),
-                                track.GetZc(),
-                            ]
-                        )
-
                         EonEntrance = tmp_EonEntrance
 
                         visible_length = intersections.second - max(
@@ -421,6 +413,13 @@ def highest_energy_track(
                             )
                             particle.time = track.GetTi()
                         else:
+                            closest_pos = np.array(
+                                [
+                                    track.GetXc(),
+                                    track.GetYc(),
+                                    track.GetZc(),
+                                ]
+                            )
                             # If the track is stopping or throughgoing,
                             # pos is point closest to detector center.
                             distance = np.sqrt((closest_pos**2).sum())
@@ -748,7 +747,7 @@ def highest_energy_bundle(
         lengths = lengths[length_mask]
 
         containment = GN_containment_types.stopping_bundle.value
-        closest_pos = []
+        highest_e = 0
         for track, MGtrack in zip(MMCTrackList, MuonGun_tracks):
             intersections = self.hull.surface.intersection(
                 MGtrack.pos, MGtrack.dir
@@ -781,30 +780,24 @@ def highest_energy_bundle(
                     raise  # re-raise unexpected errors
 
             EonEntrance += track_energy
-
-            closest_pos.append(
-                np.array(
+            if track_energy > highest_e:
+                highest_e = track_energy
+                closest_pos = np.array(
                     [
                         track.GetXc(),
                         track.GetYc(),
                         track.GetZc(),
                     ]
                 )
-                * track_energy
-            )
-            if closest_time is None:
-                closest_time = track.GetTc()
-            elif closest_time < track.GetTc():
+
                 closest_time = track.GetTc()
 
-            if intersections.second > 0:
-                visible_length = max(
-                    visible_length, intersections.second - intersections.first
-                )
-                if MGtrack.length > intersections.second:
-                    containment = (
-                        GN_containment_types.throughgoing_bundle.value
-                    )
+                if intersections.second > 0:
+                    visible_length = intersections.second - intersections.first
+                    if MGtrack.length > intersections.second:
+                        containment = (
+                            GN_containment_types.throughgoing_bundle.value
+                        )
 
         # If no intersection.second is every positive
         # the visible_length can still be negative here
@@ -823,8 +816,6 @@ def highest_energy_bundle(
             visible_length >= 0
         ), f"Visible length is negative for particle {frame['I3EventHeader']}"
 
-        closest_pos = np.sum(closest_pos, axis=0) / EonEntrance
-
         bundle.pos = dataclasses.I3Position(
             closest_pos[0], closest_pos[1], closest_pos[2]
         )