Merge branch 'detector-geometries' into 'master'

Detector geometries

See merge request !6

km3buu/geometry.py

+# Filename: geometry.py
+"""
+Detector geometry related functionalities
+"""
+
+__author__ = "Johannes Schumann"
+__copyright__ = "Copyright 2021, Johannes Schumann and the KM3NeT collaboration."
+__credits__ = []
+__license__ = "MIT"
+__maintainer__ = "Johannes Schumann"
+__email__ = "jschumann@km3net.de"
+__status__ = "Development"
+
+import numpy as np
+from abc import ABC, abstractmethod
+
+
+class DetectorVolume(ABC):
+    """
+    Detector geometry class
+    """
+    def __init__(self):
+        self._volume = -1.0
+        self._coord_origin = (0., 0., 0.)
+
+    @abstractmethod
+    def random_pos(self):
+        """
+        Generate a random position in the detector volume based on a uniform
+        event distribution
+
+        Returns
+        -------
+        tuple [m] (x, y, z)
+        """
+        pass
+
+    @abstractmethod
+    def header_entry(self):
+        """
+        Returns the header information for the detector volume geometry
+
+        Returns
+        -------
+        tuple (header_key, header_information)
+        """
+        pass
+
+    @property
+    def volume(self):
+        """
+        Detector volume
+
+        Returns
+        -------
+        float [m^3] 
+        """
+        return self._volume
+
+    @property
+    def coord_origin(self):
+        """
+        Coordinate origin
+
+        Returns
+        -------
+        tuple [m] (x, y, z)
+        """
+        return self._coord_origin
+
+
+class CanVolume(DetectorVolume):
+    """
+    Cylindrical detector geometry
+
+    Parameters
+    ----------
+    radius: float [m] (default: 403.5)
+        Cylinder radius given in metres
+    zmin: float [m] (default: 0.0)
+        Cylinder bottom z position
+    zmax: float [m] (default: 476.5)
+        Cylinder top z position
+    """
+    def __init__(self, radius=403.4, zmin=0.0, zmax=476.5):
+        super().__init__()
+        self._radius = radius
+        self._zmin = zmin
+        self._zmax = zmax
+        self._volume = self._calc_volume()
+
+    def _calc_volume(self):
+        return np.pi * (self._zmax - self._zmin) * np.power(self._radius, 2)
+
+    def random_pos(self):
+        r = self._radius * np.sqrt(np.random.uniform(0, 1))
+        phi = np.random.uniform(0, 2 * np.pi)
+        pos_x = r * np.cos(phi)
+        pos_y = r * np.sin(phi)
+        pos_z = np.random.uniform(self._zmin, self._zmax)
+        return (pos_x, pos_y, pos_z)
+
+    def header_entry(self):
+        key = "can"
+        value = "{} {} {}".format(self._zmin, self._zmax, self._radius)
+        return key, value
+
+
+class SphericalVolume(DetectorVolume):
+    """
+    Spherical detector geometry
+    
+    Parameters
+    ----------
+    radius: float [m]
+        The radius of the sphere
+    center: tuple [m]
+        Coordinate center of the sphere
+        (x, y, z)
+    """
+    def __init__(self, radius, coord_origin=(0, 0, 0)):
+        super().__init__()
+        self._radius = radius
+        self._coord_origin = coord_origin
+        self._volume = self._calc_volume()
+
+    def _calc_volume(self):
+        return 4 / 3 * np.pi * np.power(self._radius, 3)
+
+    def random_pos(self):
+        r = np.power(self._radius, 1 / 3)
+        phi = np.random.uniform(0, np.pi)
+        cosTheta = np.random.uniform(-1, 1)
+        pos_x = r * np.cos(phi) * np.sqrt(1 - np.power(cosTheta, 2))
+        pos_y = r * np.sin(phi) * np.sqrt(1 - np.power(cosTheta, 2))
+        pos_z = r * cosTheta
+        pos = (pos_x, pos_y, pos_z)
+        return tuple(np.add(self._coord_origin, pos))
+
+    def header_entry(self):
+        key = "sphere"
+        value = "radius: {} center_x: {} center_y: {} center_z: {}".format(
+            self._radius, self._coord_origin[0], self._coord_origin[1],
+            self._coord_origin[2])
+        return key, value

km3buu/output.py

@@ -27,6 +27,7 @@ import mendeleev
 from datetime import datetime
 
 from .jobcard import Jobcard, read_jobcard, PDGID_LOOKUP
+from .geometry import DetectorVolume, CanVolume
 from .config import Config, read_default_media_compositions
 from .__version__ import version
 
@@ -91,29 +92,16 @@ ROOTTUPLE_KEY = "RootTuple"
 
 EMPTY_KM3NET_HEADER_DICT = {
     "start_run": "0",
-    "XSecFile": "",
     "drawing": "volume",
-    "detector": "",
     "depth": "2475.0",
-    "muon_desc_file": "",
     "target": "",
-    "cut_primary": "0 0 0 0",
-    "cut_seamuon": "0 0 0 0",
-    "cut_in": "0 0 0 0",
     "cut_nu": "0 0 0 0",
     "spectrum": "0",
     "can": "0 0 0",
     "flux": "0 0 0",
-    "fixedcan": "0 0 0 0 0",
-    "genvol": "0 0 0 0 0",
     "coord_origin": "0 0 0",
-    "genhencut": "0 0",
     "norma": "0 0",
-    "livetime": "0 0",
-    "seabottom": "0",
-    "DAQ": "0",
     "tgen": "0",
-    "primary": "0",
     "simul": ""
 }
 
