Backends

Backends that write the numeric functions of compiled problems: NumPy, Numba, and JAX.

`BaseBackend`

Bases: BaseConfig[None, LoweredKernelsSpec]

Base of the backends, which turn the kernels of a problem into the numeric functions of its form.

A backend stores its settings in its constructor, prepare() takes nothing, and run(kernels) takes the LoweredKernelsSpec of a problem, with its form, the size of its state vector, its reductions, and its draws, and keeps the LoweredFunctionsSpec of the form in functions_. Subclasses implement the hooks _run, _get_name, the name a problem records, and _get_support, the times, spaces, laws, noise, and orders the backend writes functions of.

Attributes:

Name	Type	Description
`functions_`	`LoweredFunctionsSpec`	The numeric functions of the form of the kernels of the last run.

`JaxBackend`

Bases: BaseBackend

Backend that writes discrete maps and ordinary differential equations as pure JAX functions.

The functions return new arrays instead of writing into buffers, so JAX traces, compiles, and differentiates them. Reductions lower to sums, means, and matrix products, and a reduction weighted by a shaped parameter to the product with its block of the parameter vector reshaped into a matrix. The forward functions in functions_ take and return NumPy arrays and compute in double precision. Rule, noise, and jump laws, reductions over neighbourhoods, and laws given as Python callables have no JAX form.

Attributes:

Name	Type	Description
`functions_`	`LoweredFunctionsSpec`	The forward functions of the form of the kernels of the last run, on NumPy arrays.
`traced_`	`TracedFunction`	The pure JAX function `f(t, x, p)` of the form of the kernels of the last run, the drift of ordinary differential equations or the update of a discrete map.

Examples:

>>> from skcomplex.engine.compile import compile_system
>>> from skcomplex.engine.backends import JaxBackend
>>> from skcomplex.core.system import System
>>> from skcomplex.models.biology.common.populations import PopulationElement
>>> problem = compile_system(System(elements=[PopulationElement('a', n=2.0)]), backend=JaxBackend())
>>> problem.backend
'jax'

`NumbaBackend(cache_dir=None, parallel=True)`

Bases: BaseBackend

Backend that generates fused Numba kernels and compiled stepper loops, cached on disk.

Kinds of laws that contain a plain Python callable fall back to the NumPy functions, and the problem then has no compiled stepper loops.

Parameters:

Name	Type	Description	Default
`cache_dir`	`str \| Path \| None`	The directory of the generated kernel sources and Numba's cache, or `None` for `~/.cache/skcomplex`.	`None`
`parallel`	`bool`	Whether member and pair loops over groups of at least 1000 members run in parallel threads. Kernels that the compiled loops of jump processes and discrete maps call once per event always run serially.	`True`

Examples:

>>> from skcomplex.core.elements import GroupElement
>>> from skcomplex.engine.compile import compile_system
>>> from skcomplex.core.system import System
>>> from skcomplex.models.biology.common.populations import ExponentialGrowthInteraction, PopulationElement
>>> system = System(
...     elements=[GroupElement('herds', element=PopulationElement, size=2, values={'n': [1.0, 2.0]})],
...     interactions=[ExponentialGrowthInteraction('growth', target='herds', r=3.0)],
... )
>>> problem = compile_system(system, backend=NumbaBackend())
>>> problem.drift(0.0, problem.initial_state, problem.initial_parameters).tolist()
[3.0, 6.0]

Source code in src/skcomplex/engine/backends/_numba.py

def __init__(self, cache_dir: str | Path | None = None, parallel: bool = True) -> None:
    self.cache_dir = cache_dir
    self.parallel = parallel

`NumpyBackend`

Bases: BaseBackend

Reference backend that evaluates every kernel with NumPy array operations.

Examples:

>>> from skcomplex.engine.compile import compile_system
>>> from skcomplex.engine.backends import NumpyBackend
>>> from skcomplex.core.system import System
>>> from skcomplex.models.biology.common.populations import PopulationElement
>>> compile_system(System(elements=[PopulationElement('a')]), backend=NumpyBackend()).backend
'numpy'