scikit-hep / scikit-hep/vector
Remember to implement arrays of 2D/3D/Lorentz vectors in Numba
Nobody has claimed this yet.
- 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 appropriateVectorObject*DTypewhen it would otherwise return anp.record. - (Presumably, @DrTodd13's implementation would ensure that other slices maintain their class. If not, include that in the
__getitem__override. Currently, aVectorNumpyloses its class—becoming a plainnp.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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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