@@ -360,7 +348,7 @@ class GiBUUOutput:
 
 def write_detector_file(gibuu_output,
                         ofile="gibuu.offline.root",
-                        can=(0, 476.5, 403.4),
+                        geometry=CanVolume(),
                         livetime=3.156e7,
                         propagate_tau=True):  # pragma: no cover
     """
@@ -372,12 +360,13 @@ def write_detector_file(gibuu_output,
         Output object which wraps the information from the GiBUU output files
     ofile: str
         Output filename
-    can: tuple
-        The can dimensions which are used to distribute the events 
-        (z_min, z_max, radius)
+    geometry: DetectorVolume
+        The detector geometry which should be used
     livetime: float
         The data livetime
     """
+    if not isinstance(geometry, DetectorVolume):
+        raise TypeError("Geometry needs to be a DetectorVolume")
 
     evt = ROOT.Evt()
     outfile = ROOT.TFile.Open(ofile, "RECREATE")
@@ -439,28 +428,29 @@ def write_detector_file(gibuu_output,
 
     media = read_default_media_compositions()
     density = media["SeaWater"]["density"]
-    symbol = mendeleev.element(gibuu_output.Z).symbol
-    target = media["SeaWater"]["elements"][symbol]
+    element = mendeleev.element(gibuu_output.Z)
+    target = media["SeaWater"]["elements"][element.symbol]
     target_density = 1e3 * density * target[1]
     targets_per_volume = target_density * (1e3 * constants.Avogadro /
                                            target[0].atomic_weight)
 
-    can_volume = np.pi * (can[1] - can[0]) * np.power(can[2], 2)
-    w2 = gibuu_output.w2weights(can_volume, targets_per_volume, 4 * np.pi)
+    w2 = gibuu_output.w2weights(geometry.volume, targets_per_volume, 4 * np.pi)
 
     head = ROOT.Head()
     header_dct = EMPTY_KM3NET_HEADER_DICT.copy()
 
-    timestamp = datetime.now()
-    header_dct["simul"] = "KM3BUU {} {}".format(
-        version, timestamp.strftime("%Y%m%d %H%M%S"))
-    header_dct["can"] = "{:.1f} {:.1f} {:.1f}".format(*can)
-    header_dct["tgen"] = "{:.1f}".format(livetime)
+    header_dct["target"] = element.name
+    key, value = geometry.header_entry()
+    header_dct[key] = value
+    header_dct["coord_origin"] = "{} {} {}".format(*geometry.coord_origin)
     header_dct["flux"] = "{:d} 0 0".format(nu_type)
-    header_dct["genvol"] = "0 {:.1f} {:.1f} {:.1f} {:d}".format(
-        can[1], can[2], can_volume, gibuu_output.generated_events)
     header_dct["cut_nu"] = "{:.2f} {:.2f} -1 1".format(gibuu_output.energy_min,
                                                        gibuu_output.energy_max)
+    header_dct["tgen"] = "{:.1f}".format(livetime)
+    header_dct["norma"] = "0 {}".format(gibuu_output.generated_events)
+    timestamp = datetime.now()
+    header_dct["simul"] = "KM3BUU {} {}".format(
+        version, timestamp.strftime("%Y%m%d %H%M%S"))
 
     for k, v in header_dct.items():
         head.set_line(k, v)
@@ -471,16 +461,11 @@ def write_detector_file(gibuu_output,
         evt.id = mc_event_id
         evt.mc_run_id = mc_event_id
         # Weights
-        evt.w.push_back(can_volume)  #w1 (can volume)
+        evt.w.push_back(geometry.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)
-        pos_x = r * np.cos(phi)
-        pos_y = r * np.sin(phi)
-        pos_z = np.random.uniform(can[0], can[1])
-        vtx_pos = np.array([pos_x, pos_y, pos_z])
+        vtx_pos = np.array(geometry.random_pos())
         # Direction
         phi = np.random.uniform(0, 2 * np.pi)
         cos_theta = np.random.uniform(-1, 1)

km3buu/tests/test_geometry.py

+#!/usr/bin/env python
+# coding=utf-8
+# Filename: test_output.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 unittest
+
+from km3buu.geometry import *
+
+
+class TestGeneralGeometry(unittest.TestCase):
+    def test_abstract_init(self):
+        with self.assertRaises(TypeError) as ctx:
+            d = DetectorVolume()
+
+
+class TestSphere(unittest.TestCase):
+    def setUp(self):
+        self.detector_geometry = SphericalVolume(20, (2, 2, 2))
+
+    def test_volume(self):
+        volume = self.detector_geometry.volume
+        self.assertAlmostEqual(volume, 33510.32, 2)
+
+    def test_random_pos(self):
+        for i in range(50):
+            pos = self.detector_geometry.random_pos()
+            assert pos[0] > -18.0
+            assert pos[1] > -18.0
+            assert pos[2] > -18.0
+            assert pos[0] < 22.0
+            assert pos[1] < 22.0
+            assert pos[2] < 22.0
+
+
+class TestCan(unittest.TestCase):
+    def setUp(self):
+        self.detector_geometry = CanVolume()
+
+    def test_volume(self):
+        volume = self.detector_geometry.volume
+        self.assertAlmostEqual(volume, 243604084.28, 2)