diff --git a/km3buu/physics.py b/km3buu/physics.py index 52d931f8787989d246834deb2e8ea628755acafa..fd3faead58b5d08ad37ce1396639d31be9b6a4f3 100644 --- a/km3buu/physics.py +++ b/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) diff --git a/km3buu/tests/test_physics.py b/km3buu/tests/test_physics.py index dc3dc7ef97ac3185b7ba325375cefb3419ed86a8..256fb2b724e14fd6a6eb44565fb900b9ba34a997 100644 --- a/km3buu/tests/test_physics.py +++ b/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,