Merge branch 'km3net-dataformat' into 'master'

KM3NeT Dataformat data writeout

See merge request !4

km3buu/config.py

@@ -28,9 +28,18 @@ CONFIG_PATH = os.path.expanduser("~/.km3buu/config")
 
 log = get_logger(__name__)
 
+GENERAL_SECTION = "General"
+KM3NET_LIB_PATH_KEY = "km3net_lib_path"
+
 GIBUU_SECTION = "GiBUU"
+GIBUU_IMAGE_PATH_KEY = "image_path"
+
 PROPOSAL_SECTION = "PROPOSAL"
-GSEAGEN_SECTION = "gSeagen"
+PROPOSAL_ITP_PATH_KEY = "itp_table_path"
+
+GSEAGEN_SECTION = "gSeaGen"
+GSEAGEN_PATH_KEY = "path"
+
 GSEAGEN_MEDIA_COMPOSITION_FILE = "MediaComposition.xml"
 
 
@@ -62,8 +71,8 @@ class Config(object):
 
     @property
     def gibuu_image_path(self):
-        key = "image_path"
-        image_path = self.get(GIBUU_SECTION, key)
+        key = GIBUU_IMAGE_PATH_KEY
+        image_path = self.get(GIBUU_SECTION, GIBUU_IMAGE_PATH_KEY)
         if image_path is None or not isfile(image_path):
             dev_path = abspath(join(dirname(__file__), os.pardir, IMAGE_NAME))
             if isfile(dev_path):
@@ -86,26 +95,34 @@ class Config(object):
 
     @gibuu_image_path.setter
     def gibuu_image_path(self, value):
-        section = "GiBUU"
-        key = "image_path"
+        section = GIBUU_SECTION
+        key = GIBUU_IMAGE_PATH_KEY
         self.set(section, key, value)
 
     @property
     def proposal_itp_tables(self):
         default_path = abspath(join(dirname(__file__), "../.tables"))
-        return self.get(PROPOSAL_SECTION, "itp_table_path", default_path)
+        return self.get(PROPOSAL_SECTION, PROPOSAL_ITP_PATH_KEY, default_path)
 
     @proposal_itp_tables.setter
     def proposal_itp_tables(self, value):
-        self.set(PROPOSAL_SECTION, "itp_table_path", value)
+        self.set(PROPOSAL_SECTION, PROPOSAL_ITP_PATH_KEY, value)
 
     @property
     def gseagen_path(self):
-        return self.get(GSEAGEN_SECTION, "path")
+        return self.get(GSEAGEN_SECTION, GSEAGEN_PATH_KEY)
 
     @gseagen_path.setter
     def gseagen_path(self, value):
-        self.set(GIBUU_SECTION, "path", value)
+        self.set(GIBUU_SECTION, GSEAGEN_PATH_KEY, value)
+
+    @property
+    def km3net_lib_path(self):
+        return self.get(GENERAL_SECTION, KM3NET_LIB_PATH_KEY, None)
+
+    @km3net_lib_path.setter
+    def km3net_lib_path(self, value):
+        self.set(GENERAL_SECTION, KM3NET_LIB_PATH_KEY, value)
 
 
 def read_media_compositions(filename):

km3buu/jobcard.py

