File: test_methods.py

package info (click to toggle)
python-vector 1.6.3-1
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 2,984 kB
  • sloc: python: 40,200; makefile: 13
file content (73 lines) | stat: -rw-r--r-- 2,952 bytes parent folder | download
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)