Add feature to extract and reconstruct track and cascade separately.

src/graphnet/data/constants.py

@@ -38,6 +38,7 @@ class TRUTH:
     ICECUBE86 = [
         "energy",
         "energy_track",
+        "energy_cascade",
         "position_x",
         "position_y",
         "position_z",

src/graphnet/data/extractors/i3truthextractor.py

@@ -108,6 +108,7 @@ def __call__(
             "track_length": padding_value,
             "stopped_muon": padding_value,
             "energy_track": padding_value,
+            "energy_cascade": padding_value,
             "inelasticity": padding_value,
             "DeepCoreFilter_13": padding_value,
             "CascadeFilter_13": padding_value,
@@ -177,10 +178,15 @@ def __call__(
             try:
                 (
                     energy_track,
+                    energy_cascade,
                     inelasticity,
                 ) = self._get_primary_track_energy_and_inelasticity(frame)
             except RuntimeError:  # track energy fails on northeren tracks with ""Hadrons" has no mass implemented. Cannot get total energy."
-                energy_track, inelasticity = (padding_value, padding_value)
+                energy_track, energy_cascade, inelasticity = (
+                    padding_value,
+                    padding_value,
+                    padding_value,
+                )
 
             output.update(
                 {
@@ -197,6 +203,7 @@ def __call__(
                         frame, padding_value
                     ),
                     "energy_track": energy_track,
+                    "energy_cascade": energy_cascade,
                     "inelasticity": inelasticity,
                 }
             )
@@ -376,7 +383,7 @@ def _get_primary_particle_interaction_type_and_elasticity(
     def _get_primary_track_energy_and_inelasticity(
         self,
         frame: "icetray.I3Frame",
-    ) -> Tuple[float, float]:
+    ) -> Tuple[float, float, float]:
         """Get the total energy of tracks from primary, and inelasticity.
 
         Args:
@@ -399,9 +406,10 @@ def _get_primary_track_energy_and_inelasticity(
 
         energy_total = primary.total_energy
         energy_track = sum(track.total_energy for track in tracks)
+        energy_cascade = energy_total - energy_track
         inelasticity = 1.0 - energy_track / energy_total
 
-        return energy_track, inelasticity
+        return energy_track, energy_cascade, inelasticity
 
     # Utility methods
     def _find_data_type(self, mc: bool, input_file: str) -> str:

src/graphnet/models/task/reconstruction.py

@@ -125,6 +125,26 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.pow(10, x[:, 0] + 1.0).unsqueeze(1)
 
 
+class EnergyTCReconstruction(Task):
+    """Reconstructs track and cascade energies using stable method."""
+
+    # Requires two features: untransformed energy for track and cascade
+    default_target_labels = ["energy_track", "energy_cascade"]
+    default_prediction_labels = ["energy_track_pred", "energy_cascade_pred"]
+    nb_inputs = 2
+
+    def _forward(self, x: Tensor) -> Tensor:
+        # Transform to positive energy domain avoiding `-inf` in `log10`
+        # Transform, thereby preventing overflow and underflow error.
+        x[:, 0] = torch.nn.functional.softplus(
+            x[:, 0].clone(), beta=0.05
+        ) + eps_like(x[:, 0].clone())
+        x[:, 1] = torch.nn.functional.softplus(
+            x[:, 1].clone(), beta=0.05
+        ) + eps_like(x[:, 1].clone())
+        return x
+
+
 class EnergyReconstructionWithUncertainty(EnergyReconstruction):
     """Reconstructs energy and associated uncertainty (log(var))."""