feat: short names for to_VectorND methods (#432)

* feat: short names for to_VectorND methods * Test the alias methods
scikit-hep · Mar 5, 2024 · 50ea01f · 50ea01f
1 parent e219431
commit 50ea01f
Show file tree

Hide file tree

Showing 5 changed files with 140 additions and 34 deletions.
diff --git a/README.md b/README.md
@@ -630,7 +630,7 @@ The (current) list of properties and methods is:
 - `count_nonzero()`: can also use `numpy.count_nonzero` or `awkward.count_nonzero`, only for NumPy and Awkward vectors
 - `count()`: can also use `awkward.count`, only for Awkward vectors
 - `isclose(vector, rtol=1e-5, atol=1e-8, equal_nan=False)`: works like [np.isclose](https://numpy.org/doc/stable/reference/generated/numpy.isclose.html); arrays also have an [allclose](https://numpy.org/doc/stable/reference/generated/numpy.allclose.html) method
-- `to_Vector*D(coordinates)`: replace `*` with the reuquired vector dimension
+- `to_VectorND(coordinates)`/`to_ND(coordinates)`: replace `N` with the required vector dimension
 - `to_{coordinate-names}`: for example - `to_rhophietatau`
 - `like(other)`: projects the vector into the dimensions of `other`, for example - `two_d_vector.like(three_d_vector)`
 

diff --git a/docs/usage/intro.ipynb b/docs/usage/intro.ipynb
@@ -2616,7 +2616,7 @@
     "   * `count_nonzero()`: can also use `numpy.count_nonzero` or `awkward.count_nonzero`, only for NumPy and Awkward vectors\n",
     "   * `count()`: can also use `awkward.count`, only for Awkward vectors\n",
     "   * `isclose(vector, rtol=1e-5, atol=1e-8, equal_nan=False)`: works like [np.isclose](https://numpy.org/doc/stable/reference/generated/numpy.isclose.html); arrays also have an [allclose](https://numpy.org/doc/stable/reference/generated/numpy.allclose.html) method\n",
-    "   * `to_Vector*D(coordinates)`: replace `*` with the reuquired vector dimension\n",
+    "   * `to_VectorND(coordinates)`/`to_ND(coordinates)`: replace `N` with the required vector dimension\n",
     "   * `to_{coordinate-names}`: for example - `to_rhophietatau`\n",
     "   * `like(other)`: projects the vector into the dimensions of `other`, for example - `two_d_vector.like(three_d_vector)`\n"
    ]

diff --git a/src/vector/_methods.py b/src/vector/_methods.py
@@ -204,6 +204,43 @@ def to_Vector4D(self) -> VectorProtocolLorentz:
         """
         raise AssertionError
 
+    def to_2D(self) -> VectorProtocolPlanar:
+        """
+        Projects this vector/these vectors onto azimuthal coordinates only.
+
+        Alias for :meth:`vector._methods.VectorProtocol.to_Vector2D`.
+        """
+        raise AssertionError
+
+    def to_3D(self) -> VectorProtocolSpatial:
+        """
+        Projects this vector/these vectors onto azimuthal and longitudinal
+        coordinates only.
+
+        If 2D, a default $z$ component of $0$ is imputed.
+
+        The longitudinal coordinate can be passed as a named argument.
+
+        Alias for :meth:`vector._methods.VectorProtocol.to_Vector3D`.
+        """
+        raise AssertionError
+
+    def to_4D(self) -> VectorProtocolLorentz:
+        """
+        Projects this vector/these vectors onto azimuthal, longitudinal,
+        and temporal coordinates.
+
+        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$.
+
+        The longitudinal and temporal coordinates can be passed as named arguments.
+
+        Alias for :meth:`vector._methods.VectorProtocol.to_Vector4D`.
+        """
+        raise AssertionError
+
     def to_xy(self) -> VectorProtocolPlanar:
         """
         Converts to $x$-$y$ coordinates, possibly eliminating dimensions with a
@@ -3322,6 +3359,24 @@ def to_Vector4D(
             1,
         )
 
+    def to_2D(self) -> VectorProtocolPlanar:
+        """
+        Alias for :meth:`vector._methods.Vector2D.to_Vector2D`.
+        """
+        return self.to_Vector2D()
+
+    def to_3D(self, **kwargs: float | None) -> VectorProtocolSpatial:
+        """
+        Alias for :meth:`vector._methods.Vector2D.to_Vector3D`.
+        """
+        return self.to_Vector3D(**kwargs)
+
+    def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz:
+        """
+        Alias for :meth:`vector._methods.Vector2D.to_Vector4D`.
+        """
+        return self.to_Vector4D(**kwargs)
+
 
 class Vector3D(Vector, VectorProtocolSpatial):
     def to_Vector2D(self) -> VectorProtocolPlanar:
@@ -3382,6 +3437,24 @@ def to_Vector4D(
             1,
         )
 
+    def to_2D(self) -> VectorProtocolPlanar:
+        """
+        Alias for :meth:`vector._methods.Vector3D.to_Vector2D`.
+        """
+        return self.to_Vector2D()
+
+    def to_3D(self) -> VectorProtocolSpatial:
+        """
+        Alias for :meth:`vector._methods.Vector3D.to_Vector3D`.
+        """
+        return self.to_Vector3D()
+
+    def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz:
+        """
+        Alias for :meth:`vector._methods.Vector3D.to_Vector4D`.
+        """
+        return self.to_Vector4D(**kwargs)
+
 
 class Vector4D(Vector, VectorProtocolLorentz):
     def to_Vector2D(self) -> VectorProtocolPlanar:
@@ -3403,6 +3476,24 @@ def to_Vector3D(self) -> VectorProtocolSpatial:
     def to_Vector4D(self) -> VectorProtocolLorentz:
         return self
 
+    def to_2D(self) -> VectorProtocolPlanar:
+        """
+        Alias for :meth:`vector._methods.Vector4D.to_Vector2D`.
+        """
+        return self.to_Vector2D()
+
+    def to_3D(self) -> VectorProtocolSpatial:
+        """
+        Alias for :meth:`vector._methods.Vector4D.to_Vector3D`.
+        """
+        return self.to_Vector3D()
+
+    def to_4D(self) -> VectorProtocolLorentz:
+        """
+        Alias for :meth:`vector._methods.Vector4D.to_Vector4D`.
+        """
+        return self.to_Vector4D()
+
 
 class Planar(VectorProtocolPlanar):
     @property

diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py
@@ -32,9 +32,10 @@ def test_dimension_conversion():
             [],
         ]
     )
