4.6.9 pde.tools.numba module

Helper functions for just-in-time compilation with numba

class Counter(value=0)[source]

Bases: object

helper class for implementing JIT_COUNT

We cannot use a simple integer for this, since integers are immutable, so if one imports JIT_COUNT from this module it would always stay at the fixed value it had when it was first imported. The workaround would be to import the symbol every time the counter is read, but this is error-prone. Instead, we implement a thin wrapper class around an int, which only supports reading and incrementing the value. Since this object is now mutable it can be used easily. A disadvantage is that the object needs to be converted to int before it can be used in most expressions.

Parameters:

value (int) –

increment()[source]
flat_idx(arr, i)[source]

helper function allowing indexing of scalars as if they arrays

Parameters:
Return type:

int | float | complex

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, signature=None, parallel=False, **kwargs)[source]

apply nb.jit with predefined arguments

Parameters:
  • function (TFunc) – 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

Return type:

TFunc

make_array_constructor(arr)[source]

returns an array within a jitted function using basic information

Parameters:

arr (ndarray) – The array that should be accessible within jit

Return type:

Callable[[], ndarray]

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=None)[source]

converts a python dictionary to a numba typed dictionary

Parameters:

data (dict[str, Any] | None) –

Return type:

Dict | None

numba_environment()[source]

return information about the numba setup used

Returns:

(dict) information about the numba setup

Return type:

dict[str, Any]

ol_flat_idx(arr, i)[source]

helper function allowing indexing of scalars as if they arrays

random_seed(seed=0)[source]

sets the seed of the random number generator of numpy and numba

Parameters:

seed (int) – Sets random seed

Return type:

None