Skip to content

Commit

Permalink
feat: allow momentum coords in to_Vector*D methods + cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@Saransh-cpp
Saransh-cpp committed Feb 19, 2024
1 parent ea3c018 commit e16643d
Show file tree
Hide file tree
Showing 7 changed files with 367 additions and 264 deletions.
3 changes: 3 additions & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,3 +142,6 @@ dmypy.json
/node_modules
/package.json
/yarn.lock

# MacOS
.DS_Store
72 changes: 44 additions & 28 deletions src/vector/_methods.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,9 @@ def to_Vector3D(self) -> VectorProtocolSpatial:
Projects this vector/these vectors onto azimuthal and longitudinal
coordinates only.
If 2D, a $z$ component of $0$ is imputed.
If 2D, a default $z$ component of $0$ is imputed.
The longitudinal coordinate can be passed as a named argument.
"""
raise AssertionError

Expand All @@ -193,9 +195,12 @@ def to_Vector4D(self) -> VectorProtocolLorentz:
Projects this vector/these vectors onto azimuthal, longitudinal,
and temporal coordinates.
If 2D or 3D, a $t$ component of $0$ is imputed.
If 3D, a default $t$ component of $0$ is imputed.
If 2D, a $z$ component of $0$ is imputed along with a default
$t$ component of $0$.
If 2D, a $z$ component of $0$ is imputed.
The longitudinal and temporal coordinates can be passed as named arguments.
"""
raise AssertionError

Expand Down Expand Up @@ -2963,34 +2968,34 @@ def to_Vector3D(
self,
*,
z: float | None = None,
pz: float | None = None,
theta: float | None = None,
eta: float | None = None,
) -> VectorProtocolSpatial:
"""
Converts a 2D vector to 3D vector.
The scalar longitudinal coordinate is broadcasted for NumPy and Awkward
vectors. Only a single longitudinal coordinate should be provided. Generic
coordinate counterparts should be provided for the momentum coordinates.
vectors. Only a single longitudinal coordinate should be provided.
Examples:
>>> import vector
>>> vec = vector.VectorObject2D(x=1, y=2)
>>> vec.to_Vector3D(z=1)
VectorObject3D(x=1, y=2, z=1)
>>> vec = vector.MomentumObject2D(px=1, py=2)
>>> vec.to_Vector3D(z=4)
>>> vec.to_Vector3D(pz=4)
MomentumObject3D(px=1, py=2, pz=4)
"""
if sum(x is not None for x in (z, theta, eta)) > 1:
raise TypeError(
"At most one longitudinal coordinate (`z`, `theta`, or `eta`) may be assigned (non-None)"
"At most one longitudinal coordinate (`z`/`pz`, `theta`, or `eta`) may be assigned (non-None)"
)

l_value = 0.0
l_type: type[Longitudinal] = LongitudinalZ
if z is not None:
l_value = z
if any(coord is not None for coord in [z, pz]):
l_value = next(coord for coord in [z, pz] if coord is not None)
elif eta is not None:
l_value = eta
l_type = LongitudinalEta
Expand All @@ -3009,49 +3014,55 @@ def to_Vector4D(
self,
*,
z: float | None = None,
pz: float | None = None,
theta: float | None = None,
eta: float | None = None,
t: float | None = None,
e: float | None = None,
E: float | None = None,
energy: float | None = None,
tau: float | None = None,
m: float | None = None,
M: float | None = None,
mass: float | None = None,
) -> VectorProtocolLorentz:
"""
Converts a 2D vector to 4D vector.
The scalar longitudinal and temporal coordinates are broadcasted for NumPy and
Awkward vectors. Only a single longitudinal and temporal coordinate should be
provided. Generic coordinate counterparts should be provided for the momentum
coordinates.
provided.
Examples:
>>> import vector
>>> vec = vector.VectorObject2D(x=1, y=2)
>>> vec.to_Vector4D(z=3, t=4)
VectorObject4D(x=1, y=2, z=3, t=4)
>>> vec = vector.MomentumObject2D(px=1, py=2)
>>> vec.to_Vector4D(z=4, t=4)
>>> vec.to_Vector4D(pz=4, energy=4)
MomentumObject4D(px=1, py=2, pz=4, E=4)
"""
if sum(x is not None for x in (z, theta, eta)) > 1:
raise TypeError(
"At most one longitudinal coordinate (`z`, `theta`, or `eta`) may be assigned (non-None)"
"At most one longitudinal coordinate (`z`/`pz`, `theta`, or `eta`) may be assigned (non-None)"
)
elif sum(x is not None for x in (t, tau)) > 1:
raise TypeError(
"At most one longitudinal coordinate (`t`, `tau`) may be assigned (non-None)"
"At most one longitudinal coordinate (`t`/`e`/`E`/`energy`, `tau`/`m`/`M`/`mass`) may be assigned (non-None)"
)