-    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)
+    # test alias
+    assert ak.all(vec.to_3D(z=1).z == 1)
+    assert ak.all(vec.to_3D(eta=1).eta == 1)
+    assert ak.all(vec.to_3D(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)
@@ -53,8 +54,9 @@ def test_dimension_conversion():
     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)
+    # test alias
+    assert ak.all(vec.to_4D(z=1, t=1).x == vec.x)
+    assert ak.all(vec.to_4D(z=1, t=1).y == vec.y)
 
     # 3D -> 4D
     vec = vector.Array(

diff --git a/tests/compute/test_conversions.py b/tests/compute/test_conversions.py
@@ -500,8 +500,9 @@ def test_conversion_with_coords_object():
     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
+    # test alias
+    assert vec.to_3D(z=1).x == vec.x
+    assert vec.to_3D(z=1).y == vec.y
 
     # 2D -> 4D
     assert vec.to_Vector4D(z=1, t=1).z == 1
@@ -517,13 +518,16 @@ def test_conversion_with_coords_object():
     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
+    # test alias
+    assert vec.to_4D(z=1, t=1).x == vec.x
+    assert vec.to_4D(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
+
+    # test alias
+    assert vec.to_4D(t=1).t == 1
+    assert vec.to_4D(tau=1).tau == 1
 
     assert vec.to_Vector4D(t=1).x == vec.x
     assert vec.to_Vector4D(t=1).y == vec.y
@@ -533,18 +537,21 @@ def test_conversion_with_coords_object():
     vec = vector.MomentumObject2D(px=1, py=2)
     assert vec.to_Vector3D(pz=1).pz == 1
 
-    assert vec.to_Vector4D(pz=1, m=1).pz == 1
-    assert vec.to_Vector4D(pz=1, m=1).m == 1
-    assert vec.to_Vector4D(pz=1, mass=1).mass == 1
-    assert vec.to_Vector4D(pz=1, M=1).M == 1
+    # test both alias and original methods
+    assert vec.to_4D(pz=1, m=1).pz == 1
+    assert vec.to_4D(pz=1, m=1).m == 1
+    assert vec.to_4D(pz=1, mass=1).mass == 1
+    assert vec.to_4D(pz=1, M=1).M == 1
     assert vec.to_Vector4D(pz=1, e=1).e == 1
     assert vec.to_Vector4D(pz=1, energy=1).energy == 1
     assert vec.to_Vector4D(pz=1, E=1).E == 1
 
     vec = vector.MomentumObject3D(px=1, py=2, pz=3)
-    assert vec.to_Vector4D(m=1).m == 1
-    assert vec.to_Vector4D(mass=1).mass == 1
-    assert vec.to_Vector4D(M=1).M == 1
+
+    # test both alias and original methods
+    assert vec.to_4D(m=1).m == 1
+    assert vec.to_4D(mass=1).mass == 1
+    assert vec.to_4D(M=1).M == 1
     assert vec.to_Vector4D(e=1).e == 1
     assert vec.to_Vector4D(energy=1).energy == 1
     assert vec.to_Vector4D(E=1).E == 1
@@ -560,8 +567,9 @@ def test_conversion_with_coords_numpy():
     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)
+    # test alias
+    assert all(vec.to_3D(z=1).x == vec.x)
+    assert all(vec.to_3D(z=1).y == vec.y)
 
     # 2D -> 4D
     assert all(vec.to_Vector4D(z=1, t=1).t == 1)
@@ -577,8 +585,9 @@ def test_conversion_with_coords_numpy():
     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)
