From 5d5467063e204b3f106ebab18e5c72466e61479d Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 23 Feb 2023 23:01:02 +0530 Subject: [PATCH] feat: allow passing coordinates to to_Vector*D (#319) * feat: allow passing coordinates to to_Vector*D * Add tests * Add docstrings * import or skip awkward for light tests * Update src/vector/_methods.py * Add a note for momentum type coords * Fix repr of MomentumObject3D * Update src/vector/_methods.py Co-authored-by: Jim Pivarski * style: pre-commit fixes * Better error message --------- Co-authored-by: Jim Pivarski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- src/vector/_methods.py | 129 ++++++++++++++++++++++++++++++--- src/vector/backends/object.py | 2 +- tests/backends/test_awkward.py | 49 ++++++++++++- tests/backends/test_numpy.py | 43 +++++++++++ tests/backends/test_object.py | 37 ++++++++++ 5 files changed, 246 insertions(+), 14 deletions(-) diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 85befaee..c95d5e98 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -2966,32 +2966,103 @@ def to_Vector3D( 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. + + 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) + MomentumObject3D(px=1, py=2, pz=4) + """ if sum(x is not None for x in (z, theta, eta)) > 1: - raise TypeError("Only one non-None parameter allowed") + raise TypeError( + "At most one longitudinal coordinate (`z`, `theta`, or `eta`) may be assigned (non-None)" + ) - coord_value = 0.0 + l_value = 0.0 l_type: type[Longitudinal] = LongitudinalZ if z is not None: - coord_value = z + l_value = z elif eta is not None: - coord_value = eta + l_value = eta l_type = LongitudinalEta elif theta is not None: - coord_value = theta + l_value = theta l_type = LongitudinalTheta return self._wrap_result( type(self), - (*self.azimuthal.elements, coord_value), + (*self.azimuthal.elements, l_value), [_aztype(self), l_type, None], 1, ) - def to_Vector4D(self) -> VectorProtocolLorentz: + def to_Vector4D( + self, + *, + z: float | None = None, + theta: float | None = None, + eta: float | None = None, + t: float | None = None, + tau: 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. + + 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) + 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)" + ) + 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)" + ) + + 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 + t_type = TemporalTau + + l_value = 0.0 + l_type: type[Longitudinal] = LongitudinalZ + if z is not None: + l_value = z + elif eta is not None: + l_value = eta + l_type = LongitudinalEta + elif theta is not None: + l_value = theta + l_type = LongitudinalTheta + return self._wrap_result( type(self), - (*self.azimuthal.elements, 0, 0), - [_aztype(self), LongitudinalZ, TemporalT], + (*self.azimuthal.elements, l_value, t_value), + [_aztype(self), l_type, t_type], 1, ) @@ -3008,11 +3079,45 @@ def to_Vector2D(self) -> VectorProtocolPlanar: def to_Vector3D(self) -> VectorProtocolSpatial: return self - def to_Vector4D(self) -> VectorProtocolLorentz: + def to_Vector4D( + self, + *, + t: float | None = None, + tau: 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. + + 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) + 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)" + ) + + 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 + t_type = TemporalTau + return self._wrap_result( type(self), - self.azimuthal.elements + self.longitudinal.elements + (0,), - [_aztype(self), _ltype(self), TemporalT], + (*self.azimuthal.elements, *self.longitudinal.elements, t_value), + [_aztype(self), _ltype(self), t_type], 1, ) diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py index 6891a07f..0f44ce5d 100644 --- a/src/vector/backends/object.py +++ b/src/vector/backends/object.py @@ -1259,7 +1259,7 @@ def __repr__(self) -> str: for x in lnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.longitudinal, x)}") - return "vector.MomentumObject3D(" + ", ".join(out) + ")" + return "MomentumObject3D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import MomentumNumpy3D diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py index 6ad87664..c8e043cd 100644 --- a/tests/backends/test_awkward.py +++ b/tests/backends/test_awkward.py @@ -9,11 +9,58 @@ import vector -pytest.importorskip("awkward") +ak = pytest.importorskip("awkward") pytestmark = pytest.mark.awkward +def test_dimension_conversion(): + # 2D -> 3D + vec = vector.Array( + [ + [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.1}], + [], + ] + ) + assert ak.all(vec.to_Vector3D(z=1).z == 1) + assert ak.all(vec.to_Vector3D(eta=1).eta == 1) + assert ak.all(vec.to_Vector3D(theta=1).theta == 1) + + assert ak.all(vec.to_Vector3D(z=1).x == vec.x) + assert ak.all(vec.to_Vector3D(z=1).y == vec.y) + + # 2D -> 4D + assert ak.all(vec.to_Vector4D(z=1, t=1).t == 1) + assert ak.all(vec.to_Vector4D(z=1, t=1).z == 1) + assert ak.all(vec.to_Vector4D(eta=1, t=1).eta == 1) + assert ak.all(vec.to_Vector4D(eta=1, t=1).t == 1) + assert ak.all(vec.to_Vector4D(theta=1, t=1).theta == 1) + assert ak.all(vec.to_Vector4D(theta=1, t=1).t == 1) + assert ak.all(vec.to_Vector4D(z=1, tau=1).z == 1) + assert ak.all(vec.to_Vector4D(z=1, tau=1).tau == 1) + assert ak.all(vec.to_Vector4D(eta=1, tau=1).eta == 1) + assert ak.all(vec.to_Vector4D(eta=1, tau=1).tau == 1) + assert ak.all(vec.to_Vector4D(theta=1, tau=1).theta == 1) + assert ak.all(vec.to_Vector4D(theta=1, tau=1).tau == 1) + + assert ak.all(vec.to_Vector4D(z=1, t=1).x == vec.x) + assert ak.all(vec.to_Vector4D(z=1, t=1).y == vec.y) + + # 3D -> 4D + vec = vector.Array( + [ + [{"x": 1, "y": 1.1, "z": 1.2}, {"x": 2, "y": 2.1, "z": 2.2}], + [], + ] + ) + assert ak.all(vec.to_Vector4D(t=1).t == 1) + assert ak.all(vec.to_Vector4D(tau=1).tau == 1) + + assert ak.all(vec.to_Vector4D(t=1).x == vec.x) + assert ak.all(vec.to_Vector4D(t=1).y == vec.y) + assert ak.all(vec.to_Vector4D(t=1).z == vec.z) + + def test_type_checks(): with pytest.raises(TypeError): vector.Array( diff --git a/tests/backends/test_numpy.py b/tests/backends/test_numpy.py index 89d433eb..47a52727 100644 --- a/tests/backends/test_numpy.py +++ b/tests/backends/test_numpy.py @@ -14,6 +14,49 @@ 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( diff --git a/tests/backends/test_object.py b/tests/backends/test_object.py index bad95740..a976611f 100644 --- a/tests/backends/test_object.py +++ b/tests/backends/test_object.py @@ -11,6 +11,43 @@ import vector +def test_dimension_conversion(): + # 2D -> 3D + vec = vector.VectorObject2D(x=1, y=2) + assert vec.to_Vector3D(z=1).z == 1 + assert vec.to_Vector3D(eta=1).eta == 1 + assert vec.to_Vector3D(theta=1).theta == 1 + + assert vec.to_Vector3D(z=1).x == vec.x + assert vec.to_Vector3D(z=1).y == vec.y + + # 2D -> 4D + assert vec.to_Vector4D(z=1, t=1).z == 1 + assert vec.to_Vector4D(z=1, t=1).t == 1 + assert vec.to_Vector4D(eta=1, t=1).eta == 1 + assert vec.to_Vector4D(eta=1, t=1).t == 1 + assert vec.to_Vector4D(theta=1, t=1).theta == 1 + assert vec.to_Vector4D(theta=1, t=1).t == 1 + assert vec.to_Vector4D(z=1, tau=1).z == 1 + assert vec.to_Vector4D(z=1, tau=1).tau == 1 + assert vec.to_Vector4D(eta=1, tau=1).eta == 1 + assert vec.to_Vector4D(eta=1, tau=1).tau == 1 + assert vec.to_Vector4D(theta=1, tau=1).theta == 1 + assert vec.to_Vector4D(theta=1, tau=1).tau == 1 + + assert vec.to_Vector4D(z=1, t=1).x == vec.x + assert vec.to_Vector4D(z=1, t=1).y == vec.y + + # 3D -> 4D + vec = vector.VectorObject3D(x=1, y=2, z=3) + assert vec.to_Vector4D(t=1).t == 1 + assert vec.to_Vector4D(tau=1).tau == 1 + + assert vec.to_Vector4D(t=1).x == vec.x + assert vec.to_Vector4D(t=1).y == vec.y + assert vec.to_Vector4D(t=1).z == vec.z + + def test_constructors_2D(): vec = vector.VectorObject2D(x=1, y=2) assert vec.x == 1