Coordinate Transforms

Polar \(\leftrightarrow\) Cartesian coordinate system transformations. All functions operate on torch::Tensor objects and are accessible at the Python level as torch.ops.tpyc.* / torch.ops.cupyc.*.

The namespace is split into:

pyc::transform::separate — accepts one column tensor per coordinate
pyc::transform::combined — accepts a single stacked Nx4 tensor

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

\[ \begin{align}\begin{aligned}p_x = p_T \cos\phi, \quad p_y = p_T \sin\phi, \quad p_z = p_T \sinh\eta\\p_T = \sqrt{p_x^2 + p_y^2}, \quad \phi = \mathrm{atan2}(p_y, p_x), \quad \eta = \mathrm{asinh}\!\left(\frac{p_z}{p_T}\right)\end{aligned}\end{align} \]

Separate Column Inputs

namespace separate

Overloads operating on individual component tensors.

Functions

torch::Tensor Pt(torch::Tensor px, torch::Tensor py): Compute pT = √(px² + py²).

torch::Tensor Eta(torch::Tensor px, torch::Tensor py, torch::Tensor pz): Compute pseudorapidity.

torch::Tensor Phi(torch::Tensor px, torch::Tensor py): Compute azimuthal angle φ = atan2(py, px).

torch::Tensor PtEtaPhi(torch::Tensor px, torch::Tensor py, torch::Tensor pz): Compute (pT, η, φ) in one call.

torch::Tensor PtEtaPhiE(torch::Tensor px, torch::Tensor py, torch::Tensor pz, torch::Tensor e): Compute (pT, η, φ, E) in one call.

torch::Tensor Px(torch::Tensor pt, torch::Tensor phi): Compute px = pT cos φ.

torch::Tensor Py(torch::Tensor pt, torch::Tensor phi): Compute py = pT sin φ.

torch::Tensor Pz(torch::Tensor pt, torch::Tensor eta): Compute pz = pT sinh η.

torch::Tensor PxPyPz(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi): Compute (px, py, pz) in one call.

torch::Tensor PxPyPzE(torch::Tensor pt, torch::Tensor eta, torch::Tensor phi, torch::Tensor e): Compute (px, py, pz, E) in one call.

Combined Stacked Input

namespace combined

Overloads operating on stacked [N, 4] or [N, 3] tensors.

Functions

torch::Tensor Pt(torch::Tensor pmc)

torch::Tensor Eta(torch::Tensor pmc)

torch::Tensor Phi(torch::Tensor pmc)

torch::Tensor PtEtaPhi(torch::Tensor pmc)

torch::Tensor PtEtaPhiE(torch::Tensor pmc)

torch::Tensor Px(torch::Tensor pmu)

torch::Tensor Py(torch::Tensor pmu)

torch::Tensor Pz(torch::Tensor pmu)

torch::Tensor PxPyPz(torch::Tensor pmu)

torch::Tensor PxPyPzE(torch::Tensor pmu)

Internal Kernel Namespace

The transform_ internal namespace provides the raw implementations before wrapping in the public pyc::transform API.

namespace transform_

Internal namespace for batched coordinate transforms.

Functions

torch::Tensor Px(torch::Tensor *pt, torch::Tensor *phi)

Compute px = pT cos φ.

Parameters:

pt – [N] transverse momentum.
phi – [N] azimuthal angle.

Returns:

[N].

torch::Tensor Py(torch::Tensor *pt, torch::Tensor *phi)

Compute py = pT sin φ.

Returns:: [N].

torch::Tensor Pz(torch::Tensor *pt, torch::Tensor *eta)

Compute pz = pT sinh η.

Parameters:

pt – [N].
eta – [N] pseudorapidity.

Returns:

[N].

torch::Tensor PxPyPz(torch::Tensor *pt, torch::Tensor *eta, torch::Tensor *phi): Compute (px, py, pz) and stack into [N, 3].

torch::Tensor PxPyPzE(torch::Tensor *pt, torch::Tensor *eta, torch::Tensor *phi, torch::Tensor *energy): Compute (px, py, pz, E) and stack into [N, 4].

torch::Tensor PxPyPz(torch::Tensor *pmu)

Overload: extract (px, py, pz) from stacked [N, 4] polar tensor.

Parameters:: pmu – [N, 4] (pT, η, φ, E).
Returns:: [N, 3] (px, py, pz).

torch::Tensor PxPyPzE(torch::Tensor *pmu)

Overload: compute full Cartesian four-momentum from [N, 4] polar.

Returns:: [N, 4] (px, py, pz, E).

torch::Tensor Pt(torch::Tensor *px, torch::Tensor *py)

Compute pT = √(px² + py²).

Returns:: [N].

torch::Tensor PtEta(torch::Tensor *pt, torch::Tensor *pz): Compute (pT, η) and stack into [N, 2].

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

Compute φ = atan2(py, px) from stacked [N, 4] Cartesian tensor.

Returns:: [N].

torch::Tensor Phi(torch::Tensor *px, torch::Tensor *py): Compute φ from individual px and py tensors.

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

Compute pseudorapidity from [N, 4] Cartesian tensor.

Returns:: [N].

torch::Tensor Eta(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz): Compute pseudorapidity from individual components.

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

Compute (pT, η, φ) from [N, 4] Cartesian tensor.

Returns:: [N, 3].

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

Compute (pT, η, φ, E) from [N, 4] Cartesian tensor.

Returns:: [N, 4].

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

Compute (pT, η, φ) from individual components.

Returns:: [N, 3].

torch::Tensor PtEtaPhiE(torch::Tensor *px, torch::Tensor *py, torch::Tensor *pz, torch::Tensor *e)

Compute (pT, η, φ, E) from individual components.

Returns:: [N, 4].