Newer
Older
Daniel Guderian
committed
import os
from unittest import TestCase
import tempfile
import numpy as np
import h5py
import orcasong.core
import orcasong.extractors as extractors
__author__ = "Daniel Guderian"
test_dir = os.path.dirname(os.path.realpath(__file__))
Daniel Guderian
committed
NEUTRINO_FILE = os.path.join(
test_dir,
"data",
"mcv6.0.gsg_anti-muon-NC_1-500GeV.km3sim.jterbr00008350.jorcarec.jsh.aanet.898.root.h5",
)
DET_FILE_NEUTRINO = os.path.join(test_dir, "data", "KM3NeT_00000049_20200707.detx")
Daniel Guderian
committed
Daniel Guderian
committed
NOT_FULLY_RECONSTRUCTED_FILE = os.path.join(
test_dir,
"data",
"mcv6.0.gsg_anti-muon-NC_1-500GeV.km3sim.jterbr00008349.jorcarec.jsh.aanet.897.root.h5",
)
Daniel Guderian
committed
Daniel Guderian
committed
class TestStdRecoExtractor(TestCase):
""" Assert that the neutrino info is extracted correctly. File has 18 events. """
Daniel Guderian
committed
@classmethod
def setUpClass(cls):
Daniel Guderian
committed
# normal case, with complete recos
Daniel Guderian
committed
cls.proc = orcasong.core.FileGraph(
max_n_hits=3,
time_window=[0, 50],
hit_infos=["pos_z", "time", "channel_id"],
extractor=extractors.get_neutrino_mc_info_extr(NEUTRINO_FILE,1),
Daniel Guderian
committed
det_file=DET_FILE_NEUTRINO,
add_t0=True,
Daniel Guderian
committed
)
cls.tmpdir = tempfile.TemporaryDirectory()
cls.outfile = os.path.join(cls.tmpdir.name, "binned.h5")
cls.proc.run(infile=NEUTRINO_FILE, outfile=cls.outfile)
cls.f = h5py.File(cls.outfile, "r")
Daniel Guderian
committed
cls.reco_names_in_original_file = ["best_jshower", "best_jmuon"]
cls.reco_names = ["jshower", "jmuon"]
cls.quantities_to_test = ["dir_z", "dir_x", "pos_x"]
cls.mc_quantities_to_test = ["dir_z", "dir_x", "vertex_pos_x"]
# case where a few recos from jmuon are missing
cls.proc = orcasong.core.FileGraph(
max_n_hits=3,
time_window=[0, 50],
hit_infos=["pos_z", "time", "channel_id"],
Daniel Guderian
committed
extractor=extractors.get_neutrino_mc_info_extr(
Daniel Guderian
committed
),
det_file=DET_FILE_NEUTRINO,
add_t0=True,
keep_event_info=True,
Daniel Guderian
committed
cls.outfile_not_fully_rec = os.path.join(
cls.tmpdir.name, "binned_not_fully_rec.h5"
)
cls.proc.run(
infile=NOT_FULLY_RECONSTRUCTED_FILE, outfile=cls.outfile_not_fully_rec
)
cls.f_not_fully_rec_original = h5py.File(cls.outfile_not_fully_rec, "r")
Daniel Guderian
committed
@classmethod
def tearDownClass(cls):
cls.f.close()
cls.tmpdir.cleanup()
def test_keys(self):
self.assertSetEqual(
set(self.f.keys()),
{
"_i_event_info",
"_i_group_info",
"_i_y",
"event_info",
"group_info",
"x",
"x_indices",
"y",
},
)
Daniel Guderian
committed
# the y from the processed file
Daniel Guderian
committed
y = self.f["y"][()]
Daniel Guderian
committed
# the info from the original file
orig_info = h5py.File(NEUTRINO_FILE, "r")
orig_reco_info = orig_info["reco"]
orig_mc_info = orig_info["mc_tracks"]
Daniel Guderian
committed
# test a few reco parameter
for i in range(len(self.reco_names_in_original_file)):
for j in self.quantities_to_test:
Daniel Guderian
committed
orig_reco_values = orig_reco_info[self.reco_names_in_original_file[i]][
j
]
assert np.allclose(
orig_reco_values, y[self.reco_names[i] + "_" + j], equal_nan=True
)
# test a few mc truth parameter
Daniel Guderian
committed
orig_mc_values = orig_mc_info[self.quantities_to_test[j]][
0
] # only take fist, the primary
assert np.allclose(
orig_mc_values, y[self.mc_quantities_to_test[j]][0], equal_nan=True
)
def test_incomplete_reco_padding(self):
# make sure the x and y even in this case have the same length
x = self.f_not_fully_rec_original["x"]
y = self.f_not_fully_rec_original["y"]