@@ -94,13 +94,13 @@ def read_jobcard(filepath):
 
 
 def generate_neutrino_jobcard_template(
-    process,
-    flavour,
-    energy,
-    target,
-    write_events=False,
-    do_decay=False,
-    input_path=INPUT_PATH):  # pragma: no cover
+        process,
+        flavour,
+        energy,
+        target,
+        write_events=False,
+        do_decay=False,
+        input_path=INPUT_PATH):  # pragma: no cover
     """
     Generate a jobcard for neutrino interaction
 

km3buu/output.py

@@ -15,7 +15,7 @@ __status__ = "Development"
 import re
 import numpy as np
 from io import StringIO
-from os import listdir
+from os import listdir, environ
 from os.path import isfile, join, abspath
 from tempfile import TemporaryDirectory
 import awkward as ak
@@ -26,7 +26,18 @@ import scipy.constants as constants
 import mendeleev
 
 from .jobcard import Jobcard, read_jobcard, PDGID_LOOKUP
-from .config import read_default_media_compositions
+from .config import Config, read_default_media_compositions
+
+try:
+    import ROOT
+    libpath = environ.get("KM3NET_LIB")
+    if libpath is None:
+        libpath = Config().km3net_lib_path
+    if ROOT.gSystem.Load(join(libpath, "libKM3NeTROOT.so")) < 0:
+        raise ModuleNotFoundError("KM3NeT dataformat library not found!")
+except ModuleNotFoundError:
+    import warnings
+    warnings.warn("KM3NeT dataformat library was not loaded.", ImportWarning)
 
 EVENT_FILENAME = "FinalEvents.dat"
 ROOT_PERT_FILENAME = "EventOutput.Pert.[0-9]{8}.root"
@@ -298,7 +309,7 @@ def write_detector_file(gibuu_output,
                         livetime=3.156e7,
                         propagate_tau=True):  # pragma: no cover
     """
-    Convert the GiBUU output to a KM3NeT MC (AANET) file
+    Convert the GiBUU output to a KM3NeT MC (OfflineFormat) file
 
     Parameters
     ----------
