4.1.2.6 pde.backends.numba.utils module

Defines utilities for the numba backend.

make_get_valid	Create a function to extract the valid part of a full data array.
make_get_arr_1d	Create function that extracts a 1d array at a given position.
numba_environment	Return information about the numba setup used.
jit	Apply nb.jit with predefined arguments.
make_array_constructor	Returns an array within a jitted function using basic information.
numba_dict	Converts a python dictionary to a numba typed dictionary.
get_common_numba_dtype	Returns a numba numerical type in which all arrays can be represented.
random_seed	Sets the seed of the random number generator of numpy and numba.

class Counter(value=0)

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()

flat_idx(arr, i)

Helper function allowing indexing of scalars as if they arrays.

Parameters:

• arr (ndarray or scalar) – Array or scalar value to index

• i (int) – Index to access

Return type:

Number

get_common_numba_dtype(*args)

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, *, backend=None, **kwargs)

Apply nb.jit with predefined arguments.

Parameters:

• function (TFunc) – The function which will be jitted

• signature – Signature of the function to compile

• parallel (bool) – Allow parallel compilation of the function

• backend (NumbaBackend | None) – The backend from which compilation options will be read. If None, the default numba backend will be used.

• **kwargs – Additional arguments to numba.jit()

Returns:

Function that will be compiled using numba

Return type:

TFunc

make_array_constructor(arr)

Returns an array within a jitted function using basic information.

Parameters:

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

Return type:

Callable[[], NumericArray]

Warning

A reference to the array needs to be retained outside the numba code to prevent garbage collection from removing the array

make_get_arr_1d(dim, axis)

Create function that extracts a 1d array at a given position.

Parameters:

• dim (int) – The dimension of the space, i.e., the number of axes in the supplied array

• axis (int) – The axis that is returned as the 1d array

Returns:

A numba compiled function that takes the full array arr and an index idx (a tuple of dim integers) specifying the point where the 1d array is extract. The function returns a tuple (arr_1d, i, bc_i), where arr_1d is the 1d array, i is the index i into this array marking the current point and bc_i are the remaining components of idx, which locate the point in the orthogonal directions. Consequently, i = idx[axis] and arr[…, idx] == arr_1d[…, i].

Return type:

function

make_get_valid(grid)

Create a function to extract the valid part of a full data array.

Parameters:

grid (GridBase) – The grid for which the function is created

Returns:

Mapping a numpy array containing the full data of the grid to a

numpy array of only the valid data

Return type:

callable

numba_dict(data=None, /, **kwargs)

Converts a python dictionary to a numba typed dictionary.

Parameters:

• data (dict, optional) – Data to be converted to a dictionary. If None, an empty dictionary is created.

• **kwargs – Additional items added to the dictionary

Returns:

A dictionary of numba type

Return type:

Dict

numba_environment(backend=None)

Return information about the numba setup used.

Parameters:

backend (NumbaBackend | None) – The backend from which compilation options will be read. If None, the default numba backend will be used.

Returns:

(dict) information about the numba setup

Return type:

dict[str, Any]

ol_flat_idx(arr, i)

Helper function allowing indexing of scalars as if they arrays.

random_seed(seed=0)

Sets the seed of the random number generator of numpy and numba.

Parameters:

seed (int) – Sets random seed

Return type:

None