Merge branch 'tau_propagation' into 'master'

Tau propagation

See merge request !2

km3buu/config.py

@@ -28,6 +28,9 @@ CONFIG_PATH = os.path.expanduser("~/.km3buu/config")
 
 log = get_logger(__name__)
 
+GIBUU_SECTION = "GiBUU"
+PROPOSAL_SECTION = "PROPOSAL"
+
 
 class Config(object):
     def __init__(self, config_path=CONFIG_PATH):
@@ -57,9 +60,8 @@ class Config(object):
 
     @property
     def gibuu_image_path(self):
-        section = "GiBUU"
         key = "image_path"
-        image_path = self.get(section, key)
+        image_path = self.get(GIBUU_SECTION, 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):
@@ -77,7 +79,7 @@ class Config(object):
                     default=default_dir,
                 )
                 image_path = build_image(image_dir)
-            self.set(section, key, image_path)
+            self.set(GIBUU_SECTION, key, image_path)
         return image_path
 
     @gibuu_image_path.setter
@@ -86,6 +88,15 @@ class Config(object):
         key = "image_path"
         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)
+
+    @proposal_itp_tables.setter
+    def proposal_itp_tables(self, value):
+        self.set(PROPOSAL_SECTION, "itp_table_path")
+
 
 def read_media_compositions(filename):
     """

km3buu/output.py

@@ -99,6 +99,7 @@ class GiBUUOutput:
         Parameters
         ----------
         data_dir: str
+            Path to the GiBUU output directory
         """
         if isinstance(data_dir, TemporaryDirectory):
             self._tmp_dir = data_dir
@@ -160,8 +161,18 @@ class GiBUUOutput:
         xsec = np.divide(total_flux_events * weights, deltaE * n_files)
         return xsec
 
+    def _w2weights(self, root_tupledata):
+        pass
+
     @staticmethod
     def bjorken_x(roottuple_data):
