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

flat_idx(arr, i)[source]

helper function allowing indexing of scalars as if they arrays

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

numba_dict(data: dict = None) Optional[numba.typed.typeddict.Dict][source]

converts a python dictionary to a numba typed dictionary

numba_environment() dict[source]

return information about the numba setup used

Returns

(dict) information about the numba setup