Skip to content

Commit 9cc3a70

Browse files
sevmagsevmag
committed
cover case of empty primary list
1 parent e8740aa commit 9cc3a70

File tree

1 file changed

+68
-60
lines changed

1 file changed

+68
-60
lines changed

‎src/graphnet/data/extractors/icecube/i3calorimetry.py

Lines changed: 68 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -82,28 +82,39 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
8282
self.highest_energy_primary,
8383
)
8484

85-
MMCTrackList = frame[self.mmctracklist]
86-
# Filter tracks that are not daughters of the desired
87-
if self.daughters:
88-
MMCTrackList = [
89-
track
90-
for track in MMCTrackList
91-
if frame[self.mctree].get_primary(track.GetI3Particle())
92-
in primaries
93-
]
94-
95-
# Create a lookup dict for the tracks
96-
track_lookup = {}
97-
for track in MuonGun.Track.harvest(
98-
frame[self.mctree], MMCTrackList
99-
):
100-
track_lookup[track.id] = track
85+
if not len(primaries) == 0:
86+
87+
MMCTrackList = frame[self.mmctracklist]
88+
# Filter tracks that are not daughters of the desired
89+
if self.daughters:
90+
temp_MMCTrackList = []
91+
for track in MMCTrackList:
92+
for p in primaries:
93+
if frame[self.mctree].is_in_subtree(
94+
p.id, track.GetI3Particle().id
95+
):
96+
temp_MMCTrackList.append(track)
97+
break
98+
MMCTrackList = temp_MMCTrackList
99+
100+
# Create a lookup dict for the tracks
101+
track_lookup = {}
102+
for track in MuonGun.Track.harvest(
103+
frame[self.mctree], MMCTrackList
104+
):
105+
track_lookup[track.id] = track
106+
107+
e_dep_cascade, e_dep_track, e_ent_track = self.get_energies(
108+
frame, primaries, track_lookup
109+
)
101110

102-
e_dep_cascade, e_dep_track, e_ent_track = self.get_energies(
103-
frame, primaries, track_lookup
104-
)
111+
primary_energy = sum([p.energy for p in primaries])
112+
else:
113+
e_ent_track = np.nan
114+
e_dep_track = np.nan
115+
e_dep_cascade = np.nan
116+
primary_energy = np.nan
105117

106-
primary_energy = sum([p.energy for p in primaries])
107118
e_total = e_ent_track + e_dep_cascade
108119

109120
# In case all particles are considered and
@@ -117,55 +128,52 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
117128
)
118129
):
119130
self.warning(
120-
"No energy deposited in the hull, \
121-
Think about in creasing the padding. \
122-
\nCurrent padding: {}\
131+
"No energy deposited in the hull,"
132+
"Think about in creasing the padding."
133+
f"\nCurrent padding: {self.hull.padding}"
134+
f"\nTotal energy: {e_total}"
135+
f"\nTrack energy: {e_ent_track}"
136+
f"\nCascade energy: {e_dep_cascade}"
137+
f"\nEvent header: {frame['I3EventHeader']}"
138+
)
139+
140+
# Check only in the case that there were primaries
141+
if not len(primaries) == 0:
142+
143+
# total energy should always be less than the primary energy
144+
assert e_total <= (
145+
primary_energy * (1 + 1e-6)
146+
), "Total energy on entrance is greater than primary energy\
123147
\nTotal energy: {}\
148+
\nPrimary energy: {}\
124149
\nTrack energy: {}\
125-
\nCascade energy: {} \
126-
\nEvent header: {}\
127-
".format(
128-
self.hull.padding,
129-
e_total,
130-
e_ent_track,
131-
e_dep_cascade,
132-
frame["I3EventHeader"],
133-
)
150+
\nCascade energy: {}\
151+
{}".format(
152+
e_total,
153+
primary_energy,
154+
e_ent_track,
155+
e_dep_cascade,
156+
frame["I3EventHeader"],
134157
)
135158

136-
# total energy should always be less than the primary energy
137-
assert e_total <= (
138-
primary_energy * (1 + 1e-6)
139-
), "Total energy on entrance is greater than primary energy\
140-
\nTotal energy: {}\
141-
\nPrimary energy: {}\
142-
\nTrack energy: {}\
143-
\nCascade energy: {}\
144-
{}".format(
145-
e_total,
146-
primary_energy,
147-
e_ent_track,
148-
e_dep_cascade,
149-
frame["I3EventHeader"],
150-
)
159+
assert (
160+
primary_energy > 0
161+
), "Primary energy is 0, this should not happen.\
162+
\nTotal energy: {}\
163+
\nTrack energy: {}\
164+
\nCascade energy: {}\
165+
{}".format(
166+
e_total,
167+
e_ent_track,
168+
e_dep_cascade,
169+
frame["I3EventHeader"],
170+
)
171+
fraction_primary = e_total / primary_energy
151172

152173
cascade_fraction = None
153174
if e_total > 0:
154175
cascade_fraction = e_dep_cascade / e_total
155176

156-
assert (
157-
primary_energy > 0
158-
), "Primary energy is 0, this should not happen.\
159-
\nTotal energy: {}\
160-
\nTrack energy: {}\
161-
\nCascade energy: {}\
162-
{}".format(
163-
e_total,
164-
e_ent_track,
165-
e_dep_cascade,
166-
frame["I3EventHeader"],
167-
)
168-
fraction_primary = e_total / primary_energy
169177
output.update(
170178
{
171179
"e_entrance_track_" + self._extractor_name: e_ent_track,

0 commit comments

Comments
 (0)