4.6.8. pde.tools.numba module¶
Helper functions for just-in-time compilation with numba
- convert_scalar(arr)[source]¶
helper function that turns 0d-arrays into scalars
This helps to avoid the bug discussed in https://github.com/numba/numba/issues/6000
- get_common_numba_dtype(*args)[source]¶
returns a numba numerical type in which all arrays can be represented
- Parameters
*args – All items to be tested
Returns: numba.complex128 if any entry is complex, otherwise numba.double
- jit(function: TFunc, signature=None, parallel: bool = False, **kwargs) TFunc [source]¶
apply nb.jit with predefined arguments
- Parameters
function – The function which is jitted
signature – Signature of the function to compile
parallel (bool) – Allow parallel compilation of the function
**kwargs – Additional arguments to nb.jit
- Returns
Function that will be compiled using numba
- jit_allocate_out(func: Callable, parallel: bool = False, out_shape: Optional[Tuple[int, ...]] = None, num_args: int = 1, **kwargs) Callable [source]¶
Decorator that compiles a function with allocating an output array.
This decorator compiles a function that takes the arguments arr and out. The point of this decorator is to make the out array optional by supplying an empty array of the same shape as arr if necessary. This is implemented efficiently by using
numba.generated_jit()
.- Parameters
func – The function to be compiled
parallel (bool) – Determines whether the function is jitted with parallel=True.
out_shape (tuple) – Determines the shape of the out array. If omitted, the same shape as the input array is used.
num_args (int, optional) – Determines the number of input arguments of the function.
**kwargs – Additional arguments used in
numba.jit()
- Returns
The decorated function
- make_array_constructor(arr: numpy.ndarray) Callable[[], numpy.ndarray] [source]¶
returns an array within a jitted function using basic information
- Parameters
arr (
ndarray
) – The array that should be accessible within jit
Warning
A reference to the array needs to be retained outside the numba code to prevent garbage collection from removing the array