feat: implement to_2D/3D/4D, momentum conversions, and like in numba object backend (#750)

src/vector/backends/_numba_object.py

@@ -763,6 +763,7 @@ def vector_obj_impl(
 # properties and methods ######################################################
 
 
+@numba.extending.overload_method(VectorObject2DType, "to_2D")
 @numba.extending.overload_method(VectorObject2DType, "to_Vector2D")
 def VectorObject2D_to_Vector2D(v):
     def VectorObject2D_to_Vector2D_impl(v):
@@ -771,6 +772,7 @@ def VectorObject2D_to_Vector2D_impl(v):
     return VectorObject2D_to_Vector2D_impl
 
 
+@numba.extending.overload_method(VectorObject2DType, "to_3D")
 @numba.extending.overload_method(VectorObject2DType, "to_Vector3D")
 def VectorObject2D_to_Vector3D(v):
     if issubclass(v.instance_class, Momentum):
@@ -786,6 +788,7 @@ def VectorObject2D_to_Vector3D_impl(v):
     return VectorObject2D_to_Vector3D_impl
 
 
+@numba.extending.overload_method(VectorObject2DType, "to_4D")
 @numba.extending.overload_method(VectorObject2DType, "to_Vector4D")
 def VectorObject2D_to_Vector4D(v):
     if issubclass(v.instance_class, Momentum):
@@ -805,6 +808,7 @@ def VectorObject2D_to_Vector4D_impl(v):
     return VectorObject2D_to_Vector4D_impl
 
 
+@numba.extending.overload_method(VectorObject3DType, "to_2D")
 @numba.extending.overload_method(VectorObject3DType, "to_Vector2D")
 def VectorObject3D_to_Vector2D(v):
     if issubclass(v.instance_class, Momentum):
@@ -820,6 +824,7 @@ def VectorObject3D_to_Vector2D_impl(v):
     return VectorObject3D_to_Vector2D_impl
 
 
+@numba.extending.overload_method(VectorObject3DType, "to_3D")
 @numba.extending.overload_method(VectorObject3DType, "to_Vector3D")
 def VectorObject3D_to_Vector3D(v):
     def VectorObject3D_to_Vector3D_impl(v):
@@ -828,6 +833,7 @@ def VectorObject3D_to_Vector3D_impl(v):
     return VectorObject3D_to_Vector3D_impl
 
 
+@numba.extending.overload_method(VectorObject3DType, "to_4D")
 @numba.extending.overload_method(VectorObject3DType, "to_Vector4D")
 def VectorObject3D_to_Vector4D(v):
     if issubclass(v.instance_class, Momentum):
@@ -843,6 +849,7 @@ def VectorObject3D_to_Vector4D_impl(v):
     return VectorObject3D_to_Vector4D_impl
 
 
+@numba.extending.overload_method(VectorObject4DType, "to_2D")
 @numba.extending.overload_method(VectorObject4DType, "to_Vector2D")
 def VectorObject4D_to_Vector2D(v):
     if issubclass(v.instance_class, Momentum):
@@ -858,6 +865,7 @@ def VectorObject4D_to_Vector2D_impl(v):
     return VectorObject4D_to_Vector2D_impl
 
 
+@numba.extending.overload_method(VectorObject4DType, "to_3D")
 @numba.extending.overload_method(VectorObject4DType, "to_Vector3D")
 def VectorObject4D_to_Vector3D(v):
     if issubclass(v.instance_class, Momentum):
@@ -873,6 +881,7 @@ def VectorObject4D_to_Vector3D_impl(v):
     return VectorObject4D_to_Vector3D_impl
 
 
+@numba.extending.overload_method(VectorObject4DType, "to_4D")
 @numba.extending.overload_method(VectorObject4DType, "to_Vector4D")
 def VectorObject4D_to_Vector4D(v):
     def VectorObject4D_to_Vector4D_impl(v):
@@ -881,6 +890,35 @@ def VectorObject4D_to_Vector4D_impl(v):
     return VectorObject4D_to_Vector4D_impl
 
 
+@numba.extending.overload_method(VectorObject2DType, "like")
+@numba.extending.overload_method(VectorObject3DType, "like")
+@numba.extending.overload_method(VectorObject4DType, "like")
+def VectorObject_like(v, other):
+    if isinstance(other, VectorObject2DType):
+
+        def VectorObject_like_2D_impl(v, other):
+            return v.to_Vector2D()
+
+        return VectorObject_like_2D_impl
+
+    elif isinstance(other, VectorObject3DType):
+
+        def VectorObject_like_3D_impl(v, other):
+            return v.to_Vector3D()
+
+        return VectorObject_like_3D_impl
+
+    elif isinstance(other, VectorObject4DType):
+
+        def VectorObject_like_4D_impl(v, other):
+            return v.to_Vector4D()
+
+        return VectorObject_like_4D_impl
+
+    else:
+        raise numba.TypingError(f"'{other}' is not a vector.Vector")
+
+
 @numba.jit(nopython=True)
 def make_AzimuthalObjectXY(v):
     return AzimuthalObjectXY(v.x, v.y)
@@ -968,22 +1006,31 @@ def make_TemporalObjectTau_zero(v):
 
 def add_coordinate_change(vectortype, azcoordtype, lcoordtype, tcoordtype):
     methodname = "to_"
+    momentumname = "to_"
     if azcoordtype is AzimuthalObjectXY:
         methodname += "xy"
+        momentumname += "pxpy"
     elif azcoordtype is AzimuthalObjectRhoPhi:
         methodname += "rhophi"
+        momentumname += "ptphi"
     if lcoordtype is LongitudinalObjectZ:
         methodname += "z"
-    if lcoordtype is LongitudinalObjectTheta:
+        momentumname += "pz"
+    elif lcoordtype is LongitudinalObjectTheta:
         methodname += "theta"
-    if lcoordtype is LongitudinalObjectEta:
+        momentumname += "theta"
+    elif lcoordtype is LongitudinalObjectEta:
         methodname += "eta"
+        momentumname += "eta"
     if tcoordtype is TemporalObjectT:
         methodname += "t"
-    if tcoordtype is TemporalObjectTau:
+        momentumname += "energy"
+    elif tcoordtype is TemporalObjectTau:
         methodname += "tau"
+        momentumname += "mass"
 
     @numba.extending.overload_method(vectortype, methodname)
+    @numba.extending.overload_method(vectortype, momentumname)
     def overloader(v):
         if lcoordtype is None and tcoordtype is None:
             azcoords = make_coordobject[vectortype, azcoordtype]

tests/backends/test_numba_object.py

@@ -1800,3 +1800,117 @@ def get_mt2(v):
     assert get_et2(p) == pytest.approx(p.et2)
     assert get_mt(p) == pytest.approx(p.mt)
     assert get_mt2(p) == pytest.approx(p.mt2)
+
+
+def test_dimension_conversions_numbered():
+    @numba.njit
+    def to_2d(v):
+        return v.to_2D()
+
+    @numba.njit
+    def to_3d(v):
+        return v.to_3D()
+
+    @numba.njit
+    def to_4d(v):
+        return v.to_4D()
+
+    # 4D momentum vector
+    p4 = vector.obj(px=1.0, py=2.0, pz=3.0, E=10.0)
+    p2 = to_2d(p4)
+    p3 = to_3d(p4)
+    assert p2.px == pytest.approx(1.0)
+    assert p2.py == pytest.approx(2.0)
+    assert p3.px == pytest.approx(1.0)
+    assert p3.py == pytest.approx(2.0)
+    assert p3.pz == pytest.approx(3.0)
+    assert pytest.approx(10.0) == to_4d(p4).E
+
+    # 2D generic vector
+    v2 = vector.obj(x=1.0, y=2.0)
+    v3 = to_3d(v2)
+    v4 = to_4d(v2)
+    assert to_2d(v2).x == pytest.approx(1.0)
+    assert v3.x == pytest.approx(1.0)
+    assert v3.y == pytest.approx(2.0)
+    assert v3.z == pytest.approx(0.0)
+    assert v4.t == pytest.approx(0.0)
+
+
+def test_momentum_named_coordinate_conversions():
+    @numba.njit
+    def to_pxpy(v):
+        return v.to_pxpy()
+
+    @numba.njit
+    def to_ptphi(v):
+        return v.to_ptphi()
+
+    @numba.njit
+    def to_pxpypz(v):
+        return v.to_pxpypz()
+
+    @numba.njit
+    def to_pxpypzenergy(v):
+        return v.to_pxpypzenergy()
+
+    @numba.njit
+    def to_ptphietamass(v):
+        return v.to_ptphietamass()
+
+    p = vector.obj(px=1.0, py=2.0, pz=3.0, E=10.0)
+    res_pxpy = to_pxpy(p)
+    assert res_pxpy.px == pytest.approx(1.0)
+    assert res_pxpy.py == pytest.approx(2.0)
+
+    res_ptphi = to_ptphi(p)
+    assert res_ptphi.pt == pytest.approx(p.pt)
+    assert res_ptphi.phi == pytest.approx(p.phi)
+
+    res_pxpypz = to_pxpypz(p)
+    assert res_pxpypz.px == pytest.approx(1.0)
+    assert res_pxpypz.py == pytest.approx(2.0)
+    assert res_pxpypz.pz == pytest.approx(3.0)
+
+    res_energy = to_pxpypzenergy(p)
+    assert pytest.approx(10.0) == res_energy.E
+
+    res_mass = to_ptphietamass(p)
+    assert res_mass.mass == pytest.approx(p.mass)
+    assert res_mass.eta == pytest.approx(p.eta)
+
+
+def test_like():
+    @numba.njit
+    def project_like(v, target):
+        return v.like(target)
+
+    v2 = vector.obj(x=1.0, y=2.0)
+    v3 = vector.obj(x=3.0, y=4.0, z=5.0)
+    v4 = vector.obj(px=1.0, py=2.0, pz=3.0, E=10.0)
+
+    # 4D projected into 2D
+    res_4_to_2 = project_like(v4, v2)
+    assert res_4_to_2.px == pytest.approx(1.0)
+    assert res_4_to_2.py == pytest.approx(2.0)
+
+    # 2D projected into 4D
+    res_2_to_4 = project_like(v2, v4)
+    assert res_2_to_4.x == pytest.approx(1.0)
+    assert res_2_to_4.y == pytest.approx(2.0)
+    assert res_2_to_4.z == pytest.approx(0.0)
+    assert res_2_to_4.t == pytest.approx(0.0)
+
+    # 3D projected into 3D
+    res_3_to_3 = project_like(v3, v3)
+    assert res_3_to_3.x == pytest.approx(3.0)
+    assert res_3_to_3.y == pytest.approx(4.0)
+    assert res_3_to_3.z == pytest.approx(5.0)
+
+    # Invalid target raises TypingError
+    @numba.njit
+    def invalid_like(v):
+        return v.like(123)
+
+    with pytest.raises(numba.TypingError):
+        invalid_like(v2)