Repository navigation
Numba object backend: missing to_2D/to_3D/to_4D, momentum-named conversions, and like() #750
Description
🤖 AI text below 🤖
The Numba backend for vector.obj (src/vector/backends/_numba_object.py) covers nearly the full VectorProtocol API, but three groups of methods that work in pure Python raise TypingError inside numba.njit:
- Dimension conversions by number:
to_2D(),to_3D(),to_4D(). Theto_Vector2D(),to_Vector3D(),to_Vector4D()spellings are overloaded and work. - Momentum-named coordinate conversions:
to_pxpy(),to_ptphi(),to_pxpypz(),to_pxpypzenergy(),to_ptphietamass(), and the rest of theto_p*family. The generic spellings (to_xy(),to_rhophi(),to_xyz(),to_xyzt(),to_rhophietatau(), ...) are overloaded viaadd_coordinate_changeand work. like(). This is the recommended replacement for the removed mixed-dimension arithmetic, so it should be available in compiled code (see also Code review: bugs, performance, simplifications, and modernizations #711, item on mixed-dimension semantics).
Reproducer (Numba 0.65.1):
import numba, vector
v = vector.obj(px=1.0, py=2.0, pz=3.0, E=10.0)
numba.njit(lambda v: v.to_Vector2D())(v) # works
numba.njit(lambda v: v.to_2D())(v) # TypingError
numba.njit(lambda v: v.to_xy())(v) # works
numba.njit(lambda v: v.to_pxpy())(v) # TypingError
numba.njit(lambda v: v.like(v.to_Vector2D()))(v) # TypingErrorItems 1 and 2 are aliases of existing overloads, so they can be registered the same way the momentum temporal aliases are (_add_momentum_temporal_alias), or by looping over _repr_momentum_to_generic. Item 3 needs a small dispatch on the dimension of the argument type at typing time.
Activity
sumitsaraswat362 commentedon Oct 22, 2026
Hi @henryiii, I looked into this and was able to reproduce the missing overloads and verify a working fix locally.
Here is the implementation structure tested in
src/vector/backends/_numba_object.py:to_2D,to_3D,to_4D): Stacked@numba.extending.overload_methodforto_2D,to_3D, andto_4Ddirectly onto the correspondingto_Vector*Doverload handlers acrossVectorObject2DType,VectorObject3DType, andVectorObject4DType.to_pxpy,to_ptphi,to_pxpypz,to_pxpypzenergy,to_ptphietamass, etc.): Inadd_coordinate_change, registered the corresponding momentum method name alongside the generic method name (@numba.extending.overload_method(vectortype, momentumname)), reusing the exact coordinate transformation closure.like(other)dispatch: Added@numba.extending.overload_methodforlikeon 2D, 3D, and 4D vector types with typing-time inspection onother:VectorObject2DType-> dispatches tov.to_Vector2D()VectorObject3DType-> dispatches tov.to_Vector3D()VectorObject4DType-> dispatches tov.to_Vector4D()numba.TypingError.I verified the reproducers in the issue description and added unit test coverage for all three groups in
tests/backends/test_numba_object.py(all 48 tests passing locally, linted cleanly withruff).If this design looks good to you, I would be happy to submit a PR!