Skip to content

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

Open
sumitsaraswat362 wants to merge 1 commit into
scikit-hep:mainfrom
sumitsaraswat362:feat-750-numba-conversions-and-like
Open

sumitsaraswat362 wants to merge 1 commit into
scikit-hep:mainfrom
sumitsaraswat362:feat-750-numba-conversions-and-like

Conversation

@sumitsaraswat362

Copy link
Copy Markdown

Fixes #750.

Description

In the Numba Object backend, several coordinate conversions and helper methods available on the Python Vector objects were missing corresponding @numba.extending.overload_method registrations, causing Numba JIT typing failures.

This PR implements the missing overload dispatchers in src/vector/backends/_numba_object.py:

  1. Numbered dimension conversions (to_2D, to_3D, to_4D):
    • Registered @numba.extending.overload_method for to_2D, to_3D, and to_4D directly across VectorObject2DType, VectorObject3DType, and VectorObject4DType, mapping to the corresponding to_Vector*D handlers.
  2. Momentum coordinate conversions:
    • In add_coordinate_change, registered the corresponding momentum method name alongside the generic method name (e.g. to_pxpy, to_ptphi, to_pxpypz, to_pxpypzenergy, to_ptphietamass, etc.), reusing the exact coordinate transformation closure.
  3. like(other) dispatch:
    • Added @numba.extending.overload_method for like on 2D, 3D, and 4D vector types with typing-time inspection on other:
      • VectorObject2DType -> dispatches to v.to_Vector2D()
      • VectorObject3DType -> dispatches to v.to_Vector3D()
      • VectorObject4DType -> dispatches to v.to_Vector4D()
      • Non-vector types cleanly raise numba.core.errors.TypingError at typing time.

Testing

  • Added comprehensive unit tests in tests/backends/test_numba_object.py covering:
    • to_2D, to_3D, to_4D dimension changes inside JIT functions.
    • Coordinate conversions to Cartesian, Polar, Cylindrical, and Momentum representations across 2D, 3D, and 4D vectors inside JIT functions.
    • like() conversions matching 2D, 3D, and 4D vector targets.
    • Rejection of non-vector targets raising numba.core.errors.TypingError.
  • All 48 tests pass locally. Linted cleanly with ruff.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Numba object backend: missing to_2D/to_3D/to_4D, momentum-named conversions, and like()

1 participant