1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73
|
# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/vector for details.
from __future__ import annotations
import vector
from vector import (
MomentumNumpy2D,
MomentumNumpy3D,
MomentumNumpy4D,
MomentumObject2D,
MomentumObject3D,
MomentumObject4D,
VectorObject4D,
)
def test_handler_of():
object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0)
object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0)
protocol = vector._methods._handler_of(object_a, object_b)
assert protocol == object_a
def test_momentum_coordinate_transforms():
numpy_vec = vector.array(
{
"px": [1.0, 2.0, 3.0],
"py": [-1.0, 2.0, 3.0],
},
)
object_vec = MomentumObject2D(px=0.0, py=0.0)
for t1 in "pxpy", "ptphi":
for t2 in "pz", "eta", "theta":
for t3 in "mass", "energy":
transformed_object = getattr(object_vec, "to_" + t1)()
assert isinstance(transformed_object, MomentumObject2D)
assert hasattr(transformed_object, t1[:2])
assert hasattr(transformed_object, t1[2:])
transformed_object = getattr(object_vec, "to_" + t1 + t2)()
assert isinstance(transformed_object, MomentumObject3D)
assert hasattr(transformed_object, t1[:2])
assert hasattr(transformed_object, t1[2:])
assert hasattr(transformed_object, t2)
transformed_object = getattr(object_vec, "to_" + t1 + t2 + t3)()
assert isinstance(transformed_object, MomentumObject4D)
assert hasattr(transformed_object, t1[:2])
assert hasattr(transformed_object, t1[2:])
assert hasattr(transformed_object, t2)
assert hasattr(transformed_object, t3)
transformed_numpy = getattr(numpy_vec, "to_" + t1)()
assert isinstance(transformed_numpy, MomentumNumpy2D)
assert hasattr(transformed_numpy, t1[:2])
assert hasattr(transformed_numpy, t1[2:])
transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2)()
assert isinstance(transformed_numpy, MomentumNumpy3D)
assert hasattr(transformed_numpy, t1[:2])
assert hasattr(transformed_numpy, t1[2:])
assert hasattr(transformed_numpy, t2)
transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2 + t3)()
assert isinstance(transformed_numpy, MomentumNumpy4D)
assert hasattr(transformed_numpy, t1[:2])
assert hasattr(transformed_numpy, t1[2:])
assert hasattr(transformed_numpy, t2)
assert hasattr(transformed_numpy, t3)
|