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 <[email protected]> * style: pre-commit fixes * Better error message --------- Co-authored-by: Jim Pivarski <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
scikit-hep · Feb 23, 2023 · 5d54670 · 5d54670
1 parent 17af58f
commit 5d54670
Show file tree

Hide file tree

Showing 5 changed files with 246 additions and 14 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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