Numba

numba

Optional bounded NumPy CPU lowering with explicit, reusable packed inputs.

Importing this module does not import Numba. Compilation is eager and never falls back to Python. Generated functions have one homogeneous float signature; normal compiler evaluators and other array namespaces remain independent.

BufferLayout(name, shape, offset) dataclass

A named C-order slice in a homogeneous flat input array.

size property

Number of scalar entries, including one for scalar buffers.

CompilationSeconds(generation, optional_import, native_compilation) dataclass

Wall times for source generation, optional import and eager compilation.

NumbaRhs(model, dtype, state_names, state_layout, parameter_layout, source, compilation_seconds, kernel) dataclass

One eagerly compiled, homogeneous, NumPy CPU RHS signature.

kernel(t, y, p) accepts C-contiguous one-dimensional NumPy arrays in the layouts below. The signature also accepts read-only arrays. It checks lengths, does not mutate inputs and returns a fresh flat array. Call evaluate for strict Python-side time/dtype/layout validation, or use kernel from a Numba consumer after validating the same contract. Each call allocates its result and expression temporaries. No process-global or on-disk cache is used.

pack_state and pack_parameters explicitly copy once into owned, read-only arrays. Reuse packed parameters while values stay fixed. No input arrays or training data are retained by this artifact.

evaluate(t, y, parameters)

Validate packed inputs and execute the single compiled signature.

Returns:

Type	Description
_Array	Fresh flat NumPy derivative array in state_names order.

Raises:

Type	Description
TypeError	Inputs are foreign, differently typed or noncontiguous.

Source code in src/op_system/numba.py

def evaluate(self, t: object, y: object, parameters: object) -> _Array:
    """Validate packed inputs and execute the single compiled signature.

    Returns:
        Fresh flat NumPy derivative array in ``state_names`` order.

    Raises:
        TypeError: Inputs are foreign, differently typed or noncontiguous.
    """
    state = _input(y, (len(self.state_names),), self.dtype)
    params = _input(
        parameters, (sum(item.size for item in self.parameter_layout),), self.dtype
    )
    if not state.flags.c_contiguous or not params.flags.c_contiguous:
        msg = "Numba inputs must be C-contiguous; pack or copy explicitly"
        raise TypeError(msg)
    time = _input(t, (), self.dtype)[()]
    return self.kernel(time, state, params)

pack_parameters(values)

Copy scalar and shaped parameters into one reusable read-only array.

Returns:

Type	Description
_Array	Owned read-only NumPy parameter array; empty for parameter-free models.

Source code in src/op_system/numba.py

def pack_parameters(self, values: Mapping[str, object]) -> _Array:
    """Copy scalar and shaped parameters into one reusable read-only array.

    Returns:
        Owned read-only NumPy parameter array; empty for parameter-free models.
    """
    return _pack(self.parameter_layout, values, self.dtype)

pack_state(values)

Copy structured state buffers into the declared C-order flat layout.

Returns:

Type	Description
_Array	Owned read-only NumPy state array.

Source code in src/op_system/numba.py

def pack_state(self, values: Mapping[str, object]) -> _Array:
    """Copy structured state buffers into the declared C-order flat layout.

    Returns:
        Owned read-only NumPy state array.
    """
    return _pack(self.state_layout, values, self.dtype)

compile_numba(compiled, *, dtype='float64')

Eagerly lower a bounded scalar/fully vectorized model to NumPy CPU code.

Arithmetic, static indexing, elementwise abs/exp/log/sqrt/sin/cos/tanh, maximum/minimum/where, transposes and single-axis sums are supported. No fastmath, dynamic evaluator closure, automatic accelerator selection or fallback is used. State/parameter scalar counts are capped at one million each and generated source at 100000 characters. Runtime input values are not captured.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhs	Natural-coordinate compiler artifact with retained source IR.	required
dtype	Literal['float32', 'float64']	One homogeneous native floating precision.	'float64'

Returns:

Type	Description
NumbaRhs	Artifact with one eagerly compiled signature, source and explicit layouts.

Raises:

