Add load and save methods to event class

README.md

@@ -1,5 +1,5 @@
 ## icecube_data_reader
 
-[![codecov](https://codecov.io/github/nu-isance/icecube_data_reader/graph/badge.svg?token=G1SR407LQU)](https://codecov.io/github/nu-isance/icecube_data_reader)
+[![codecov](https://codecov.io/github/nu-isance/icecube_data_reader/branch/master/graph/badge.svg?token=G1SR407LQU)](https://codecov.io/github/nu-isance/icecube_data_reader)
 
 A collection of classes to load IceCube track data. Not affiliated with The IceCube Collaboration.

icecube_data_reader/event_types.py

@@ -12,34 +12,48 @@ def __repr__(self):
 
     # Construct to add str(EventType())
 
-
 class EventType:
     def __str__(self):
         return self._str
 
+    def __eq__(self, other):
+        return self.S == other.S
+
 
+@dataclass(eq=False)
 class IC40(EventType, metaclass=Meta):
     _str = "IC40"
+    S = 1
 
 
+@dataclass(eq=False)
 class IC59(EventType, metaclass=Meta):
     _str = "IC59"
+    S = 2
 
 
+@dataclass(eq=False)
 class IC79(EventType, metaclass=Meta):
     _str = "IC79"
+    S = 3
 
 
+@dataclass(eq=False)
 class IC86(EventType, metaclass=Meta):
     _str = "IC86"
+    S = 4
 
 
+@dataclass(eq=False)
 class IC86_I(EventType, metaclass=Meta):
     _str = "IC86_I"
+    S = 5
 
 
+@dataclass(eq=False)
 class IC86_II(EventType, metaclass=Meta):
     _str = "IC86_II"
+    S = 6
 
 
 # And what have they ever given us in return...
@@ -66,3 +80,60 @@ class DR2:
 @dataclass
 class DR1:
     available_irfs = [IC40, IC59, IC79, IC86_I, IC86_II]
+
+class Refrigerator:
+    """Collect all event types"""
+
+    detectors = [IC40, IC59, IC79, IC86, IC86_I, IC86_II]
+
+    '''
+    @classmethod
+    def python2dm(cls, python):
+        """Returns EventType corresponding to python event-type string"""
+
+        for dm in cls.detectors:
+            if dm.P == python:
+                return dm
+        else:
+            raise ValueError(f"No detector {python} available.")
+    '''
+    '''
+    @classmethod
+    def stan2dm(cls, stan):
+        """Returns EventType corresponding to stan event-type"""
+
+        for dm in cls.detectors:
+            if dm.S == stan:
+                return dm
+        else:
+            raise ValueError(f"No detector {stan} available.")
+    '''
+    @classmethod
+    def int2str(cls, int_):
+        """Returns python event-type string corresponding to integer event-type"""
+
+        for dm in cls.detectors:
+            if int_ == dm._str:
+                return dm._str
+        else:
+            raise ValueError(f"No detector {int_} available.")
+
+    @classmethod
+    def str2int(cls, str_):
+        """Returns stan event-type corresponding to python event-type string"""
+
+        for dm in cls.detectors:
+            if str_ == str(dm):
+                return dm.S
+        else:
+            raise ValueError(f"No detector {str_} available.")
+
+    @classmethod
+    def int2dm(cls, int_):
+
+        for dm in cls.detectors:
+            if int_ == dm.S:
+                return dm
+        else:
+            raise ValueError(f"No detector {int_} available.")
+

icecube_data_reader/events.py

@@ -2,7 +2,7 @@
 Class to organise IceCube event data
 """
 
-from abc import ABC
+from abc import ABC, abstractmethod
 import os
 from os.path import join
 from pathlib import Path
@@ -11,13 +11,15 @@
 from astropy.coordinates import SkyCoord
 from astropy import units as u
 from astropy.time import Time
-
-import logging
+import h5py
+from time import time as thyme
 
 from icecube_data_reader.downloader import data_directory, I3_14, available_datasets, IceCubeData
-from icecube_data_reader.event_types import IC40, IC59, IC79, IC86, suffixes
+from icecube_data_reader.event_types import IC40, IC59, IC79, IC86, suffixes, EventType, Refrigerator
 from icecube_data_reader.lifetime import LifeTime
 
+from typing import Self
+import logging
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -50,16 +52,42 @@ def ra(self):
     @property
     def dec(self):
         return self._dec
+
+    @property
+    def unit_vector(self):
+        return self._unit_vector
 
     @property
-    def type(self):
-        return self._type
+    def event_type(self):
+        return self._event_type
+
+    @property
+    def int_event_type(self):
+        return self._int_event_type
+
 
     @property
     def N(self):
-        return self.type.size
+        return self.event_type.size
 
+    @u.quantity_input
     def apply_energy_cut(self, Emin: u.GeV, Emax: u.GeV = np.inf * u.GeV):
+        """Select events based on (reconstructed) energy
+
+        :param Emin: Minimum energy
+        :type Emin: u.GeV
+        :param Emax: Maximum energy, defaults to np.inf*u.GeV
+        :type Emax: u.GeV, optional
+        """        
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_file():
+        pass
+
+    @abstractmethod
+    def to_file():
         pass
 
 
@@ -79,13 +107,14 @@ def __init__(
         self,
         energy: u.GeV,
         coord: SkyCoord,
-        type: np.ndarray,
+        event_type: np.ndarray,
         ang_err: u.deg,
         mjd: Time,
     ):
         self._energy = energy
         self._coord = coord
-        self._type = type
+        self._event_type = event_type
+        self._int_event_type = np.array([_.S for _ in event_type])
         self._ang_err = ang_err
         self._mjd = mjd
         self._coord.representation_type = "cartesian"
@@ -129,12 +158,127 @@ def select(self, mask: npt.NDArray[np.bool_]):
         self._energy = self._energy[mask]
         self._coord = self._coord[mask]
         self._unit_vector = self._unit_vector[mask]
-        self._type = self._type[mask]
+        self._event_type = self._event_type[mask]
+        self._int_event_type = self._int_event_type[mask]
         self._ang_err = self._ang_err[mask]
         self._mjd = self._mjd[mask]
 
+    def to_file(
+            self,
+            path: Path,
+            append: bool = False,
+            group_name: str | None = None,
+            overwrite: bool = False
+        ) -> Path:
+        """Write events to file.
+        Keyworded arguments control behaviour with existing files.
+        If not overwrite, but `path` exists, append a timestamp to the file name.
+
+        :param path: File path
+        :type path: Path
+        :param append: If true, append to existing file, defaults to False
+        :type append: bool, optional
+        :param group_name: If provided, create new group in path and write events there, defaults to None
+        :type group_name: str | None, optional
+        :param overwrite: If true, overwrite existing path, defaults to False
+        :type overwrite: bool, optional
+
+        :return: path Object
+        :rtype: Path
+        """
+
+        self._file_keys = ["energy", "unit_vector", "event_type", "ang_err", "mjd"]
+        self._file_values = [
+            self.energy.to(u.GeV).value,
+            self.unit_vector,
+            self.int_event_type,
+            self.ang_err.to(u.deg).value,
+            self.mjd.mjd,
+        ]
+
+        if append:
+            with h5py.File(path, "r+") as f:
+                if group_name is None:
+                    event_folder = f.create_group("events")
+                else:
+                    event_folder = f.create_group(group_name)
+
+                for key, value in zip(self._file_keys, self._file_values):
+                    event_folder.create_dataset(key, data=value)
+
+        else:
+            dirname = os.path.dirname(path)
+            filename = os.path.basename(path)
+            if dirname:
+                if not os.path.exists(dirname):
+                    logger.warning(
+                        f"{dirname} does not exist, saving instead to {os.getcwd()}"
+                    )
+                    dirname = os.getcwd()
+            else:
+                dirname = os.getcwd()
+            path = Path(dirname) / Path(filename)
+            if os.path.exists(filename) and not overwrite:
+                logger.warning(f"File {filename} already exists.")
+                file = os.path.splitext(filename)[0]
+                ext = os.path.splitext(filename)[1]
+                file += f"_{int(thyme())}"
+                filename = file + ext
+
+            path = Path(dirname) / Path(filename)
+
+            with h5py.File(path, "w") as f:
+                if group_name is None:
+                    event_folder = f.create_group("events")
+                else:
+                    event_folder = f.create_group(group_name)
+
+                for key, value in zip(self._file_keys, self._file_values):
+                    event_folder.create_dataset(key, data=value)
+        return path
+
+    @classmethod
+    def from_file(
+        cls,
+        filename: Path,
+        group_name:str = None,
+    ) -> Self:
+        """Load events from .h5 file
+
+        :param filename: File to load events from
+        :type filename: Path
+        :param group_name: Name of events group, if provided when writing to file, defaults to None
+        :type group_name: str, optional
+        """        
+        with h5py.File(filename, "r") as f:
+            if group_name is None:
+                events_folder = f["events"]
+            else:
+                events_folder = f[group_name]
+
+            energy = events_folder["energy"][()] * u.GeV
+            uv = events_folder["unit_vector"][()]
+            int_event_type = events_folder["event_type"][()]
+            event_type = np.array([Refrigerator.int2dm(_) for _ in int_event_type])
+            ang_err = events_folder["ang_err"][()] * u.deg
+
+            # For backwards compatibility
+            try:
+                mjd = events_folder["mjd"][()]
+            except KeyError:
+                mjd = [99.0] * len(energy)
+
+        coord = SkyCoord(
+            uv.T[0], uv.T[1], uv.T[2], representation_type="cartesian", frame="icrs"
+        )
+        mjd = Time(mjd, format="mjd")
+        coord.representation_type = "spherical"
+        events = cls(energy, coord, event_type, ang_err, mjd)
+
+        return events
+
     @classmethod
-    def from_event_files(cls, *seasons):
+    def from_event_files(cls, *seasons: EventType) -> Self:
         """
         Load data of provided seasons.
         If none are provided, use all.
@@ -171,13 +315,13 @@ def from_event_files(cls, *seasons):
         dec = []
         mjd = []
         ang_err = []
-        type = []
+        event_type = []
 
-        def _append_data(s):
+        def _append_data(s, suffering: str = ""):
             data = np.loadtxt(
                 join(
                     data_directory,
-                    f"{str(Path(directory) / Path(sub_directory))}/events/{s}_exp.csv",
+                    f"{str(Path(directory) / Path(sub_directory))}/events/{str(s)+suffering}_exp.csv",
                 )
             )
 
@@ -186,12 +330,12 @@ def _append_data(s):
             ra.append(data[:, cls.ra_])
             dec.append(data[:, cls.dec_])
             ang_err.append(data[:, cls.ang_err_])
-            type.append(len(data[:, cls.energy_]) * [s])
+            event_type.append(len(data[:, cls.energy_]) * [s])
 
         for s in seasons:
             if s == IC86:
                 for suffering in suffixes:
-                    _append_data(str(s) + suffering)
+                    _append_data(s, suffering)
             else:
                 _append_data(s)
 
@@ -200,10 +344,10 @@ def _append_data(s):
         ra = np.concatenate(ra) << u.deg
         dec = np.concatenate(dec) << u.deg
         ang_err = np.concatenate(ang_err) << u.deg
-        type = np.concatenate(type)
+        event_type = np.concatenate(event_type)
         coord = SkyCoord(ra=ra, dec=dec, frame="icrs")
 
-        events = cls(energy, coord, type, ang_err, mjd)
+        events = cls(energy, coord, event_type, ang_err, mjd)
         events._ra = ra
         events._dec = dec
 

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "numpy",
     "astropy",
     "scipy",
+    "h5py",
     "tqdm",
     "requests",
     "requests_cache",