@@ -312,15 +323,16 @@ def write_detector_file(gibuu_output,
     livetime: float
         The data livetime
     """
-    import aa, ROOT
 
-    aafile = ROOT.EventFile()
-    aafile.set_output(ofile)
+    evt = ROOT.Evt()
+    outfile = ROOT.TFile.Open(ofile, "RECREATE")
+    tree = ROOT.TTree("E", "KM3NeT Evt Tree")
+    tree.Branch("Evt", evt, 32000, 4)
     mc_trk_id = 0
 
     def add_particles(particle_data, pos_offset, rotation, start_mc_trk_id,
                       timestamp):
-        nonlocal aafile
+        nonlocal evt
         for i in range(len(particle_data.E)):
             trk = ROOT.Trk()
             trk.id = start_mc_trk_id + i
@@ -342,7 +354,7 @@ def write_detector_file(gibuu_output,
             trk.mother_id = 0
             trk.type = int(particle_data.barcode[i])
             trk.E = particle_data.E[i]
-            aafile.evt.mc_trks.push_back(trk)
+            evt.mc_trks.push_back(trk)
 
     is_cc = False
 
@@ -381,13 +393,13 @@ def write_detector_file(gibuu_output,
     w2 = gibuu_output.w2weights(can_volume, targets_per_volume, 4 * np.pi)
 
     for mc_event_id, event in enumerate(event_data):
-        aafile.evt.clear()
-        aafile.evt.id = mc_event_id
-        aafile.evt.mc_run_id = mc_event_id
+        evt.clear()
+        evt.id = mc_event_id
+        evt.mc_run_id = mc_event_id
         # Weights
-        aafile.evt.w.push_back(can_volume)  #w1 (can volume)
-        aafile.evt.w.push_back(w2[mc_event_id])  #w2
-        aafile.evt.w.push_back(-1.0)  #w3 (= w2*flux)
+        evt.w.push_back(can_volume)  #w1 (can volume)
+        evt.w.push_back(w2[mc_event_id])  #w2
+        evt.w.push_back(-1.0)  #w3 (= w2*flux)
         # Vertex Position
         r = can[2] * np.sqrt(np.random.uniform(0, 1))
         phi = np.random.uniform(0, 2 * np.pi)
@@ -432,7 +444,7 @@ def write_detector_file(gibuu_output,
             lep_out_trk.dir.set(*p_dir)
             lep_out_trk.E = event.lepOut_E
             lep_out_trk.t = timestamp
-            aafile.evt.mc_trks.push_back(lep_out_trk)
+            evt.mc_trks.push_back(lep_out_trk)
 
         # bjorken_y = 1.0 - float(event.lepOut_E / event.lepIn_E)
         nu_in_trk.setusr('bx', bjorkenx[mc_event_id])
@@ -440,7 +452,7 @@ def write_detector_file(gibuu_output,
         nu_in_trk.setusr('ichan', ichan)
         nu_in_trk.setusr("cc", is_cc)
 
-        aafile.evt.mc_trks.push_back(nu_in_trk)
+        evt.mc_trks.push_back(nu_in_trk)
 
         if tau_secondaries is not None:
             event_tau_sec = tau_secondaries[mc_event_id]
@@ -449,6 +461,7 @@ def write_detector_file(gibuu_output,
 
         add_particles(event, vtx_pos, R, mc_trk_id, timestamp)
 
-        aafile.write()
-
-    del aafile
+        tree.Fill()
+    outfile.Write()
+    outfile.Close()
+    del outfile

km3buu/tests/test_config.py

+#!/usr/bin/env python
+# coding=utf-8
+# Filename: test_ctrl.py
+
+__author__ = "Johannes Schumann"
+__copyright__ = "Copyright 2020, Johannes Schumann and the KM3NeT collaboration."
+__credits__ = []
+__license__ = "MIT"
+__maintainer__ = "Johannes Schumann"
+__email__ = "jschumann@km3net.de"
+__status__ = "Development"
+
+import os
+import unittest
+from unittest.mock import patch
+import numpy as np
+from km3buu.jobcard import *
+from km3buu.config import Config, read_media_compositions
+
+from tempfile import NamedTemporaryFile
+
+TEST_CONFIG = """
+[General]
+km3net_lib_path=/tmp/km3net_lib_test
+[GiBUU]
+image_path=%s
+[PROPOSAL]
+itp_table_path=/tmp/.tables
+[gSeaGen]
+path=/tmp/gseagen
+"""
+
+
+class TestConfig(unittest.TestCase):
+    def setUp(self):
+        self.cfg_tmpfile = NamedTemporaryFile(delete=False)
+        self.mock_image_file = NamedTemporaryFile(delete=False)
+        with open(self.cfg_tmpfile.name, "w") as f:
+            f.write(TEST_CONFIG % self.mock_image_file.name)
+        self.cfg = Config(self.cfg_tmpfile.name)
+
+    def tearDown(self):
+        os.remove(self.cfg_tmpfile.name)
+        os.remove(self.mock_image_file.name)
+
+    def test_general_section(self):
+        assert self.cfg.km3net_lib_path == "/tmp/km3net_lib_test"
+
+    def test_gibuu_section(self):
+        assert self.cfg.gibuu_image_path == self.mock_image_file.name
+
+    def test_proposal_section(self):
+        assert self.cfg.proposal_itp_tables == "/tmp/.tables"
+
+    def test_gseagen_path(self):
+        assert self.cfg.gseagen_path == "/tmp/gseagen"

km3buu/tests/test_output.py

@@ -15,19 +15,25 @@ from unittest.mock import patch
 import numpy as np
 import pytest
 import km3io
-from km3buu.output import *
-from os import listdir
+from os import listdir, environ
 from os.path import abspath, join, dirname
 from km3net_testdata import data_path
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
+from km3buu.output import *
+from km3buu.config import Config
+
 TESTDATA_DIR = data_path("gibuu")
 
 try:
-    import aa, ROOT
-    AANET_AVAILABLE = True
-except ImportError:
-    AANET_AVAILABLE = False
+    import ROOT
+    libpath = environ.get("KM3NET_LIB")
+    if libpath is None:
+        libpath = Config().km3net_lib_path
+    KM3NET_LIB_AVAILABLE = (ROOT.gSystem.Load(join(libpath,
+                                                   "libKM3NeTROOT.so")) >= 0)
+except ModuleNotFoundError:
+    KM3NET_LIB_AVAILABLE = False
 
 
 class TestXSection(unittest.TestCase):
@@ -73,7 +79,8 @@ class TestGiBUUOutput(unittest.TestCase):
             w2[:3], [7.63360911e+01, 3.60997502e-01, 1.13273189e+03])
 
 
-@pytest.mark.skipif(not AANET_AVAILABLE, reason="aanet required")
+@pytest.mark.skipif(not KM3NET_LIB_AVAILABLE,
+                    reason="KM3NeT dataformat required")
 class TestAANET(unittest.TestCase):
     def setUp(self):
         output = GiBUUOutput(TESTDATA_DIR)