deepcopy of mctree to avoid changes between extractors

Author: Aske-Rosted

src/graphnet/data/extractors/combine_extractors.py

@@ -7,7 +7,12 @@
 from typing import List, Dict
 
 if has_icecube_package() or TYPE_CHECKING:
-    from icecube import icetray  # pyright: reportMissingImports=false
+    from icecube import (
+        icetray,
+        dataclasses,
+    )  # pyright: reportMissingImports=false
+
+from copy import deepcopy
 
 
 class CombinedExtractor(I3Extractor):
@@ -44,5 +49,10 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, float]:
         """
         output = {}
         for extractor in self._extractors:
+            if hasattr(extractor, "mctree") and frame.Has(extractor.mctree):
+                mctree_backup = deepcopy(frame[extractor.mctree])
             output.update(extractor(frame))
+            if hasattr(extractor, "mctree") and frame.Has(extractor.mctree):
+                frame.Delete(extractor.mctree)
+                frame[extractor.mctree] = mctree_backup
         return output

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

@@ -88,7 +88,7 @@ def __init__(
             self._borders = borders
 
         self._extend_boundary = extend_boundary
-        self._mctree = mctree
+        self.mctree = mctree
 
     def set_gcd(self, i3_file: str, gcd_file: Optional[str] = None) -> None:
         """Extract GFrame and CFrame from i3/gcd-file pair.
@@ -130,8 +130,8 @@ def __call__(
         self, frame: "icetray.I3Frame", padding_value: Any = -1
     ) -> Dict[str, Any]:
         """Extract truth-level information."""
-        is_mc = frame_is_montecarlo(frame, self._mctree)
-        is_noise = frame_is_noise(frame, self._mctree)
+        is_mc = frame_is_montecarlo(frame, self.mctree)
+        is_noise = frame_is_noise(frame, self.mctree)
         sim_type = self._find_data_type(is_mc, self._i3_file, frame)
 
         output = {
@@ -289,7 +289,7 @@ def __call__(
     def _extract_dbang_decay_length(
         self, frame: "icetray.I3Frame", padding_value: float = -1
     ) -> float:
-        mctree = frame[self._mctree]
+        mctree = frame[self.mctree]
         try:
             p_true = mctree.primaries[0]
             p_daughters = mctree.get_daughters(p_true)
@@ -418,12 +418,12 @@ def _get_primary_particle_interaction_type_and_elasticity(
             try:
                 MCInIcePrimary = frame["MCInIcePrimary"]
             except KeyError:
-                MCInIcePrimary = frame[self._mctree][0]
+                MCInIcePrimary = frame[self.mctree][0]
             if (
                 MCInIcePrimary.energy != MCInIcePrimary.energy
             ):  # This is a nan check. Only happens for some muons
                 # where second item in MCTree is primary. Weird!
-                MCInIcePrimary = frame[self._mctree][1]
+                MCInIcePrimary = frame[self.mctree][1]
                 # For some strange reason the second entry is identical in
                 # all variables and has no nans (always muon)
         else:
@@ -472,7 +472,7 @@ def _get_primary_track_energy_and_inelasticity(
             Tuple containing the energy of tracks from primary, and the
                 corresponding inelasticity.
         """
-        mc_tree = frame[self._mctree]
+        mc_tree = frame[self.mctree]
         primary = mc_tree.primaries[0]
         daughters = mc_tree.get_daughters(primary)
         tracks = []

src/graphnet/data/readers/i3reader.py

@@ -12,9 +12,14 @@
 from graphnet.utilities.filesys import find_i3_files
 from .graphnet_file_reader import GraphNeTFileReader
 
+from copy import deepcopy
+
 
 if has_icecube_package():
-    from icecube import icetray, dataio  # pyright: reportMissingImports=false
+    from icecube import (
+        icetray,
+        dataio,
+    )  # pyright: reportMissingImports=false
 
 
 class I3Reader(GraphNeTFileReader):
@@ -87,22 +92,54 @@ def __call__(
         # Open I3 file
         i3_file_io = dataio.I3File(file_path.i3_file, "r")
         data = list()
-        while i3_file_io.more():
-            try:
-                frame = i3_file_io.pop_physics()
-            except Exception as e:
-                if "I3" in str(e):
+        try:
+            while i3_file_io.more():
+                try:
+                    frame = i3_file_io.pop_physics()
+                except Exception as e:
+                    if "I3" in str(e):
+                        continue
+                # check if frame should be skipped
+                if self._skip_frame(frame):
                     continue
-            # check if frame should be skipped
-            if self._skip_frame(frame):
-                continue
-
-            # Try to extract data from I3Frame
-            results = [extractor(frame) for extractor in self._extractors]
-
-            data_dict = OrderedDict(zip(self.extracor_names, results))
 
-            data.append(data_dict)
+                # Try to extract data from I3Frame
+                results = []
+                try:
+                    for extractor in self._extractors:
+                        if hasattr(extractor, "mctree") and frame.Has(
+                            extractor.mctree
+                        ):
+                            mctree_backup = deepcopy(frame[extractor.mctree])
+
+                        results.append(extractor(frame))
+                        if hasattr(extractor, "mctree") and frame.Has(
+                            extractor.mctree
+                        ):
+                            frame.Delete(extractor.mctree)
+                            frame[extractor.mctree] = mctree_backup
+
+                    data_dict = OrderedDict(zip(self.extracor_names, results))
+
+                except KeyError as e:
+                    if "Deserialization failed for object" in str(e):
+                        self.warning(
+                            f"KeyError {e} in file {file_path.i3_file}"
+                            " - skipping frame."
+                        )
+                        continue
+                    else:
+                        raise e
+
+                data.append(data_dict)
+        except KeyError as e:
+            if "Deserialization failed for object" in str(e):
+                self.warning(
+                    f"KeyError {e} in file {file_path.i3_file}"
+                    " - skipping remaining frames."
+                )
+            else:
+                raise e
         return data
 
     def find_files(self, path: Union[str, List[str]]) -> List[I3FileSet]: