diff --git a/README.md b/README.md index 2635fd2..e142632 100644 --- a/README.md +++ b/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. diff --git a/icecube_data_reader/event_types.py b/icecube_data_reader/event_types.py index 3afcc57..d4f9774 100644 --- a/icecube_data_reader/event_types.py +++ b/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.") + diff --git a/icecube_data_reader/events.py b/icecube_data_reader/events.py index 23ba483..77854c6 100644 --- a/icecube_data_reader/events.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 7d5b04b..09ff39e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "numpy", "astropy", "scipy", + "h5py", "tqdm", "requests", "requests_cache", diff --git a/tests/test_events.py b/tests/test_events.py index 772e313..5d3c607 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -1,6 +1,58 @@ from icecube_data_reader.events import IceTrackDR2Events +from icecube_data_reader.event_types import IC40 +from astropy import units as u +import numpy as np +from pathlib import Path +import os + +import pytest def test_event_number(): events = IceTrackDR2Events.from_event_files() assert events.N == 1643355 + +@pytest.fixture +def test_saving(output_directory): + events = IceTrackDR2Events.from_event_files(IC40) + path = events.to_file(Path(output_directory) / "ic40_events.h5") + return (path, events) + +def test_overwriting(output_directory, test_saving): + path, events = test_saving + time_appended = events.to_file(path) + assert os.path.basename(time_appended) in os.listdir(output_directory) + + +def test_loading(test_saving): + path, ic40 = test_saving + loaded = IceTrackDR2Events.from_file(path) + + assert np.all(np.isclose(loaded.energy.to_value(u.GeV), ic40.energy.to_value(u.GeV))) + +def test_selecting(test_saving): + idx = 5 + _, events = test_saving + N = events.N + mask = np.zeros(N, dtype=bool) + mask[idx] = True + + e = events.energy[idx].to_value(u.GeV) + et = events.int_event_type[idx] + events.select(mask) + + assert pytest.approx(events.energy[0].to_value(u.GeV)) == e + assert et == events.int_event_type[0] + + +def test_erroneous_selecting(test_saving): + idx = 5 + events = IceTrackDR2Events.from_event_files(IC40) + N = events.N + mask = np.zeros(N, dtype=bool) + mask[idx] = True + + with pytest.raises(ValueError, match="Mask needs to be of the same length as N."): + mask = np.zeros(N+1, dtype=bool) + mask[1] = 1 + events.select(mask) \ No newline at end of file diff --git a/tests/test_refrigerator.py b/tests/test_refrigerator.py new file mode 100644 index 0000000..52eee2f --- /dev/null +++ b/tests/test_refrigerator.py @@ -0,0 +1,21 @@ +from icecube_data_reader.event_types import ( + Refrigerator, DR2 +) + +import pytest + +def test_int2dm(): + for int_, dm in enumerate(DR2.available_irfs, 1): + assert Refrigerator.int2dm(int_) == dm + + with pytest.raises(ValueError, match="No detector 0 available."): + Refrigerator.int2dm(0) + +def test_str2int(): + for int_, dm in enumerate(DR2.available_irfs, 1): + assert Refrigerator.str2int(str(dm)) == dm.S + +def test_dm_equality(): + for dm in DR2.available_irfs: + assert dm == dm + assert dm == dm() \ No newline at end of file