Allow Basix functions to be called from Numba kernels

[garth-wells]

Tabulation and the push-forward/pull-back maps are exposed as native functions that a Numba kernel compiled in nopython mode can call directly, so that loops over large numbers of points run at native speed instead of paying a Python call per point.

Until now basix.numba_helpers only provided DOF transformations, reimplemented in Python. Anything else had to be reimplemented the same way, which is what #201 asks to avoid.

Interface

from basix.numba_helpers import ElementHandle, tabulate_fn

element = basix.create_element(basix.ElementFamily.P, basix.CellType.triangle, 2)
handle = ElementHandle(element)
tabulate = tabulate_fn(element.dtype)

@numba.njit(cache=True, nogil=True)
def kernel(address, X, basis):
    if tabulate(address, 0, X.ctypes.data, X.shape[0], basis.ctypes.data) != 0:
        raise RuntimeError("Basix tabulation failed.")

• The entry points use no Python API and let no C++ exception escape into JIT-compiled code, reporting failure through their return value instead. Kernels calling them can release the GIL.
• They are registered with LLVM by symbol name and wrapped as numba.types.ExternalFunction, so kernels that call them are still cacheable. Binding the raw address through ctypes instead makes Numba refuse to cache the kernel (Cannot cache compiled function ... uses dynamic globals).
• ElementHandle holds the address of the C++ element together with a reference to the element, so the element cannot be collected while a kernel is using it.
• The map entry point writes into a buffer supplied by the caller and performs no allocation, unlike push_forward/pull_back. It derives the reference and physical value sizes from the element and the geometric dimension, which differ when gdim != tdim — a contravariant Piola mapped element on a triangle in 3D has reference value size 2 and physical value size 3.

FiniteElement.tabulate_shape is added to the Python interface, since the size of the tabulation buffer cannot otherwise be computed from Python.

Demo

demo/python/demo_numba_kernels.py assembles the interpolation matrix from a degree 2 Lagrange space to a set of scattered points, which is the use case raised in #201. One kernel pulls each point back to the reference cell, tabulates, and scatters the basis values into a matrix row. On 200,000 points:

Numba kernel: 0.13 us/point
Python loop:  8.71 us/point
Speedup:      67.8x

Its demo test is guarded with importorskip, since Numba is not installed on every platform Basix supports.

Testing

43 cases in test/test_numba.py: tabulation against FiniteElement.tabulate for float32 and float64, derivative orders 0 to 2 and five elements; maps against push_forward/pull_back including three cases where gdim != tdim; a cache test that promotes NumbaWarning to an error; handle lifetime and dtype rejection; and an error path using a custom L2 Piola element, where an unsupported map is reported rather than letting the C++ exception escape.

Closes #201.

🤖 Generated with Claude Code

[jhale]

not by the mathematics - not by the calculations performed within the call.

[jhale]

This is a bit roundabout and I'm not sure about the GIL because this doesn't appear to be threaded?

I would propose "the numba code directly calls the Basix C++ function without passing back through the Python interpreter, leading to no Python function call overhead."

[jhale]

This comment is 'standard knowledge'.

[jorgensd]

Not just a bounding box tree, it would also have to check for containment within the cell by using a pull back or closest point projection.

[jorgensd]

Shouldn't this rather use the basis functions to compute the jacobian, so this generalizes to higher order or non-affien geometries?