diff --git a/README.md b/README.md index 7cdb0769..437af103 100644 --- 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 index f289a94f..89bfcb66 100644 --- 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 index 8c8d3829..528059b9 100644 --- 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 index cf8a757f..bfa8614c 100644 --- 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 index 8aa6f2b9..f8dd1392 100644 --- 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)