Add boolean flag for whether or not to use an small epsilon on charge

Author: astrojarred

src/graphnet/models/detector/magic.py

@@ -1,6 +1,6 @@
 """MAGIC-specific `Detector` class(es)."""
 
-from typing import Dict, Callable
+from typing import Dict, Callable, Any
 import torch
 import os
 
@@ -18,6 +18,17 @@ class MAGIC(Detector):
     # By default, treat the telescope ID as a spatial-like z-coordinate
     xyz = ["x_cam", "y_cam", "tel_id"]
 
+    def __init__(self, use_charge_epsilon: bool = True, **kwargs: Any) -> None:
+        """Construct detector for the MAGIC telescopes.
+
+        Args:
+            use_charge_epsilon: Whether to add a small epsilon value to charge
+                before taking log10 to avoid log(0). Defaults to True (uses
+                1e-6). Set to False if data is guaranteed to have zero-padding.
+        """
+        super().__init__(**kwargs)
+        self._charge_epsilon = 1e-6 if use_charge_epsilon else 0.0
+
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension.
 
@@ -47,5 +58,4 @@ def _time(self, x: torch.tensor) -> torch.tensor:
 
     def _charge(self, x: torch.tensor) -> torch.tensor:
         """Add a small epsilon to avoid log(0)."""
-        epsilon = 1e-6
-        return torch.log10(x + epsilon)
+        return torch.log10(x + self._charge_epsilon)