Merge pull request #218 from asogaard/marc-analysis

Code for Marc Jacquart's analysis

datasets/upgrade_nmo_sensitivity.yml

@@ -0,0 +1,6 @@
+name: upgrade_nmo_sensitivity
+server: cobalt.icecube.wisc.edu
+sources: |
+  /data/user/jvmead/ug_sim/genie2/genie/step4/140021/
+  /data/user/jvmead/ug_sim/genie2/genie/step4/141021/
+date: 20220607

examples/convert_i3_to_sqlite.py

@@ -1,7 +1,7 @@
 """Minimum working example (MWE) to use SQLiteDataConverter
 """
 
-from graphnet.data.i3extractor import (
+from graphnet.data.extractors import (
     I3FeatureExtractorIceCube86,
     I3FeatureExtractorIceCubeUpgrade,
     I3RetroExtractor,

src/graphnet/components/loss_functions.py

@@ -57,6 +57,30 @@ def _forward(self, prediction: Tensor, target: Tensor) -> Tensor:
         """Syntax similar to `.forward` for implentation in inheriting classes."""
 
 
+class MSELoss(LossFunction):
+    """Mean squared error loss."""
+
+    def _forward(self, prediction: Tensor, target: Tensor) -> Tensor:
+        """Implementation of loss calculation."""
+        # Check(s)
+        assert prediction.dim() == 2
+        assert prediction.size() == target.size()
+
+        elements = torch.mean((prediction - target) ** 2, dim=-1)
+        return elements
+
+
+class RMSELoss(MSELoss):
+    """Root mean squared error loss."""
+
+    def _forward(self, prediction: Tensor, target: Tensor) -> Tensor:
+        """Implementation of loss calculation."""
+        # Check(s)
+        elements = super()._forward(prediction, target)
+        elements = torch.sqrt(elements)
+        return elements
+
+
 class LogCoshLoss(LossFunction):
     """Log-cosh loss function.
 
@@ -281,7 +305,7 @@ def _forward(self, prediction: Tensor, target: Tensor) -> Tensor:
         return self._evaluate(p, t)
 
 
-class EuclideanDistance(LossFunction):
+class EuclideanDistanceLoss(LossFunction):
     def _forward(self, prediction: Tensor, target: Tensor) -> Tensor:
         """Calculates the 3D Euclidean distance between predicted and target.
 

src/graphnet/data/constants.py

@@ -26,6 +26,7 @@ class FEATURES:
 class TRUTH:
     ICECUBE86 = [
         "energy",
+        "energy_track",
         "position_x",
         "position_y",
         "position_z",
@@ -36,6 +37,7 @@ class TRUTH:
         "sim_type",
         "interaction_type",
         "interaction_time",  # Added for vertex reconstruction
+        "inelasticity",
         "stopped_muon",
     ]
     DEEPCORE = ICECUBE86

src/graphnet/data/dataconverter.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from graphnet.data.i3extractor import (
+from graphnet.data.extractors import (
     I3Extractor,
     I3ExtractorCollection,
     I3TruthExtractor,

src/graphnet/data/extractors/__init__.py

@@ -0,0 +1,4 @@
+from .i3extractor import *
+from .i3featureextractor import *
+from .i3truthextractor import *
+from .i3retroextractor import *

src/graphnet/data/extractors/i3extractor.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from graphnet.utilities.logging import get_logger
+
+logger = get_logger()
+
+try:
+    from icecube import (
+        icetray,
+        dataio,
+    )  # pyright: reportMissingImports=false
+except ImportError:
+    logger.info("icecube package not available.")
+
+
+class I3Extractor(ABC):
+    """Extracts relevant information from physics frames."""
+
+    def __init__(self, name):
+
+        # Member variables
+        self._i3_file = None
+        self._gcd_file = None
+        self._gcd_dict = None
+        self._calibration = None
+        self._name = name
+
+    def set_files(self, i3_file, gcd_file):
+        # @TODO: Is it necessary to set the `i3_file`? It is only used in one
+        #        place in `I3TruthExtractor`, and there only in a way that might
+        #        be solved another way.
+        self._i3_file = i3_file
+        self._gcd_file = gcd_file
+        self._load_gcd_data()
+
+    def _load_gcd_data(self):
+        """Loads the geospatial information contained in the gcd-file."""
+        gcd_file = dataio.I3File(self._gcd_file)
+        g_frame = gcd_file.pop_frame(icetray.I3Frame.Geometry)
+        c_frame = gcd_file.pop_frame(icetray.I3Frame.Calibration)
+        self._gcd_dict = g_frame["I3Geometry"].omgeo
+        self._calibration = c_frame["I3Calibration"]
+
+    @abstractmethod
+    def __call__(self, frame) -> dict:
+        """Extracts relevant information from frame."""
+        pass
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+
+class I3ExtractorCollection(list):
+    """Class to manage multiple I3Extractors."""
+
+    def __init__(self, *extractors):
+        # Check(s)
+        for extractor in extractors:
+            assert isinstance(extractor, I3Extractor)
+
+        # Base class constructor
+        super().__init__(extractors)
+
+    def set_files(self, i3_file, gcd_file):
+        for extractor in self:
+            extractor.set_files(i3_file, gcd_file)
+
+    def __call__(self, frame) -> List[dict]:
+        return [extractor(frame) for extractor in self]

src/graphnet/data/extractors/i3featureextractor.py

@@ -0,0 +1,198 @@
+from graphnet.data.extractors.i3extractor import I3Extractor
+from graphnet.utilities.logging import get_logger
+
+logger = get_logger()
+try:
+    from icecube import (
+        dataclasses,
+    )  # pyright: reportMissingImports=false
+except ImportError:
+    logger.info("icecube package not available.")
+
+
+class I3FeatureExtractor(I3Extractor):
+    def __init__(self, pulsemap):
+        self._pulsemap = pulsemap
+        super().__init__(pulsemap)
+
+    def _get_om_keys_and_pulseseries(self, frame):
+        """Gets the indicies for the gcd_dict and the pulse series
+
+        Args:
+            frame (i3 physics frame): i3 physics frame
+
+        Returns:
+            om_keys (index): the indicies for the gcd_dict
+            data    (??)   : the pulse series
+        """
+        data = frame[self._pulsemap]
+        try:
+            om_keys = data.keys()
+        except:  # noqa: E722
+            try:
+                if "I3Calibration" in frame.keys():
+                    data = frame[self._pulsemap].apply(frame)
+                    om_keys = data.keys()
+                else:
+                    frame["I3Calibration"] = self._calibration
+                    data = frame[self._pulsemap].apply(frame)
+                    om_keys = data.keys()
+                    del frame[
+                        "I3Calibration"
+                    ]  # Avoid adding unneccesary data to frame
+            except:  # noqa: E722
+                data = dataclasses.I3RecoPulseSeriesMap.from_frame(
+                    frame, self._pulsemap
+                )
+                om_keys = data.keys()
+        return om_keys, data
+
+
+class I3FeatureExtractorIceCube86(I3FeatureExtractor):
+    def __call__(self, frame) -> dict:
+        """Extract features to be used as inputs to GNN models."""
+
+        output = {
+            "charge": [],
+            "dom_time": [],
+            "dom_x": [],
+            "dom_y": [],
+            "dom_z": [],
+            "width": [],
+            "pmt_area": [],
+            "rde": [],
+        }
+
+        try:
+            om_keys, data = self._get_om_keys_and_pulseseries(frame)
+        except KeyError:
+            return output
+
+        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
+            z = self._gcd_dict[om_key].position.z
+            area = self._gcd_dict[om_key].area
+            rde = self._get_relative_dom_efficiency(frame, om_key)
+
+            # Loop over pulses for each OM
+            pulses = data[om_key]
+            for pulse in pulses:
+                output["charge"].append(pulse.charge)
+                output["dom_time"].append(pulse.time)
+                output["width"].append(pulse.width)
+                output["pmt_area"].append(area)
+                output["rde"].append(rde)
+                output["dom_x"].append(x)
+                output["dom_y"].append(y)
+                output["dom_z"].append(z)
+
+        return output
+
+    def _get_relative_dom_efficiency(self, frame, om_key):
+        if (
+            "I3Calibration" in frame
+        ):  # Not available for e.g. mDOMs in IceCube Upgrade
+            rde = frame["I3Calibration"].dom_cal[om_key].relative_dom_eff
+        else:
+            try:
+                rde = self._calibration.dom_cal[om_key].relative_dom_eff
+            except:  # noqa: E722
+                rde = -1
+        return rde
+
+
+class I3FeatureExtractorIceCubeDeepCore(I3FeatureExtractorIceCube86):
+    """..."""
+
+
+class I3FeatureExtractorIceCubeUpgrade(I3FeatureExtractorIceCube86):
+    def __call__(self, frame) -> dict:
+        """Extract features to be used as inputs to GNN models."""
+
+        output = {
+            "string": [],
+            "pmt_number": [],
+            "dom_number": [],
+            "pmt_dir_x": [],
+            "pmt_dir_y": [],
+            "pmt_dir_z": [],
+            "dom_type": [],
+        }
+
+        try:
+            om_keys, data = self._get_om_keys_and_pulseseries(frame)
+        except KeyError:  # Target pulsemap does not exist in `frame`
+            return output
+
+        for om_key in om_keys:
+            # Common values for each OM
+            pmt_dir_x = self._gcd_dict[om_key].orientation.x
+            pmt_dir_y = self._gcd_dict[om_key].orientation.y
+            pmt_dir_z = self._gcd_dict[om_key].orientation.z
+            string = om_key[0]
+            dom_number = om_key[1]
+            pmt_number = om_key[2]
+            dom_type = self._gcd_dict[om_key].omtype
+
+            # Loop over pulses for each OM
+            pulses = data[om_key]
+            for _ in pulses:
+                output["string"].append(string)
+                output["pmt_number"].append(pmt_number)
+                output["dom_number"].append(dom_number)
+                output["pmt_dir_x"].append(pmt_dir_x)
+                output["pmt_dir_y"].append(pmt_dir_y)
+                output["pmt_dir_z"].append(pmt_dir_z)
+                output["dom_type"].append(dom_type)
+
+        # Add features from IceCube86
+        output_icecube86 = super().__call__(frame)
+        output.update(output_icecube86)
+        return output
+
+
+class I3PulseNoiseTruthFlagIceCubeUpgrade(I3FeatureExtractorIceCube86):
+    def __call__(self, frame) -> dict:
+        """Extract features to be used as inputs to GNN models."""
+
+        output = {
+            "string": [],
+            "pmt_number": [],
+            "dom_number": [],
+            "pmt_dir_x": [],
+            "pmt_dir_y": [],
+            "pmt_dir_z": [],
+            "dom_type": [],
+            "truth_flag": [],
+        }
+
+        try:
+            om_keys, data = self._get_om_keys_and_pulseseries(frame)
+        except KeyError:  # Target pulsemap does not exist in `frame`
+            return output
+
+        for om_key in om_keys:
+            # Common values for each OM
+            pmt_dir_x = self._gcd_dict[om_key].orientation.x
+            pmt_dir_y = self._gcd_dict[om_key].orientation.y
+            pmt_dir_z = self._gcd_dict[om_key].orientation.z
+            string = om_key[0]
+            dom_number = om_key[1]
+            pmt_number = om_key[2]
+            dom_type = self._gcd_dict[om_key].omtype
+
+            # Loop over pulses for each OM
+            pulses = data[om_key]
+            for truth_flag in pulses:
+                output["string"].append(string)
+                output["pmt_number"].append(pmt_number)
+                output["dom_number"].append(dom_number)
+                output["pmt_dir_x"].append(pmt_dir_x)
+                output["pmt_dir_y"].append(pmt_dir_y)
+                output["pmt_dir_z"].append(pmt_dir_z)
+                output["dom_type"].append(dom_type)
+                output["truth_flag"].append(truth_flag)
+
+        return output