+    # test alias
+    assert all(vec.to_4D(z=1, t=1).x == vec.x)
+    assert all(vec.to_4D(z=1, t=1).y == vec.y)
 
     # 3D -> 4D
     vec = vector.VectorNumpy3D(
@@ -588,9 +597,10 @@ def test_conversion_with_coords_numpy():
     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)
+    # test alias
+    assert all(vec.to_4D(t=1).x == vec.x)
+    assert all(vec.to_4D(t=1).y == vec.y)
+    assert all(vec.to_4D(t=1).z == vec.z)
 
     # check if momentum coords work
     vec = vector.MomentumNumpy2D(
@@ -599,10 +609,11 @@ def test_conversion_with_coords_numpy():
     )
     assert all(vec.to_Vector3D(pz=1).pz == 1)
 
-    assert all(vec.to_Vector4D(pz=1, m=1).pz == 1)
-    assert all(vec.to_Vector4D(pz=1, m=1).m == 1)
-    assert all(vec.to_Vector4D(pz=1, mass=1).mass == 1)
-    assert all(vec.to_Vector4D(pz=1, M=1).M == 1)
+    # test both alias and original methods
+    assert all(vec.to_4D(pz=1, m=1).pz == 1)
+    assert all(vec.to_4D(pz=1, m=1).m == 1)
+    assert all(vec.to_4D(pz=1, mass=1).mass == 1)
+    assert all(vec.to_4D(pz=1, M=1).M == 1)
     assert all(vec.to_Vector4D(pz=1, e=1).e == 1)
     assert all(vec.to_Vector4D(pz=1, energy=1).energy == 1)
     assert all(vec.to_Vector4D(pz=1, E=1).E == 1)
@@ -611,9 +622,11 @@ def test_conversion_with_coords_numpy():
         [(1.0, 1.0, 1.0), (2.0, 2.0, 2.0)],
         dtype=[("px", float), ("py", float), ("pz", float)],
     )
-    assert all(vec.to_Vector4D(m=1).m == 1)
-    assert all(vec.to_Vector4D(mass=1).mass == 1)
-    assert all(vec.to_Vector4D(M=1).M == 1)
+
+    # test both alias and original methods
+    assert all(vec.to_4D(m=1).m == 1)
+    assert all(vec.to_4D(mass=1).mass == 1)
+    assert all(vec.to_4D(M=1).M == 1)
     assert all(vec.to_Vector4D(e=1).e == 1)
     assert all(vec.to_Vector4D(energy=1).energy == 1)
     assert all(vec.to_Vector4D(E=1).E == 1)