Coverage for src/evutils/jit.py: 86%

1"""Lazy numba compilation helper. 

2 

3Importing numba costs several hundred milliseconds, which would be paid by 

4``import evutils.io`` even when no numba-accelerated code path is ever used 

5(e.g. pure reading through the C parsers). :func:`lazy_njit` defers both the 

6numba import and the JIT compilation to the first call of the decorated 

7function. 

8""" 

9from __future__ import annotations 

10 

11import functools 

12from typing import Callable, TypeVar, Any 

13 

14F = TypeVar("F", bound=Callable[..., Any]) 

15 

16def lazy_njit(fn: F) -> F: 

17 """``numba.njit``, but imported and compiled on first call. 

18 

19 The wrapped function behaves like the ``@nb.njit``-decorated original; 

20 only the timing of the numba import/compilation differs. The per-call 

21 overhead after the first call is a single ``is None`` check. 

22 """ 

23 compiled: Callable[..., Any] | None = None 

24 

25 @functools.wraps(fn) 

26 def wrapper(*args: Any, **kwargs) -> Any: 

27 nonlocal compiled 

28 if compiled is None: 

29 import numba as nb 

30 compiled = nb.njit(fn) 

31 return compiled(*args, **kwargs) 

32 

33 return wrapper # type: ignore[return-value] 

34 

35def lazy_njit_unwrapped_events(fn: F) -> F: 

36 """Decorator that unwraps SoA or AoS events into constituent arrays,  

37 then calls a lazily compiled numba function. Numba handles the  

38 specialization for strided vs contiguous memory. 

39 """ 

40 compiled: Callable[..., Any] | None = None 

41 

42 @functools.wraps(fn) 

43 def wrapper(events: Any, *args: Any, **kwargs) -> Any: 

44 nonlocal compiled 

45 if compiled is None: 

46 import numba as nb 

47 compiled = nb.njit(fn) 

48 

49 # We import SoaArray here to avoid any top-level circular imports 

50 from .types import SoaArray 

51 import numpy as np 

52 

53 if isinstance(events, SoaArray): 

54 arrays = tuple(getattr(events, f) for f in ("t", "x", "y", "p") if hasattr(events, f)) 

55 # Just ensure we have exactly 4 fields for these dense kernels 

56 if len(arrays) != 4: 

57 # If it's a trigger array or something else, fall back to ordered fields 

58 arrays = tuple(getattr(events, f) for f in events._fields) 

59 elif isinstance(events, np.ndarray) and events.dtype.names is not None: 

60 if set(events.dtype.names).issuperset({"t", "x", "y", "p"}): 

61 arrays = tuple(events[f] for f in ("t", "x", "y", "p")) 

62 else: 

63 # Not a (t, x, y, p) event array: pass fields through in dtype 

64 # order, but warn -- these kernels expect (t, x, y, p). 

65 import warnings 

66 warnings.warn( 

67 f"array with fields {events.dtype.names} does not provide " 

68 "(t, x, y, p); passing fields positionally to the kernel", 

69 stacklevel=2, 

70 ) 

71 arrays = tuple(events[f] for f in events.dtype.names) 

72 else: 

73 raise TypeError(f"Unsupported event format: {type(events)}") 

74 

75 return compiled(*arrays, *args, **kwargs) 

76 

77 return wrapper # type: ignore[return-value]