+        """
+        Calculate Bjorken x variable for the GiBUU events
+
+        Parameters
+        ----------
+            roottuple_data: awkward1.highlevel.Array                
+        """
         d = roottuple_data
         k_in = np.vstack([
             np.array(d.lepIn_E),
@@ -185,6 +196,14 @@ class GiBUUOutput:
 
     @staticmethod
     def bjorken_y(roottuple_data):
+        """
+        Calculate Bjorken y scaling variable for the GiBUU events
+        (Lab. frame)
+
+        Parameters
+        ----------
+            roottuple_data: awkward1.highlevel.Array                
+        """
         d = roottuple_data
         y = 1 - np.divide(np.array(d.lepOut_E), np.array(d.lepIn_E))
         return y
@@ -238,7 +257,8 @@ class GiBUUOutput:
 def write_detector_file(gibuu_output,
                         ofile="gibuu.aanet.root",
                         can=(0, 476.5, 403.4),
-                        livetime=3.156e7):
+                        livetime=3.156e7,
+                        propagate_tau=True):
     """
     Convert the GiBUU output to a KM3NeT MC (AANET) file
 
@@ -258,6 +278,26 @@ def write_detector_file(gibuu_output,
     aafile = ROOT.EventFile()
     aafile.set_output(ofile)
     mc_event_id = 0
+    mc_trk_id = 0
+
+    def add_particles(particle_data, rotation):
+        for i in range(len(particle_data.E)):
+            trk = ROOT.Trk()
+            trk.id = mc_trk_id
+            mc_trk_id += 1
+            mom = np.array([
+                particle_data.Px[i], particle_data.Py[i], particle_data.Pz[i]
+            ])
+            p_dir = R.apply(mom / np.linalg.norm(mom))
+            prtcl_pos_offset = np.array(
+                [particle_data.x[i], particle_data.y[i], particle_data.z[i]])
+            trk.pos.set(*np.add(vtx_pos, prtcl_pos_offset))
+            trk.dir.set(*p_dir)
+            trk.mother_id = 0
+            trk.type = int(particle_data.barcode[i])
+            trk.E = particle_data.E[i]
+            trk.t = timestamp
+            aafile.evt.mc_trks.push_back(trk)
 
     is_cc = False
 
@@ -280,6 +320,13 @@ def write_detector_file(gibuu_output,
         event_data = fobj["RootTuple"].arrays()
         bjorkenx = GiBUUOutput.bjorken_x(event_data)
         bjorkeny = GiBUUOutput.bjorken_y(event_data)
+
+        tau_secondaries = None
+        if propagate_tau and abs(nu_type) == 16:
+            from .propagation import propagate_lepton
+            tau_secondaries = propagate_lepton(event_data,
+                                               np.sign(nu_type) * 15)
+
         for i, event in enumerate(event_data):
             aafile.evt.clear()
             aafile.evt.id = mc_event_id
@@ -308,7 +355,8 @@ def write_detector_file(gibuu_output,
             timestamp = np.random.uniform(0, livetime)
 
             nu_in_trk = ROOT.Trk()
-            nu_in_trk.id = 0
+            nu_in_trk.id = mc_trk_id
+            mc_trk_id += 1
             nu_in_trk.mother_id = -1
             nu_in_trk.type = nu_type
             nu_in_trk.pos.set(*vtx_pos)
@@ -316,16 +364,20 @@ def write_detector_file(gibuu_output,
             nu_in_trk.E = event.lepIn_E
             nu_in_trk.t = timestamp
 
-            lep_out_trk = ROOT.Trk()
-            lep_out_trk.id = 1
-            lep_out_trk.mother_id = 0
-            lep_out_trk.type = sec_lep_type
-            lep_out_trk.pos.set(*vtx_pos)
-            mom = np.array([event.lepOut_Px, event.lepOut_Py, event.lepOut_Pz])
-            p_dir = R.apply(mom / np.linalg.norm(mom))
-            lep_out_trk.dir.set(*p_dir)
-            lep_out_trk.E = event.lepOut_E
-            lep_out_trk.t = timestamp
+            if not propagate_tau:
+                lep_out_trk = ROOT.Trk()
+                lep_out_trk.id = mc_trk_id
+                mc_trk_id += 1
+                lep_out_trk.mother_id = 0
+                lep_out_trk.type = sec_lep_type
+                lep_out_trk.pos.set(*vtx_pos)
+                mom = np.array(
+                    [event.lepOut_Px, event.lepOut_Py, event.lepOut_Pz])
+                p_dir = R.apply(mom / np.linalg.norm(mom))
+                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)
 
             # bjorken_y = 1.0 - float(event.lepOut_E / event.lepIn_E)
             nu_in_trk.setusr('bx', bjorkenx[i])
@@ -334,20 +386,13 @@ def write_detector_file(gibuu_output,
             nu_in_trk.setusr("cc", is_cc)
 
             aafile.evt.mc_trks.push_back(nu_in_trk)
-            aafile.evt.mc_trks.push_back(lep_out_trk)
 
-            for i in range(len(event.E)):
-                trk = ROOT.Trk()
-                trk.id = i + 2
-                mom = np.array([event.Px[i], event.Py[i], event.Pz[i]])
-                p_dir = R.apply(mom / np.linalg.norm(mom))
-                trk.pos.set(*vtx_pos)
-                trk.dir.set(*p_dir)
-                trk.mother_id = 0
-                trk.type = int(event.barcode[i])
-                trk.E = event.E[i]
-                trk.t = timestamp
-                aafile.evt.mc_trks.push_back(trk)
+            if tau_secondaries is not None:
+                event_tau_sec = tau_secondaries[i]
+                add_particles(event_tau_sec)
+
+            add_particles(event)
+
             aafile.write()
 
     del aafile

km3buu/propagation.py

+# Filename: propagation.py
+"""
+Propagation functions for km3buu
+
+"""
+
+__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 numpy as np
+import scipy.constants as const
+import proposal as pp
+from particle import Particle
+import awkward1 as ak
+from collections import defaultdict
+
+from .config import Config
+
+PROPOSAL_LEPTON_DEFINITIONS = {
+    11: pp.particle.EMinusDef,
+    -11: pp.particle.EPlusDef,
+    12: pp.particle.NuEDef,
+    -12: pp.particle.NuEBarDef,
+    13: pp.particle.MuMinusDef,
+    -13: pp.particle.MuPlusDef,
+    14: pp.particle.NuMuDef,
+    -14: pp.particle.NuMuBarDef,
+    15: pp.particle.TauMinusDef,
+    -15: pp.particle.TauPlusDef,
+    16: pp.particle.NuTauDef,
+    -16: pp.particle.NuTauBarDef
+}
+
+
+def _setup_propagator(max_distance, lepton_definition):
+    sector = pp.SectorDefinition()
+    sector.medium = pp.medium.AntaresWater(1.0)
+    geometry = pp.geometry.Sphere(pp.Vector3D(), max_distance, 0)
+    sector.geometry = geometry
+    sector.scattering_model = pp.scattering.ScatteringModel.Highland
+
+    cfg = Config()
+    interpolation = pp.InterpolationDef()
+    interpolation.path_to_tables = cfg.proposal_itp_tables
+    interpolation.path_to_tables_readonly = cfg.proposal_itp_tables
+
+    return pp.Propagator(sector_defs=[sector],
+                         particle_def=lepton_definition,
+                         detector=geometry,
+                         interpolation_def=interpolation)
+
+
+def propagate_lepton(lepout_data, pdgid):
+    """
+    Lepton propagation based on PROPOSAL
+
+    Parameters
+    ----------
+    lepout_data: awkward1.highlevel.Array
+        Lepton data in the GiBUU output shape containing the fields 
+        'lepOut_E, lepOut_Px, lepOut_Py, lepOut_Pz'
+    pdgid:
+        The pdgid of the propagated lepton
+
+    Returns
+    -------
+    awkward1.highlevel.Array (E, Px, Py, Pz, x, y, z)
+    """
+    lepton_info = Particle.from_pdgid(pdgid)
+    prop_range = const.c * lepton_info.lifetime * 1e11 * np.max(
+        lepout_data.lepOut_E) / lepton_info.mass  #[cm]
+
+    lepton_def = PROPOSAL_LEPTON_DEFINITIONS[pdgid]()
+    lepton = pp.particle.DynamicData(lepton_def.particle_type)
+    propagated_data = defaultdict(list)
+
+    propagator = _setup_propagator(10 * prop_range, lepton_def)
+
+    for entry in lepout_data:
+        lepton.energy = entry.lepOut_E * 1e3
+        lepton.direction = pp.Vector3D(entry.lepOut_Px, entry.lepOut_Py,
+                                       entry.lepOut_Pz)
+        lepton.direction.normalize()
+        secondaries = propagator.propagate(lepton, 5 * prop_range)
+
+        E = np.array(secondaries.energy) / 1e3
+        pdgid = [p.type for p in secondaries.particles]
+        Px = [p.direction.x * p.momentum for p in secondaries.particles]
+        Py = [p.direction.y * p.momentum for p in secondaries.particles]
+        Pz = [p.direction.z * p.momentum for p in secondaries.particles]
+        x = [p.position.x for p in secondaries.particles]
+        y = [p.position.y for p in secondaries.particles]
+        z = [p.position.z for p in secondaries.particles]
+
+        propagated_data["E"].append(E)
+        propagated_data["Px"].append(Px)
+        propagated_data["Py"].append(Py)
+        propagated_data["Pz"].append(Pz)
+        propagated_data["barcode"].append(pdgid)
+        propagated_data["x"].append(x)
+        propagated_data["y"].append(y)
+        propagated_data["z"].append(z)
+
+    return ak.Array(propagated_data)

km3buu/tests/test_propagation.py

+#!/usr/bin/env python
+# coding=utf-8
+# Filename: test_propagation.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 csv
+import unittest
+import numpy as np
+import uproot4
+from os.path import abspath, join, dirname
+from thepipe.logger import get_logger
+from km3net_testdata import data_path
+
+import proposal as pp
+
+from km3buu.output import GiBUUOutput
+from km3buu.propagation import propagate_lepton
+
+TESTDATA_DIR = data_path("gibuu")
+
+pp.RandomGenerator.get().set_seed(1234)
+
+
+class TestTauPropagation(unittest.TestCase):
+    def setUp(self):
+        log = get_logger("ctrl.py")
+        log.setLevel("INFO")
+        self.gibuu_output = GiBUUOutput(TESTDATA_DIR)
+        fname = join(TESTDATA_DIR, self.gibuu_output.root_pert_files[0])
+        fobj = uproot4.open(fname)
+        data = fobj["RootTuple"].arrays()
+        self.sec = propagate_lepton(data, 15)
+
+    def test_secondary_momenta(self):
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].E),
+                                             [0.5, 1.3, 0.3],
+                                             decimal=1)
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].Px),
+                                             [-467.4, 320.7, -245.5],
+                                             decimal=1)
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].Py),
+                                             [127.2, -822.4, 217.5],
+                                             decimal=1)
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].Pz),
+                                             [179., 967.1, -41.1],
+                                             decimal=1)
+
+    def test_secondary_types(self):
+        np.testing.assert_array_equal(np.array(self.sec[0].barcode),
+                                      [13, 16, -14])
+
+    def test_secondary_positions(self):
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].x), [0, 0],
+                                             decimal=1)
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].y), [0, 0],
+                                             decimal=1)
+        np.testing.assert_array_almost_equal(np.array(self.sec[0].z), [0, 0],
+                                             decimal=1)

requirements.txt

@@ -10,3 +10,4 @@ uproot4
 awkward1>=0.4.4
 pandas
 mendeleev
+proposal