Minor updates as part of oscNext development

src/graphnet/data/dataconverter.py

@@ -203,7 +203,9 @@ def __init__(
     @final
     def __call__(
         self,
-        directories: Union[str, List[str]],
+        i3_files = None,
+        gcd_files = None,
+        directories: Optional[Union[str, List[str]]] = None,
         recursive: Optional[bool] = True,
     ) -> None:
         """Convert I3-files in `directories.
@@ -213,13 +215,27 @@ def __call__(
                 should be converted to an intermediate file format.
             recursive: Whether or not to search the directories recursively.
         """
+
+        # User can either directly pass files, or instead pass directories in which files can be searched for
         # Find all I3 and GCD files in the specified directories.
-        i3_files, gcd_files = find_i3_files(
-            directories, self._gcd_rescue, recursive
-        )
+        if i3_files is not None :
+            assert gcd_files is not None
+            assert directories is None
+            assert isinstance(i3_files, list)
+            assert isinstance(gcd_files, list)
+            #TOD support GCD rescue
+        else :
+            assert i3_files is None
+            assert gcd_files is None
+            i3_files, gcd_files = find_i3_files(
+                directories, self._gcd_rescue, recursive
+            )
+
+        # Check I3/GCD files
         if len(i3_files) == 0:
             self.error(f"No files found in {directories}.")
             return
+        assert len(i3_files) == len(gcd_files), "Number of GCD files must match number of i3 files"
 
         # Save a record of the found I3 files in the output directory.
         self._save_filenames(i3_files)

src/graphnet/data/dataset/dataset.py

@@ -653,14 +653,12 @@ def _get_labels(self, truth_dict: Dict[str, Any]) -> Dict[str, Any]:
                 "muon": int(abs_pid == 13),
                 "muon_stopped": int(truth_dict.get("stopped_muon") == 1),
                 "noise": int((abs_pid == 1) & (sim_type != "data")),
-                "neutrino": int(
-                    (abs_pid != 13) & (abs_pid != 1)
-                ),  # @TODO: `abs_pid in [12,14,16]`?
+                "neutrino": int(abs_pid in [12,14, 16]),
                 "v_e": int(abs_pid == 12),
                 "v_u": int(abs_pid == 14),
                 "v_t": int(abs_pid == 16),
                 "track": int(
-                    (abs_pid == 14) & (truth_dict["interaction_type"] == 1)
+                    (abs_pid == 14) & (truth_dict.get("interaction_type") == 1)
                 ),
                 "dbang": self._get_dbang_label(truth_dict),
                 "corsika": int(abs_pid > 20),

src/graphnet/data/extractors/i3featureextractor.py

@@ -98,6 +98,7 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, List[Any]]:
         event_time = frame["I3EventHeader"].start_time.mod_julian_day_double
 
         for om_key in om_keys:
+
             # Common values for each OM
             x = self._gcd_dict[om_key].position.x
             y = self._gcd_dict[om_key].position.y
@@ -112,6 +113,12 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, List[Any]]:
             pmt_number = om_key[2]
             dom_type = self._gcd_dict[om_key].omtype
 
+            # Definition of DOM positions changed in 2017 data (and onwards) due to rounding in a database
+            # Round all DOM positions to nearest cm so that there are no season-to-season differences for the training to latch on to
+            x = round( x / icetray.I3Units.centimeter ) * icetray.I3Units.centimeter
+            y = round( y / icetray.I3Units.centimeter ) * icetray.I3Units.centimeter
+            z = round( z / icetray.I3Units.centimeter ) * icetray.I3Units.centimeter
+
             # DOM flags
             if bright_doms:
                 is_bright_dom = 1 if om_key in bright_doms else 0

src/graphnet/models/task/reconstruction.py

@@ -161,6 +161,19 @@ def _forward(self, x: Tensor) -> Tensor:
         return pred
 
 
+class LengthReconstruction(StandardLearnedTask):
+    """Reconstructs track length."""
+
+    # Requires one feature: untransformed energy
+    default_target_labels = ["length"]
+    default_prediction_labels = ["length_pred"]
+    nb_inputs = 1
+
+    def _forward(self, x: Tensor) -> Tensor:
+        # Leave as it is
+        return x
+
+
 class VertexReconstruction(StandardLearnedTask):
     """Reconstructs vertex position and time."""
 
@@ -175,10 +188,13 @@ class VertexReconstruction(StandardLearnedTask):
     nb_inputs = 4
 
     def _forward(self, x: Tensor) -> Tensor:
-        # Scale xyz to roughly the right order of magnitude, leave time
-        x[:, 0] = x[:, 0] * 1e2
-        x[:, 1] = x[:, 1] * 1e2
-        x[:, 2] = x[:, 2] * 1e2
+        # Scale xyzt to roughly the right order of magnitude
+        x[:, 0] = x[:, 0] / 1e2 # 100s of m
+        x[:, 1] = x[:, 1] / 1e2 # 100s of m
+        x[:, 2] = x[:, 2] / 1e2 # 100s of m
+        x[:, 3] = x[:, 3] / 1e4 # 1000s of ns, peak at around 10,000 ns
+
+        #TODO remove this?
 
         return x
 
@@ -197,10 +213,9 @@ class PositionReconstruction(StandardLearnedTask):
 
     def _forward(self, x: Tensor) -> Tensor:
         # Scale to roughly the right order of magnitude
-        x[:, 0] = x[:, 0] * 1e2
-        x[:, 1] = x[:, 1] * 1e2
-        x[:, 2] = x[:, 2] * 1e2
-
+        x[:, 0] = x[:, 0] / 1e2
+        x[:, 1] = x[:, 1] / 1e2
+        x[:, 2] = x[:, 2] / 1e2
         return x