t_value = 0.0
t_type: type[Temporal] = TemporalT
if t is not None:
t_value = t
elif tau is not None:
t_value = tau
if any(coord is not None for coord in [tau, m, M, mass]):
t_type = TemporalTau
t_value = next(coord for coord in [tau, m, M, mass] if coord is not None)
elif any(coord is not None for coord in [t, e, E, energy]):
t_value = next(coord for coord in [t, e, E, energy] if coord is not None)

l_value = 0.0
l_type: type[Longitudinal] = LongitudinalZ
if z is not None:
l_value = z
if any(coord is not None for coord in [z, pz]):
l_value = next(coord for coord in [z, pz] if coord is not None)
elif eta is not None:
l_value = eta
l_type = LongitudinalEta
Expand Down Expand Up @@ -3083,36 +3094,41 @@ def to_Vector4D(
self,
*,
t: float | None = None,
e: float | None = None,
E: float | None = None,
energy: float | None = None,
tau: float | None = None,
m: float | None = None,
M: float | None = None,
mass: float | None = None,
) -> VectorProtocolLorentz:
"""
Converts a 3D vector to 4D vector.
The scalar temporal coordinate are broadcasted for NumPy and Awkward vectors.
Only a single temporal coordinate should be provided. Generic coordinate
counterparts should be provided for the momentum coordinates.
Only a single temporal coordinate should be provided.
Examples:
>>> import vector
>>> vec = vector.VectorObject3D(x=1, y=2, z=3)
>>> vec.to_Vector4D(t=4)
VectorObject4D(x=1, y=2, z=3, t=4)
>>> vec = vector.MomentumObject3D(px=1, py=2, pz=3)
>>> vec.to_Vector4D(tau=4)
>>> vec.to_Vector4D(M=4)
MomentumObject4D(px=1, py=2, pz=3, mass=4)
"""
if sum(x is not None for x in (t, tau)) > 1:
raise TypeError(
"At most one longitudinal coordinate (`t`, `tau`) may be assigned (non-None)"
"At most one longitudinal coordinate (`t`/`e`/`E`/`energy`, `tau`/`m`/`M`/`mass`) may be assigned (non-None)"
)

t_value = 0.0
t_type: type[Temporal] = TemporalT
if t is not None:
t_value = t
elif tau is not None:
t_value = tau
if any(coord is not None for coord in [tau, m, M, mass]):
t_type = TemporalTau
t_value = next(coord for coord in [tau, m, M, mass] if coord is not None)
elif any(coord is not None for coord in [t, e, E, energy]):
t_value = next(coord for coord in [t, e, E, energy] if coord is not None)

