Update visible energy definition

km3buu/physics.py

@@ -16,6 +16,7 @@ __email__ = "jschumann@km3net.de"
 __status__ = "Development"
 
 import numpy as np
+import awkward as ak
 from particle import Particle
 
 from .config import read_default_media_compositions
@@ -107,16 +108,50 @@ HE_PARAMS = {
 
 
 @np.vectorize
+def _get_particle_rest_mass(pdgid):
+    return Particle.from_pdgid(pdgid).mass / 1e3
+
+
 def visible_energy_fraction(pdgid, energy):
     """
     Returns the visible energy fraction in the one particle approximation (OPA)
+    how it is used in JSirene (i.e. JPythia.hh)
 
     Parameters
     ----------
     pdgid: int
         PDGID of the given particle
-    energy: float
-        Kinetic energy of the given particle
+    energy: float [GeV]
+        Total energy of the given particle 
+    """
+    pdgid = ak.to_numpy(pdgid)
+    retval = np.zeros_like(pdgid, dtype=np.float64)
+    mask = np.isin(pdgid, [11, -11, 22, 111, 221])
+    retval[mask] = 1.0
+    mask = np.isin(pdgid, [-211, 211])
+    mass = np.array(_get_particle_rest_mass(pdgid[mask]))
+    tmp = np.sqrt(ak.to_numpy(energy)[mask]**2 - mass**2)
+    retval[mask] = high_energy_weight(tmp)
+    mask = np.isin(pdgid, [13])
+    if np.any(mask):
+        mass = Particle.from_pdgid(13).mass / 1e3
+        tmp = np.sqrt(ak.to_numpy(energy)[mask]**2 - mass**2)
+        retval[mask] = muon_range_seawater(tmp, mass) / 4.7
+    return retval
+
+
+@np.vectorize
+def km3_opa_fraction(pdgid, energy):
+    """
+    Returns the visible energy fraction in the one particle approximation (OPA)
+    how it is used in KM3
+
+    Parameters
+    ----------
+    pdgid: int
+        PDGID of the given particle
+    energy: float [GeV]
+        Kinetic energy of the given particle 
     """
 
     # Cover trivial cases, i.e. 'non-visible' particles and ultra-high energy
@@ -233,7 +268,7 @@ def neutron_weight(energy):
                 energy + 1e4 * NEUTRON_PARAMS["Ni"] * energy**2 +
                 1e4 * NEUTRON_PARAMS["Nj"] * energy**3) / norm
 
-
+@np.vectorize
 def high_energy_weight(energy):
     """
     High energy weight (valid above 40 GeV)

km3buu/tests/test_physics.py

@@ -49,8 +49,7 @@ class TestVisEnergyParticle(unittest.TestCase):
 
     def test_particles(self):
         for i, pdgid in enumerate(self.particles):
-            vfunc = np.vectorize(visible_energy_fraction)
-            val = vfunc(pdgid, self.ref_values[0, :])
+            val = km3_opa_fraction(pdgid, self.ref_values[0, :])
             assert np.allclose(self.ref_values[i + 1, :],
                                val,
                                rtol=0.05,