Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow momentum coords in to_Vector*D methods + cleanup #423

Merged
merged 4 commits into from
Feb 29, 2024
Merged

feat: allow momentum coords in to_Vector*D methods + cleanup #423

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7092950
feat: allow momentum coords in to_Vector*D methods + cleanup
Saransh-cpp Feb 19, 2024
b12df12
Update tests/compute/test_demotion.py
Saransh-cpp Feb 23, 2024
9809cce
Check for multiple momentum + generic coords
Saransh-cpp Feb 29, 2024
1bca97e
Uniformity
Saransh-cpp Feb 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
View file
Open in desktop
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
80 changes: 48 additions & 32 deletions src/vector/_methods.py
View file
Open in desktop
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 @@ -3163,34 +3168,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:
if sum(x is not None for x in (z, pz, 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 @@ -3209,49 +3214,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,
Comment on lines 3216 to +3227
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If somebody passes both z and pz, or both z and theta for that matter, does the function complain? It shouldn't just pick one and use it—providing (possibly) conflicting arguments should be an error, to prevent silent mistakes.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review! Yes, the functions was checking for multiple generic coordinates, and I've added momentum coordinated to that check too.

) -> 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:
if sum(x is not None for x in (z, pz, 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:
elif sum(x is not None for x in (t, tau, m, M, mass, e, E, energy)) > 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 @@ -3283,36 +3294,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:
if sum(x is not None for x in (t, tau, m, M, mass, e, E, energy)) > 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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,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
Loading