add primary type info

Author: Aske-Rosted

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

@@ -126,11 +126,17 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                 )
             except IndexError:
                 parent_type = 0
-
             if primary_energy > 0:
                 primary_fraction = EonEntrance / primary_energy
             else:
                 primary_fraction = -1
+
+            if visible_length != -1:
+                primary = frame[self.mctree].get_primary(HEParticle.id)
+                primary_is_nu, primary_type = primary.is_neutrino, primary.type
+            else:
+                primary_type = 0
+                primary_is_nu = False
             output.update(
                 {
                     "e_fraction_" + self._extractor_name: primary_fraction,
@@ -153,6 +159,8 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]:
                     "particle_type_" + self._extractor_name: HEParticle.type,
                     "containment_" + self._extractor_name: containment,
                     "parent_type_" + self._extractor_name: parent_type,
+                    "primary_type_" + self._extractor_name: primary_type,
+                    "primary_is_nu_" + self._extractor_name: primary_is_nu,
                 }
             )