diff --git a/km3io/__init__.py b/km3io/__init__.py index e7f2ff8cf4ef3220c973536955dc0072ceb8f0f9..769a772643f3f2c95e3b5f318a035575afe55ad8 100644 --- a/km3io/__init__.py +++ b/km3io/__init__.py @@ -4,3 +4,4 @@ version = get_distribution(__name__).version from .offline import OfflineReader from .daq import DAQReader +from .gseagen import GSGReader diff --git a/km3io/gseagen.py b/km3io/gseagen.py new file mode 100644 index 0000000000000000000000000000000000000000..01e3a0fd1d2ea7acfdef888aef6e94e56417c741 --- /dev/null +++ b/km3io/gseagen.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python +# coding=utf-8 +# Filename: gseagen.py +# Author: Johannes Schumann <jschumann@km3net.de> + +import uproot +import numpy as np +import warnings +from .tools import Branch, BranchMapper, cached_property +MAIN_TREE_NAME = "Events" + + +class GSGReader: + """reader for gSeaGen ROOT files""" + def __init__(self, file_path=None, fobj=None): + """ GSGReader class is a gSeaGen ROOT file wrapper + + Parameters + ---------- + file_path : path-like object + Path to the file of interest. It can be a str or any python + path-like object that points to the file. + """ + if file_path is not None: + self._fobj = uproot.open(file_path) + else: + self._fobj = fobj + + @cached_property + def header(self): + header_key = 'Header' + if header_key in self._fobj: + header = {} + for k, v in self._fobj[header_key].items(): + v = v.array()[0] + if isinstance(v, bytes): + try: + v = v.decode('utf-8') + except UnicodeDecodeError: + pass + header[k.decode("utf-8")] = v + return header + else: + warnings.warn("Your file header has an unsupported format") + + @cached_property + def events(self): + gseagen_events_mapper = BranchMapper(name="Events", + key="Events", + extra={}, + exclude={}, + update={}, + attrparser=lambda x: x, + flat=True) + return Branch(self._fobj, gseagen_events_mapper) diff --git a/tests/samples/gseagen.root b/tests/samples/gseagen.root new file mode 100644 index 0000000000000000000000000000000000000000..6fcfed703f65fab07031bd69ce5a1ef24c23ecd4 Binary files /dev/null and b/tests/samples/gseagen.root differ diff --git a/tests/test_gseagen.py b/tests/test_gseagen.py new file mode 100644 index 0000000000000000000000000000000000000000..84e523b76132620ca288f5d2ac7f7af84c910010 --- /dev/null +++ b/tests/test_gseagen.py @@ -0,0 +1,149 @@ +import os +import re +import unittest + +from km3io.gseagen import GSGReader + +SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "samples") +GSG_FILE = os.path.join(SAMPLES_DIR, "gseagen.root") +GSG_READER = GSGReader(GSG_FILE) + + +class TestGSGHeader(unittest.TestCase): + def setUp(self): + self.header = GSG_READER.header + + def test_str_byte_type(self): + assert isinstance(self.header['gSeaGenVer'], str) + assert isinstance(self.header['GenieVer'], str) + assert isinstance(self.header['gSeaGenVer'], str) + assert isinstance(self.header['InpXSecFile'], str) + assert isinstance(self.header['Flux1'], str) + assert isinstance(self.header['Flux2'], str) + + def test_values(self): + assert self.header["RunNu"] == 1 + assert self.header["RanSeed"] == 3662074 + self.assertAlmostEqual(self.header["NTot"], 1000.) + self.assertAlmostEqual(self.header["EvMin"], 5.) + self.assertAlmostEqual(self.header["EvMax"], 50.) + self.assertAlmostEqual(self.header["CtMin"], -1.) + self.assertAlmostEqual(self.header["CtMax"], 1.) + self.assertAlmostEqual(self.header["Alpha"], 1.4) + assert self.header["NBin"] == 1 + self.assertAlmostEqual(self.header["Can1"], 0.0) + self.assertAlmostEqual(self.header["Can2"], 476.5) + self.assertAlmostEqual(self.header["Can3"], 403.4) + self.assertAlmostEqual(self.header["OriginX"], 0.0) + self.assertAlmostEqual(self.header["OriginY"], 0.0) + self.assertAlmostEqual(self.header["OriginZ"], 0.0) + self.assertAlmostEqual(self.header["HRock"], 93.03036681) + self.assertAlmostEqual(self.header["HSeaWater"], 684.39143353) + self.assertAlmostEqual(self.header["RVol"], 611.29143353) + self.assertAlmostEqual(self.header["HBedRock"], 0.0) + self.assertAlmostEqual(self.header["NAbsLength"], 3.0) + self.assertAlmostEqual(self.header["AbsLength"], 93.33) + self.assertAlmostEqual(self.header["SiteDepth"], 2425.0) + self.assertAlmostEqual(self.header["SiteLatitude"], 0.747) + self.assertAlmostEqual(self.header["SiteLongitude"], 0.10763) + self.assertAlmostEqual(self.header["SeaBottomRadius"], 6368000.) + assert round(self.header["GlobalGenWeight"]-6.26910765e+08, 0) == 0 + self.assertAlmostEqual(self.header["RhoSW"], 1.03975) + self.assertAlmostEqual(self.header["RhoSR"], 2.65) + self.assertAlmostEqual(self.header["TGen"], 31556926.) + assert not self.header["PropMode"] + assert self.header["NNu"] == 2 + self.assertListEqual(self.header["NuList"].tolist(), [-14, 14]) + + +class TestGSGEvents(unittest.TestCase): + def setUp(self): + self.events = GSG_READER.events + + def test_event_single_values(self): + event = self.events[0] + assert 1 == event.iEvt + self.assertAlmostEqual(event.PScale, 1.0) + assert event.TargetA == 28 + assert event.TargetZ == 14 + assert event.InterId == 2 + assert event.ScattId == 3 + self.assertAlmostEqual(event.Bx, 0.21398980096236023) + self.assertAlmostEqual(event.By, 0.7736389792036026) + assert event.VerInCan == 0 + self.assertAlmostEqual(event.WaterXSec, 9.814545541159586e-42) + self.assertAlmostEqual(event.WaterIntLen, 169191612204.44382) + self.assertAlmostEqual(event.PEarth, 0.9998206835107583) + self.assertAlmostEqual(event.ColumnDepth, 31352421806833.758) + self.assertAlmostEqual(event.XSecMean, 9.49810757204515e-42) + self.assertAlmostEqual(event.Agen, 934842.6115446005) + self.assertAlmostEqual(event.WGen, 480360599.4207591) + self.assertAlmostEqual(event.WEvt, 5.629637364719835) + self.assertAlmostEqual(event.E_nu, 29.210801854810974) + assert event.Pdg_nu == -14 + self.assertAlmostEqual(event.Vx_nu, 202.86806811399845) + self.assertAlmostEqual(event.Vy_nu, 226.64032038584787) + self.assertAlmostEqual(event.Vz_nu, -0.5526929982180058) + self.assertAlmostEqual(event.Dx_nu, -0.7554430076659393) + self.assertAlmostEqual(event.Dy_nu, 0.3034298552575011) + self.assertAlmostEqual(event.Dz_nu, 0.5807204018346965) + self.assertAlmostEqual(event.T_nu, 0.0) + self.assertAlmostEqual(event.E_pl, 6.648776215872709) + assert event.Pdg_pl == -13 + self.assertAlmostEqual(event.Vx_pl, 202.86806811399845) + self.assertAlmostEqual(event.Vy_pl, 226.64032038584787) + self.assertAlmostEqual(event.Vz_pl, -0.5526929982180058) + self.assertAlmostEqual(event.Dx_pl, -0.8781745795144783) + self.assertAlmostEqual(event.Dy_pl, 0.2818890466120743) + self.assertAlmostEqual(event.Dz_pl, 0.3864556550171106) + self.assertAlmostEqual(event.T_pl, 0.0) + assert event.NTracks == 1 + + def test_event_list_values(self): + event = self.events[1] + self.assertListEqual(event.Id_tr.tolist(), [4, 5, 10, 11, 12]) + self.assertListEqual(event.Pdg_tr.tolist(), [22, -13, 2112, -211, 111]) + [ + self.assertAlmostEqual(x, y) for x, y in zip( + event.E_tr, + [0.00618, 4.88912206, 2.33667201, 1.0022909, 1.17186997]) + ] + [ + self.assertAlmostEqual(x, y) + for x, y in zip(event.Vx_tr, [ + -337.67895799, -337.67895799, -337.67895799, -337.67895799, + -337.67895799 + ]) + ] + [ + self.assertAlmostEqual(x, y) + for x, y in zip(event.Vy_tr, [ + -203.90999969, -203.90999969, -203.90999969, -203.90999969, + -203.90999969 + ]) + ] + [ + self.assertAlmostEqual(x, y) + for x, y in zip(event.Vz_tr, [ + 416.08845294, 416.08845294, 416.08845294, 416.08845294, + 416.08845294 + ]) + ] + [ + self.assertAlmostEqual(x, y) + for x, y in zip(event.Dx_tr, [ + 0.06766196, -0.63563065, -0.70627586, -0.76364544, -0.80562216 + ]) + ] + [ + self.assertAlmostEqual(x, y) for x, y in zip( + event.Dy_tr, + [0.33938809, -0.4846643, 0.50569058, -0.04136113, 0.10913917]) + ] + [ + self.assertAlmostEqual(x, y) + for x, y in zip(event.Dz_tr, [ + -0.93820978, -0.6008945, -0.49543056, -0.64430963, -0.58228994 + ]) + ] + [self.assertAlmostEqual(x, y) for x, y in zip(event.T_tr, 5 * [0.])]