diff --git a/km3buu/physics.py b/km3buu/physics.py index fd3faead58b5d08ad37ce1396639d31be9b6a4f3..c943f2c5cd4eb42a59a5a6f01d3d5628667f8d66 100644 --- a/km3buu/physics.py +++ b/km3buu/physics.py @@ -107,51 +107,89 @@ HE_PARAMS = { } -@np.vectorize def _get_particle_rest_mass(pdgid): - return Particle.from_pdgid(pdgid).mass / 1e3 + @np.vectorize + def vfunc(x): + mass = Particle.from_pdgid(x).mass + if mass is None: + return 0 + return mass / 1e3 + + pdgids, invmap = np.unique(ak.to_numpy(pdgid), return_inverse=True) + masses = vfunc(pdgids) + return masses[invmap] -def visible_energy_fraction(pdgid, energy): +def get_kinetic_energy(energy, pdgid): """ - Returns the visible energy fraction in the one particle approximation (OPA) + Returns the kinetic energy + + Parameters + ---------- + energy: float [GeV] + Total energy of the given particle + pdgid: int + PDGID of the given particle + """ + mass = np.array(_get_particle_rest_mass(pdgid)) + return np.sqrt(ak.to_numpy(energy)**2 - mass**2) + + +def visible_energy(energy, pdgid): + """ + Returns the visible energy in the one particle approximation (OPA) how it is used in JSirene (i.e. JPythia.hh) Parameters ---------- + energy: float [GeV] + Total energy of the given particle pdgid: int PDGID of the given particle + """ + return get_kinetic_energy(energy, pdgid) * visible_energy_fraction( + energy, pdgid) + + +def visible_energy_fraction(energy, pdgid): + """ + Returns the visible energy fraction in the one particle approximation (OPA) + how it is used in JSirene (i.e. JPythia.hh) + + Parameters + ---------- energy: float [GeV] Total energy of the given particle + pdgid: int + PDGID 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) + ekin = get_kinetic_energy(energy[mask], pdgid[mask]) + retval[mask] = high_energy_weight(ekin) 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 + ekin = np.sqrt(ak.to_numpy(energy)[mask]**2 - mass**2) + retval[mask] = muon_range_seawater(ekin, mass) / 4.7 return retval @np.vectorize -def km3_opa_fraction(pdgid, energy): +def km3_opa_fraction(energy, pdgid): """ 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 + pdgid: int + PDGID of the given particle """ # Cover trivial cases, i.e. 'non-visible' particles and ultra-high energy @@ -268,6 +306,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): """ diff --git a/km3buu/tests/test_physics.py b/km3buu/tests/test_physics.py index 256fb2b724e14fd6a6eb44565fb900b9ba34a997..29efb84ccf8800d1806eb17f54d26ab04f4bf7af 100644 --- a/km3buu/tests/test_physics.py +++ b/km3buu/tests/test_physics.py @@ -49,7 +49,7 @@ class TestVisEnergyParticle(unittest.TestCase): def test_particles(self): for i, pdgid in enumerate(self.particles): - val = km3_opa_fraction(pdgid, self.ref_values[0, :]) + val = km3_opa_fraction(self.ref_values[0, :], pdgid) assert np.allclose(self.ref_values[i + 1, :], val, rtol=0.05,