diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 7511cda6..d8febb3e 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -206,6 +206,13 @@ def to_xy(self) -> VectorProtocolPlanar: """ raise AssertionError + def to_pxpy(self) -> VectorProtocolPlanar: + """ + Converts to $px$-$py$ coordinates, possibly eliminating dimensions with a + projection. + """ + raise AssertionError + def to_rhophi(self) -> VectorProtocolPlanar: r""" Converts to $\rho$-$\phi$ coordinates, possibly eliminating dimensions with a @@ -213,6 +220,13 @@ def to_rhophi(self) -> VectorProtocolPlanar: """ raise AssertionError + def to_ptphi(self) -> VectorProtocolPlanar: + r""" + Converts to $pt$-$\phi$ coordinates, possibly eliminating dimensions with a + projection. + """ + raise AssertionError + def to_xyz(self) -> VectorProtocolSpatial: """ Converts to $x$-$y$-$z$ coordinates, possibly eliminating or imputing @@ -234,6 +248,27 @@ def to_xyeta(self) -> VectorProtocolSpatial: """ raise AssertionError + def to_pxpypz(self) -> VectorProtocolSpatial: + """ + Converts to $px$-$py$-$pz$ coordinates, possibly eliminating or imputing + dimensions with a projection. + """ + raise AssertionError + + def to_pxpytheta(self) -> VectorProtocolSpatial: + r""" + Converts to $px$-$py$-$\theta$ coordinates, possibly eliminating or imputing + dimensions with a projection. + """ + raise AssertionError + + def to_pxpyeta(self) -> VectorProtocolSpatial: + r""" + Converts to $px$-$py$-$\eta$ coordinates, possibly eliminating or imputing + dimensions with a projection. + """ + raise AssertionError + def to_rhophiz(self) -> VectorProtocolSpatial: r""" Converts to $\rho$-$\phi$-$z$ coordinates, possibly eliminating or imputing @@ -255,6 +290,27 @@ def to_rhophieta(self) -> VectorProtocolSpatial: """ raise AssertionError + def to_ptphipz(self) -> VectorProtocolSpatial: + r""" + Converts to $pt$-$\phi$-$pz$ coordinates, possibly eliminating or imputing + dimensions with a projection. + """ + raise AssertionError + + def to_ptphitheta(self) -> VectorProtocolSpatial: + r""" + Converts to $pt$-$\phi$-$\theta$ coordinates, possibly eliminating or + imputing dimensions with a projection. + """ + raise AssertionError + + def to_ptphieta(self) -> VectorProtocolSpatial: + r""" + Converts to $pt$-$\phi$-$\eta$ coordinates, possibly eliminating or + imputing dimensions with a projection. + """ + raise AssertionError + def to_xyzt(self) -> VectorProtocolLorentz: """ Converts to $x$-$y$-$z$-$t$ coordinates, possibly imputing dimensions with @@ -297,6 +353,48 @@ def to_xyetatau(self) -> VectorProtocolLorentz: """ raise AssertionError + def to_pxpypzenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$pz$-$energy$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_pxpythetaenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$\theta$-$energy$ coordinates, possibly imputing + dimensions with a projection. + """ + raise AssertionError + + def to_pxpyetaenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$\eta$-$energy$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_pxpypzmass(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$pz$-$mass$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_pxpythetamass(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$\theta$-$energy$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_pxpyetamass(self) -> VectorProtocolLorentz: + r""" + Converts to $px$-$py$-$\eta$-$mass$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + def to_rhophizt(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$z$-$t$ coordinates, possibly imputing dimensions @@ -339,6 +437,48 @@ def to_rhophietatau(self) -> VectorProtocolLorentz: """ raise AssertionError + def to_ptphipzenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$pz$-$energy$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_ptphithetaenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$\theta$-$energy$ coordinates, possibly imputing + dimensions with a projection. + """ + raise AssertionError + + def to_ptphietaenergy(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$\eta$-$energy$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_ptphipzmass(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$pz$-$mass$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_ptphithetamass(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$\theta$-$mass$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + + def to_ptphietamass(self) -> VectorProtocolLorentz: + r""" + Converts to $pt$-$\phi$-$\theta$-$mass$ coordinates, possibly imputing dimensions + with a projection. + """ + raise AssertionError + def unit(self: SameVectorType) -> SameVectorType: """ Returns vector(s) normalized to unit length, which is `rho == 1` for 2D @@ -2656,6 +2796,9 @@ def to_xy(self) -> VectorProtocolPlanar: 1, ) + def to_pxpy(self) -> VectorProtocolPlanar: + return self.to_xy() + def to_rhophi(self) -> VectorProtocolPlanar: from vector._compute import planar @@ -2666,6 +2809,9 @@ def to_rhophi(self) -> VectorProtocolPlanar: 1, ) + def to_ptphi(self) -> VectorProtocolPlanar: + return self.to_rhophi() + def to_xyz(self) -> VectorProtocolSpatial: from vector._compute import planar, spatial @@ -2708,6 +2854,15 @@ def to_xyeta(self) -> VectorProtocolSpatial: 1, ) + def to_pxpypz(self) -> VectorProtocolSpatial: + return self.to_xyz() + + def to_pxpytheta(self) -> VectorProtocolSpatial: + return self.to_xytheta() + + def to_pxpyeta(self) -> VectorProtocolSpatial: + return self.to_xyeta() + def to_rhophiz(self) -> VectorProtocolSpatial: from vector._compute import planar, spatial @@ -2750,6 +2905,15 @@ def to_rhophieta(self) -> VectorProtocolSpatial: 1, ) + def to_ptphipz(self) -> VectorProtocolSpatial: + return self.to_rhophiz() + + def to_ptphitheta(self) -> VectorProtocolSpatial: + return self.to_rhophitheta() + + def to_ptphieta(self) -> VectorProtocolSpatial: + return self.to_rhophieta() + def to_xyzt(self) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial @@ -2852,6 +3016,24 @@ def to_xyetatau(self) -> VectorProtocolLorentz: 1, ) + def to_pxpypzenergy(self) -> VectorProtocolLorentz: + return self.to_xyzt() + + def to_pxpythetaenergy(self) -> VectorProtocolLorentz: + return self.to_xythetat() + + def to_pxpyetaenergy(self) -> VectorProtocolLorentz: + return self.to_xyetat() + + def to_pxpypzmass(self) -> VectorProtocolLorentz: + return self.to_xyztau() + + def to_pxpythetamass(self) -> VectorProtocolLorentz: + return self.to_xythetatau() + + def to_pxpyetamass(self) -> VectorProtocolLorentz: + return self.to_xyetatau() + def to_rhophizt(self) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial @@ -2954,6 +3136,24 @@ def to_rhophietatau(self) -> VectorProtocolLorentz: 1, ) + def to_ptphipzenergy(self) -> VectorProtocolLorentz: + return self.to_rhophizt() + + def to_ptphithetaenergy(self) -> VectorProtocolLorentz: + return self.to_rhophithetat() + + def to_ptphietaenergy(self) -> VectorProtocolLorentz: + return self.to_rhophietat() + + def to_ptphipzmass(self) -> VectorProtocolLorentz: + return self.to_rhophiztau() + + def to_ptphithetamass(self) -> VectorProtocolLorentz: + return self.to_rhophithetatau() + + def to_ptphietamass(self) -> VectorProtocolLorentz: + return self.to_rhophietatau() + class Vector2D(Vector, VectorProtocolPlanar): def to_Vector2D(self) -> VectorProtocolPlanar: diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py index 9797be1b..1c1ca931 100644 --- a/tests/backends/test_awkward.py +++ b/tests/backends/test_awkward.py @@ -11,6 +11,7 @@ import pytest import vector +from vector import VectorObject2D, VectorObject3D, VectorObject4D ak = pytest.importorskip("awkward") @@ -700,3 +701,175 @@ def test_demotion(): assert all(v3 + v1 == v1_v3) assert all(v2 + v3 == v2_v3) assert all(v3 + v2 == v2_v3) + + +def test_handler_of(): + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(awkward_a, object_b) + # chooses awkward backend and converts the vector to 2D + assert all(protocol == awkward_a.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a) + # chooses awkward backend and the vector is already of the + # lower dimension + assert all(protocol == awkward_a) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ) + awkward_b = vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) + # chooses awkward backend and the 2D awkward vector + # (first encountered awkward vector) + assert all(protocol == awkward_a) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [-10.0, 20.0, 30.0], + }, + ) + awkward_b = vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(awkward_b, object_b, awkward_a) + # chooses awkward backend and converts awkward_b to 2D + # (first encountered awkward vector) + assert all(protocol == awkward_b.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + awkward_b = vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) + # chooses awkward backend and converts the vector to 2D + # (the first awkward vector encountered is used as the base) + assert all(protocol == awkward_a.to_Vector2D()) + + numpy_a = vector.array( + { + "x": [1.1, 1.2, 1.3, 1.4, 1.5], + "y": [2.1, 2.2, 2.3, 2.4, 2.5], + "z": [3.1, 3.2, 3.3, 3.4, 3.5], + } + ) + awkward_b = vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_b, numpy_a, awkward_b) + # chooses awkward backend and converts the vector to 2D + assert all(protocol == awkward_b.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + numpy_a = vector.array( + { + "x": [1.1, 1.2, 1.3, 1.4, 1.5], + "y": [2.1, 2.2, 2.3, 2.4, 2.5], + } + ) + awkward_b = vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject3D.from_xyz(1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b, numpy_a) + # chooses awkward backend and converts the vector to 2D + # (the first awkward vector encountered is used as the base) + assert all(protocol == awkward_a.to_Vector2D()) + + +def test_momentum_coordinate_transforms(): + awkward_vec = vector.zip( + { + "px": [1.0, 2.0, 3.0], + "py": [-1.0, 2.0, 3.0], + }, + ) + + for t1 in "pxpy", "ptphi": + for t2 in "pz", "eta", "theta": + for t3 in "mass", "energy": + transformed_object = getattr(awkward_vec, "to_" + t1)() + assert isinstance( + transformed_object, vector.backends.awkward.MomentumAwkward2D + ) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + + transformed_object = getattr(awkward_vec, "to_" + t1 + t2)() + assert isinstance( + transformed_object, vector.backends.awkward.MomentumAwkward3D + ) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + assert hasattr(transformed_object, t2) + + transformed_object = getattr(awkward_vec, "to_" + t1 + t2 + t3)() + assert isinstance( + transformed_object, vector.backends.awkward.MomentumAwkward4D + ) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + assert hasattr(transformed_object, t2) + assert hasattr(transformed_object, t3) diff --git a/tests/test_methods.py b/tests/test_methods.py index f5b5762d..b31c9977 100644 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -5,12 +5,18 @@ from __future__ import annotations -import pytest - import vector -from vector import VectorObject2D, VectorObject3D, VectorObject4D - -awkward = pytest.importorskip("awkward") +from vector import ( + MomentumNumpy2D, + MomentumNumpy3D, + MomentumNumpy4D, + MomentumObject2D, + MomentumObject3D, + MomentumObject4D, + VectorObject2D, + VectorObject3D, + VectorObject4D, +) def test_handler_of(): @@ -49,136 +55,51 @@ def test_handler_of(): protocol = vector._methods._handler_of(object_a, object_b) assert protocol == object_b - awkward_a = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject2D.from_xy(1.0, 1.0) - protocol = vector._methods._handler_of(awkward_a, object_b) - # chooses awkward backend and converts the vector to 2D - assert all(protocol == awkward_a.to_Vector2D()) - - awkward_a = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - }, - ) - object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) - protocol = vector._methods._handler_of(object_b, awkward_a) - # chooses awkward backend and the vector is already of the - # lower dimension - assert all(protocol == awkward_a) - - awkward_a = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - }, - ) - awkward_b = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) - protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) - # chooses awkward backend and the 2D awkward vector - # (first encountered awkward vector) - assert all(protocol == awkward_a) - - awkward_a = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [-10.0, 20.0, 30.0], - }, - ) - awkward_b = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject2D.from_xy(1.0, 1.0) - protocol = vector._methods._handler_of(awkward_b, object_b, awkward_a) - # chooses awkward backend and converts awkward_b to 2D - # (first encountered awkward vector) - assert all(protocol == awkward_b.to_Vector2D()) - - awkward_a = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 1.0, 1.0], - }, - ) - awkward_b = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject2D.from_xy(1.0, 1.0) - protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) - # chooses awkward backend and converts the vector to 2D - # (the first awkward vector encountered is used as the base) - assert all(protocol == awkward_a.to_Vector2D()) - - numpy_a = vector.array( - { - "x": [1.1, 1.2, 1.3, 1.4, 1.5], - "y": [2.1, 2.2, 2.3, 2.4, 2.5], - "z": [3.1, 3.2, 3.3, 3.4, 3.5], - } - ) - awkward_b = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject2D.from_xy(1.0, 1.0) - protocol = vector._methods._handler_of(object_b, numpy_a, awkward_b) - # chooses awkward backend and converts the vector to 2D - assert all(protocol == awkward_b.to_Vector2D()) - awkward_a = vector.zip( +def test_momentum_coordinate_transforms(): + numpy_vec = vector.array( { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 1.0, 1.0], + "px": [1.0, 2.0, 3.0], + "py": [-1.0, 2.0, 3.0], }, ) - numpy_a = vector.array( - { - "x": [1.1, 1.2, 1.3, 1.4, 1.5], - "y": [2.1, 2.2, 2.3, 2.4, 2.5], - } - ) - awkward_b = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, - ) - object_b = VectorObject3D.from_xyz(1.0, 1.0, 1.0) - protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b, numpy_a) - # chooses awkward backend and converts the vector to 2D - # (the first awkward vector encountered is used as the base) - assert all(protocol == awkward_a.to_Vector2D()) + object_vec = MomentumObject2D(px=0.0, py=0.0) + + for t1 in "pxpy", "ptphi": + for t2 in "pz", "eta", "theta": + for t3 in "mass", "energy": + transformed_object = getattr(object_vec, "to_" + t1)() + assert isinstance(transformed_object, MomentumObject2D) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + + transformed_object = getattr(object_vec, "to_" + t1 + t2)() + assert isinstance(transformed_object, MomentumObject3D) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + assert hasattr(transformed_object, t2) + + transformed_object = getattr(object_vec, "to_" + t1 + t2 + t3)() + assert isinstance(transformed_object, MomentumObject4D) + assert hasattr(transformed_object, t1[:2]) + assert hasattr(transformed_object, t1[2:]) + assert hasattr(transformed_object, t2) + assert hasattr(transformed_object, t3) + + transformed_numpy = getattr(numpy_vec, "to_" + t1)() + assert isinstance(transformed_numpy, MomentumNumpy2D) + assert hasattr(transformed_numpy, t1[:2]) + assert hasattr(transformed_numpy, t1[2:]) + + transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2)() + assert isinstance(transformed_numpy, MomentumNumpy3D) + assert hasattr(transformed_numpy, t1[:2]) + assert hasattr(transformed_numpy, t1[2:]) + assert hasattr(transformed_numpy, t2) + + transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2 + t3)() + assert isinstance(transformed_numpy, MomentumNumpy4D) + assert hasattr(transformed_numpy, t1[:2]) + assert hasattr(transformed_numpy, t1[2:]) + assert hasattr(transformed_numpy, t2) + assert hasattr(transformed_numpy, t3)