scikit-hep / scikit-hep/vector

Remember to implement arrays of 2D/3D/Lorentz vectors in Numba

Open
#43 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature
Dominant language
Python
Stars
99
Forks
42
Avg merge
1d 13h
Merged PRs (30d)
6

Description

PR #42 had to stop early because we need Numba to let us recognize np.ndarray subclasses inside of a Numba-compiled function. This feature is being implemented, but it is not finished or in any released version of Numba yet: numba/numba#6148.

When that is done, we can follow the procedure described in https://github.com/numba/numba/pull/6148#issuecomment-676541700 to:

  • Override __getitem__ to return the appropriate VectorObject*DType when it would otherwise return a np.record.
  • (Presumably, @DrTodd13's implementation would ensure that other slices maintain their class. If not, include that in the __getitem__ override. Currently, a VectorNumpy loses its class—becoming a plain np.ndarray—when it is range-sliced; this will have to be tested.)
  • (It might also be necessary to override an array iterator.)
  • Override __array_ufunc__ for the relevant ufuncs (np.absolute, np.add, np.equal, etc.).

That's it! Not nearly as much work as vector.backends.numba_object.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing PR #42 and the Numba feature in numba/numba#6148, then follow the procedure in its linked issue comment. Inspect VectorNumpy slicing and the relevant getitem and array_ufunc behavior, including range slices, record dtypes, and the listed ufuncs. Done means 2D, 3D, and Lorentz vector arrays work correctly inside Numba-compiled functions.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, tooling
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.