return self._wrap_result(
type(self),
Expand Down
30 changes: 30 additions & 0 deletions tests/backends/test_awkward.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,36 @@ def test_dimension_conversion():
assert ak.all(vec.to_Vector4D(t=1).y == vec.y)
assert ak.all(vec.to_Vector4D(t=1).z == vec.z)

# check if momentum coords work
vec = vector.Array(
[
[{"px": 1, "py": 1.1}, {"px": 2, "py": 2.1}],
[],
]
)
assert ak.all(vec.to_Vector3D(pz=1).pz == 1)

assert ak.all(vec.to_Vector4D(pz=1, m=1).pz == 1)
assert ak.all(vec.to_Vector4D(pz=1, m=1).m == 1)
assert ak.all(vec.to_Vector4D(pz=1, mass=1).mass == 1)
assert ak.all(vec.to_Vector4D(pz=1, M=1).M == 1)
assert ak.all(vec.to_Vector4D(pz=1, e=1).e == 1)
assert ak.all(vec.to_Vector4D(pz=1, energy=1).energy == 1)
assert ak.all(vec.to_Vector4D(pz=1, E=1).E == 1)

vec = vector.Array(
[
[{"px": 1, "py": 1.1, "pz": 1.2}, {"px": 2, "py": 2.1, "pz": 2.2}],
[],
]
)
assert ak.all(vec.to_Vector4D(m=1).m == 1)
assert ak.all(vec.to_Vector4D(mass=1).mass == 1)
assert ak.all(vec.to_Vector4D(M=1).M == 1)
assert ak.all(vec.to_Vector4D(e=1).e == 1)
assert ak.all(vec.to_Vector4D(energy=1).energy == 1)
assert ak.all(vec.to_Vector4D(E=1).E == 1)


def test_type_checks():
with pytest.raises(TypeError):
Expand Down
161 changes: 0 additions & 161 deletions tests/backends/test_numpy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,49 +14,6 @@
import vector.backends.numpy


def test_dimension_conversion():
# 2D -> 3D
vec = vector.VectorNumpy2D(
[(1.0, 1.0), (2.0, 2.0)],
dtype=[("x", float), ("y", float)],
)
assert all(vec.to_Vector3D(z=1).z == 1)
assert all(vec.to_Vector3D(eta=1).eta == 1)
assert all(vec.to_Vector3D(theta=1).theta == 1)

assert all(vec.to_Vector3D(z=1).x == vec.x)
assert all(vec.to_Vector3D(z=1).y == vec.y)

# 2D -> 4D
assert all(vec.to_Vector4D(z=1, t=1).t == 1)
assert all(vec.to_Vector4D(z=1, t=1).z == 1)
assert all(vec.to_Vector4D(eta=1, t=1).eta == 1)
assert all(vec.to_Vector4D(eta=1, t=1).t == 1)
assert all(vec.to_Vector4D(theta=1, t=1).theta == 1)
assert all(vec.to_Vector4D(theta=1, t=1).t == 1)
assert all(vec.to_Vector4D(z=1, tau=1).z == 1)
assert all(vec.to_Vector4D(z=1, tau=1).tau == 1)
assert all(vec.to_Vector4D(eta=1, tau=1).eta == 1)
assert all(vec.to_Vector4D(eta=1, tau=1).tau == 1)
assert all(vec.to_Vector4D(theta=1, tau=1).theta == 1)
assert all(vec.to_Vector4D(theta=1, tau=1).tau == 1)

assert all(vec.to_Vector4D(z=1, t=1).x == vec.x)
assert all(vec.to_Vector4D(z=1, t=1).y == vec.y)

# 3D -> 4D
vec = vector.VectorNumpy3D(
[(1.0, 1.0, 1.0), (2.0, 2.0, 2.0)],
dtype=[("x", float), ("y", float), ("z", float)],
)
assert all(vec.to_Vector4D(t=1).t == 1)
assert all(vec.to_Vector4D(tau=1).tau == 1)

assert all(vec.to_Vector4D(t=1).x == vec.x)
assert all(vec.to_Vector4D(t=1).y == vec.y)
assert all(vec.to_Vector4D(t=1).z == vec.z)


def test_type_checks():
with pytest.raises(TypeError):
vector.backends.numpy.VectorNumpy2D(
Expand Down Expand Up @@ -486,121 +443,3 @@ def test_count_nonzero_4d():
assert numpy.count_nonzero(v2, axis=1, keepdims=True).tolist() == [[3], [2]]
assert numpy.count_nonzero(v2, axis=0).tolist() == [2, 2, 1]
assert numpy.count_nonzero(v2, axis=0, keepdims=True).tolist() == [[2, 2, 1]]


def test_demotion():
v1 = vector.array(
{
"x": [10.0, 20.0, 30.0],
"y": [-10.0, 20.0, 30.0],
},
)
v2 = vector.array(
{
"x": [10.0, 20.0, 30.0],
"y": [-10.0, 20.0, 30.0],
"z": [5.0, 1.0, 1.0],
},
)
v3 = vector.array(
{
"x": [10.0, 20.0, 30.0],
"y": [-10.0, 20.0, 30.0],
"z": [5.0, 1.0, 1.0],
"t": [16.0, 31.0, 46.0],
},
)

v1_v2 = vector.array(
{
"x": [20.0, 40.0, 60.0],
"y": [-20.0, 40.0, 60.0],
},
)
v2_v3 = vector.array(
{
"x": [20.0, 40.0, 60.0],
"y": [-20.0, 40.0, 60.0],
"z": [10.0, 2.0, 2.0],
},
)
v1_v3 = vector.array(
{
"x": [20.0, 40.0, 60.0],
"y": [-20.0, 40.0, 60.0],
},
)

# order should not matter
assert all(v1 + v2 == v1_v2)
assert all(v2 + v1 == v1_v2)
assert all(v1 + v3 == v1_v3)
assert all(v3 + v1 == v1_v3)
assert all(v2 + v3 == v2_v3)
assert all(v3 + v2 == v2_v3)

v1 = vector.array(
{
"px": [10.0, 20.0, 30.0],
"py": [-10.0, 20.0, 30.0],
},
)
v2 = vector.array(
{
"px": [10.0, 20.0, 30.0],
"py": [-10.0, 20.0, 30.0],
"pz": [5.0, 1.0, 1.0],
},
)
v3 = vector.array(
{
"px": [10.0, 20.0, 30.0],
"py": [-10.0, 20.0, 30.0],
"pz": [5.0, 1.0, 1.0],
"t": [16.0, 31.0, 46.0],
},
)

p_v1_v2 = vector.array(
{
"px": [20.0, 40.0, 60.0],
"py": [-20.0, 40.0, 60.0],
},
)
p_v2_v3 = vector.array(
{
"px": [20.0, 40.0, 60.0],
"py": [-20.0, 40.0, 60.0],
"pz": [10.0, 2.0, 2.0],
},
)
p_v1_v3 = vector.array(
{
"px": [20.0, 40.0, 60.0],
"py": [-20.0, 40.0, 60.0],
},
)

# order should not matter
assert all(v1 + v2 == p_v1_v2)
assert all(v2 + v1 == p_v1_v2)
assert all(v1 + v3 == p_v1_v3)
assert all(v3 + v1 == p_v1_v3)
assert all(v2 + v3 == p_v2_v3)
assert all(v3 + v2 == p_v2_v3)

v2 = vector.array(
{
"x": [10.0, 20.0, 30.0],
"y": [-10.0, 20.0, 30.0],
"z": [5.0, 1.0, 1.0],
},
)

# momentum + generic = generic
assert all(v1 + v2 == v1_v2)
assert all(v2 + v1 == v1_v2)
assert all(v1 + v3 == v1_v3)
assert all(v3 + v1 == v1_v3)
assert all(v2 + v3 == v2_v3)
assert all(v3 + v2 == v2_v3)
Loading

0 comments on commit e16643d

Please sign in to comment.