Type	Description
ImportError	Install op_system[numba] to enable the optional compiler.
ValueError	Precision/size bounds differ or Numba JIT is disabled.
NotImplementedError	The model needs an unsupported operation or layout.

Source code in src/op_system/numba.py

def compile_numba(  # ruff: ignore[too-many-locals]
    compiled: CompiledRhs, *, dtype: Literal["float32", "float64"] = "float64"
) -> NumbaRhs:
    """Eagerly lower a bounded scalar/fully vectorized model to NumPy CPU code.

    Arithmetic, static indexing, elementwise abs/exp/log/sqrt/sin/cos/tanh,
    maximum/minimum/where, transposes and single-axis sums are supported. No
    fastmath, dynamic evaluator closure, automatic accelerator selection or fallback
    is used. State/parameter scalar counts are capped at one million each and
    generated source at 100000 characters. Runtime input values are not captured.

    Args:
        compiled: Natural-coordinate compiler artifact with retained source IR.
        dtype: One homogeneous native floating precision.

    Returns:
        Artifact with one eagerly compiled signature, source and explicit layouts.

    Raises:
        ImportError: Install ``op_system[numba]`` to enable the optional compiler.
        ValueError: Precision/size bounds differ or Numba JIT is disabled.
        NotImplementedError: The model needs an unsupported operation or layout.
    """
    started = time.perf_counter()
    if dtype not in {"float32", "float64"}:
        msg = "Numba dtype must be float32 or float64"
        raise ValueError(msg)
    rhs = compiled._rhs  # ruff: ignore[private-member-access]
    if (
        rhs is None
        or compiled.operators
        or compiled.history_requirements
        or compiled.gauge_fix is not None
        or rhs.time_varying_params
    ):
        msg = (
            "Numba requires natural source IR without operators, history, "
            "gauge fixing or time tables"
        )
        raise NotImplementedError(msg)
    if not 0 < len(compiled.state_names) <= 1_000_000:
        msg = "Numba state count must be between 1 and 1000000"
        raise ValueError(msg)
    expressions = _Expressions()
    plan = build_vector_plan(rhs, compile_expression=expressions)
    if plan is None:
        plan = _scalar_plan(compiled, expressions)
    if plan.param_templates:
        msg = "Numba requires a supported scalar or fully vectorized template plan"
        raise NotImplementedError(msg)
    layouts = [BufferLayout(name, (), i) for i, name in enumerate(compiled.param_names)]
    offset = len(layouts)
    for name, _, shape in plan.extra_param_buffers:
        layouts.append(BufferLayout(name, shape, offset))
        offset += math.prod(shape)
    if offset > 1_000_000:
        msg = "Numba packed parameter count exceeds 1000000"
        raise ValueError(msg)
    parameters = tuple(layouts)
    source = _Source(dtype, expressions.sources).generate(plan, parameters)
    model = StructuralModelIdentity.from_compiled(compiled)
    generated = time.perf_counter()
    try:
        numba = importlib.import_module("numba")
    except ImportError as exc:
        msg = "Optional Numba lowering requires installation with op_system[numba]"
        raise ImportError(msg) from exc
    imported = time.perf_counter()
    if numba.config.DISABLE_JIT:
        msg = "Numba lowering requires JIT; NUMBA_DISABLE_JIT is enabled"
        raise ValueError(msg)
    namespace: dict[str, object] = {
        "np": np,
        "__builtins__": {"ValueError": ValueError},
    }
    exec(compile(source, "<op_system_numba>", "exec"), namespace)  # ruff: ignore[exec-builtin]
    scalar = getattr(numba, dtype)
    array = numba.types.Array(scalar, 1, "C", readonly=True)
    signature = scalar[::1](scalar, array, array)
    kernel = numba.njit(
        signature, nogil=True, fastmath=False, error_model="numpy", boundscheck=True
    )(namespace["_kernel"])
    return NumbaRhs(
        model,
        dtype,
        compiled.state_names,
        tuple(BufferLayout(t.base, t.shape, t.offset) for t in plan.state_templates),
        parameters,
        source,
        CompilationSeconds(
            generated - started, imported - generated, time.perf_counter() - imported
        ),
        cast("_NativeKernel", kernel),
    )