Physics Kernels

Relativistic kinematics implemented as batched CUDA/C++ kernels. All functions are available through the pyc::physics C++ API (see pyc Interface) and as torch.ops.tpyc / torch.ops.cupyc Python operators.

The namespace is organised into two coordinate-system variants:

pyc::physics::cartesian — inputs as Cartesian \((p_x, p_y, p_z, E)\)
pyc::physics::polar — inputs as polar \((p_T, \eta, \phi, E)\)

Each variant exposes separate (per-column tensor arguments) and combined (single stacked Nx4 tensor) overloads.

The formulas below are taken directly from physics/physics.cxx (CPU) and physics/physics.cu (CUDA):

\[\begin{split}|\vec{p}|^2 &= p_x^2 + p_y^2 + p_z^2 \\ |\vec{p}| &= \sqrt{|\vec{p}|^2} \\ \beta^2 &= \frac{|\vec{p}|^2}{E^2}, \quad \beta = \frac{|\vec{p}|}{E} \\ m^2 &= E^2 - |\vec{p}|^2, \quad m = \mathrm{sign}(m^2)\sqrt{|m^2|} \\ m_T^2 &= E^2 - p_z^2, \quad m_T = \mathrm{sign}(m_T^2)\sqrt{|m_T^2|} \\ \theta &= \mathrm{atan2}\!\left(\sqrt{p_x^2+p_y^2},\, p_z\right) \\ \Delta R &= \sqrt{(\Delta\eta)^2 + (\Delta\phi)^2}\end{split}\]

Cartesian – Separate Inputs

namespace separate

Kinematics from individual Cartesian component tensors.

Functions

torch::Tensor P2(torch::Tensor px, torch::Tensor py, torch::Tensor pz): Compute |p|² = px² + py² + pz².

torch::Tensor P(torch::Tensor px, torch::Tensor py, torch::Tensor pz): Compute |p|.

torch::Tensor Beta2(torch::Tensor px, torch::Tensor py, torch::Tensor pz, torch::Tensor e): Compute β² = |p|²/E².

torch::Tensor Beta(torch::Tensor px, torch::Tensor py, torch::Tensor pz, torch::Tensor e): Compute β = |p|/E.

torch::Tensor M2(torch::Tensor px, torch::Tensor py, torch::Tensor pz, torch::Tensor e): Compute m² = E² - |p|².

torch::Tensor M(torch::Tensor px, torch::Tensor py, torch::Tensor pz, torch::Tensor e): Compute invariant mass m.

torch::Tensor Mt2(torch::Tensor pz, torch::Tensor e): Compute transverse mass squared.

torch::Tensor Mt(torch::Tensor pz, torch::Tensor e): Compute transverse mass.

torch::Tensor Theta(torch::Tensor px, torch::Tensor py, torch::Tensor pz): Compute polar angle θ.

torch::Tensor DeltaR(torch::Tensor px1, torch::Tensor px2, torch::Tensor py1, torch::Tensor py2, torch::Tensor pz1, torch::Tensor pz2): Compute ΔR between two particles.

Cartesian – Combined Input

namespace combined

Functions

torch::Tensor P2(torch::Tensor pmc)

torch::Tensor P(torch::Tensor pmc)

torch::Tensor Beta2(torch::Tensor pmc)

torch::Tensor Beta(torch::Tensor pmc)

torch::Tensor M2(torch::Tensor pmc)

torch::Tensor M(torch::Tensor pmc)

torch::Tensor Mt2(torch::Tensor pmc)

torch::Tensor Mt(torch::Tensor pmc)

torch::Tensor Theta(torch::Tensor pmc)

torch::Tensor DeltaR(torch::Tensor pmc1, torch::Tensor pmc2)

Polar – Separate Inputs

namespace separate

Functions

torch::Tensor P2(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi)

torch::Tensor P(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi)

torch::Tensor Beta2(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi, torch::Tensor e)

torch::Tensor Beta(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi, torch::Tensor e)

torch::Tensor M2(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi, torch::Tensor e)

torch::Tensor M(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi, torch::Tensor e)

torch::Tensor Mt2(torch::Tensor pt, torch::Tensor eta, torch::Tensor e)

torch::Tensor Mt(torch::Tensor pt, torch::Tensor eta, torch::Tensor e)

torch::Tensor Theta(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi)

torch::Tensor DeltaR(torch::Tensor eta1, torch::Tensor eta2, torch::Tensor phi1, torch::Tensor phi2)

Polar – Combined Input

namespace combined

Functions

torch::Tensor P2(torch::Tensor pmu)

torch::Tensor P(torch::Tensor pmu)

torch::Tensor Beta2(torch::Tensor pmu)

torch::Tensor Beta(torch::Tensor pmu)

torch::Tensor M2(torch::Tensor pmu)

torch::Tensor M(torch::Tensor pmu)

torch::Tensor Mt2(torch::Tensor pmu)

torch::Tensor Mt(torch::Tensor pmu)

torch::Tensor Theta(torch::Tensor pmu)

torch::Tensor DeltaR(torch::Tensor pmu1, torch::Tensor pmu2)

Internal Kernel Namespace

The physics_ internal namespace provides the raw implementations before wrapping in the public pyc::physics API.

namespace physics_

Internal namespace for batched relativistic kinematics.

Functions

torch::Tensor P2(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz)

Compute |p|² from individual components.

Parameters:

px – [N],
py – [N],
pz – [N].

Returns:

[N] |p|².

torch::Tensor P2(torch::Tensor *pmc)

Compute |p|² from stacked [N, 4] Cartesian momenta.

Parameters:: pmc – [N, 4].
Returns:: [N] |p|².

torch::Tensor P(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz): Compute |p| = √(px² + py² + pz²).

torch::Tensor P(torch::Tensor *pmc)

torch::Tensor Beta2(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz, torch::Tensor *e): Compute β² = |p|²/E².

torch::Tensor Beta2(torch::Tensor *pmc)

torch::Tensor Beta(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz, torch::Tensor *e): Compute β = |p|/E.

torch::Tensor Beta(torch::Tensor *pmc)

torch::Tensor M2(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz, torch::Tensor *e): Compute m² = E² - |p|².

torch::Tensor M2(torch::Tensor *pmc)

torch::Tensor M(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz, torch::Tensor *e): Compute invariant mass m = √|m²|.

torch::Tensor M(torch::Tensor *pmc)

torch::Tensor Mt2(torch::Tensor *pz, torch::Tensor *e)

Compute transverse mass squared.

Parameters:

pz – [N] z-momentum.
e – [N] energy.

Returns:

[N] mT².

torch::Tensor Mt2(torch::Tensor *pmc)

torch::Tensor Mt(torch::Tensor *pz, torch::Tensor *e): Compute transverse mass.

torch::Tensor Mt(torch::Tensor *pmc)

torch::Tensor Theta(torch::Tensor *pmc): Compute polar angle θ from stacked [N, 4] Cartesian momenta.

torch::Tensor Theta(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz): Compute polar angle θ from individual components.

torch::Tensor DeltaR(torch::Tensor *pmu1, torch::Tensor *pmu2)

Compute ΔR from two stacked [N, 4] polar four-momenta.

Parameters:

pmu1 – [N, 4] (pT, η, φ, E).
pmu2 – [N, 4].

Returns:

[N] ΔR = √(Δη² + Δφ²).

torch::Tensor DeltaR(torch::Tensor *eta1, torch::Tensor *eta2, torch::Tensor *phi1, torch::Tensor *phi2): Compute ΔR from individual η and φ tensors.