Making Numba Modules

Source Code

The Python code shown below can be downloaded here.

NumbaModel is a Basilisk module type that achieves near-C execution speed while being written entirely in Python. It is a good fit for numerically intensive inner loops (attitude control laws, filter update steps, trajectory integrators) that are too slow as plain Python modules.

Under the hood, NumbaModel uses Numba to JIT-compile the user-supplied UpdateStateImpl static method into a cfunc (a C-callable function pointer) that the Basilisk C++ scheduler calls directly every tick. The compilation happens once, at Reset time.

Note

Because UpdateStateImpl is compiled by Numba in nopython mode, a subset of Python and NumPy is available inside it. Read the What is Allowed in UpdateStateImpl section carefully before writing complex logic.

Defining a Numba Module

Subclass NumbaModel from Basilisk.architecture.numbaModel and implement UpdateStateImpl as a @staticmethod. All message attributes and memory fields must be created before Reset; the parameter names of UpdateStateImpl tell the framework which of those attributes to wire up.

import numpy as np
from Basilisk.architecture.numbaModel import NumbaModel
from Basilisk.architecture import messaging, bskLogging
from Basilisk.utilities import SimulationBaseClass, macros


def run():
    """
    Illustration of NumbaModel: two JIT-compiled modules exchanging messages.
    """

    scSim = SimulationBaseClass.SimBaseClass()
    proc  = scSim.CreateNewProcess("proc")
    proc.addTask(scSim.CreateNewTask("task", macros.sec2nano(1.0)))

    # AccumulatorModel tracks elapsed time in memory; it has no input messages.
    prod = AccumulatorModel()
    prod.ModelTag = "accumulator"
    scSim.AddModelToTask("task", prod)

    # FilterModel reads from AccumulatorModel and scales by a configurable gain.
    filt = FilterModel()
    filt.ModelTag = "filter"
    filt.memory.gain = 2.0                  # override default before Reset
    filt.dataInMsg.subscribeTo(prod.dataOutMsg)
    scSim.AddModelToTask("task", filt)

    recorder = filt.dataOutMsg.recorder()
    scSim.AddModelToTask("task", recorder)

    scSim.InitializeSimulation()
    scSim.ConfigureStopTime(macros.sec2nano(3.0))  # 4 ticks: t = 0, 1, 2, 3 s
    scSim.ExecuteSimulation()

    print("ticks recorded by accumulator:", int(prod.memory.ticks))
    print("filtered dataVector at t=3s:  ", recorder.dataVector[-1])


# =============================================================================
# Module 1: AccumulatorModel
#
# Demonstrates: memory namespace, CurrentSimNanos, no readers.
# =============================================================================
class AccumulatorModel(NumbaModel):
    """Outputs the cumulative simulation time (seconds) and a tick counter."""

    def __init__(self):
        """Declare the output message and persistent accumulator state."""
        super().__init__()
        self.dataOutMsg = messaging.CModuleTemplateMsg()

        # Fields declared here become a C-level structured buffer after Reset.
        # The dtype is inferred from the value: 0.0 -> float64, 0 -> int64,
        # np.int32(0) -> int32, np.array(...) -> array field.
        self.memory.total_s = 0.0
        self.memory.ticks   = 0

    @staticmethod
    def UpdateStateImpl(dataOutMsgPayload, CurrentSimNanos, memory):
        """Accumulate elapsed time and publish the running totals."""
        # Everything in this method is compiled by Numba in nopython mode.
        # See the "What is allowed" section in the documentation.
        memory.total_s += np.float64(CurrentSimNanos) * 1e-9
        memory.ticks   += 1
        dataOutMsgPayload.dataVector[0] = memory.total_s
        dataOutMsgPayload.dataVector[1] = float(memory.ticks)


# =============================================================================
# Module 2: FilterModel
#
# Demonstrates: scalar reader with IsLinked guard, bskLogger, memory gain.
# =============================================================================
class FilterModel(NumbaModel):
    """Scales each element of the input dataVector by a configurable gain."""

    def __init__(self):
        """Set up one reader, one writer, and the configurable gain term."""
        super().__init__()
        self.dataInMsg  = messaging.CModuleTemplateMsgReader()
        self.dataOutMsg = messaging.CModuleTemplateMsg()
        self.memory.gain = 1.0               # default; caller may override

    @staticmethod
    def UpdateStateImpl(dataInMsgPayload, dataInMsgIsLinked,
                        dataOutMsgPayload, bskLogger, memory):
        """Scale the input vector when linked and warn otherwise."""
        if not dataInMsgIsLinked:
            bskLogger.bskLog(bskLogging.BSK_WARNING, "FilterModel: input not linked")
            return
        for i in range(3):
            dataOutMsgPayload.dataVector[i] = dataInMsgPayload.dataVector[i] * memory.gain


if __name__ == "__main__":
    run()

Parameter Naming Convention

Unlike regular Python or C++ modules, NumbaModel does not require you to write a Reset or UpdateState method. Instead, the framework introspects the parameter names of UpdateStateImpl at Reset time and automatically wires each one to the correct buffer.

You may override Reset (e.g. to perform additional Python-level setup), but you must call super().Reset(CurrentSimNanos) so that the cfunc is compiled and registered:

def Reset(self, CurrentSimNanos=0):
    # ... additional setup ...
    super().Reset(CurrentSimNanos)   # must be called

Do not override UpdateState. The C++ scheduler calls the JIT-compiled cfunc registered by NumbaModel.Reset directly.

The recognised UpdateStateImpl parameter names are:

Parameter name pattern	Type inside UpdateStateImpl	What it provides
<name>OutMsgPayload	Numba Record (struct-like)	Writable payload of self.<name>OutMsg
<name>InMsgPayload	Numba Record or array of Records	Read-only snapshot of self.<name>InMsg
<name>InMsgIsLinked	bool (scalar reader), uint8[:] (list), or Record (dict)	True/non-zero when the corresponding reader is subscribed
CurrentSimNanos	uint64	Current simulation time in nanoseconds
moduleID	int64	This module’s unique integer ID
memory	Numba Record	Persistent state fields (see Persistent Memory)
bskLogger	BskLoggerProxy	Logging proxy (see Logging with bskLogger)
rng	xoshiro256++ PRNG	Per-module PRNG seeded from self.RNGSeed at Reset (see Per-module Random Number Generator (rng))

Any parameter name not matching one of the patterns above raises an AttributeError at Reset time with a descriptive message.

Note

Parameter order inside UpdateStateImpl does not matter; the framework matches by name, not position.

Output Messages

Declare output messages as messaging.<Type>Msg() attributes in __init__:

self.dataOutMsg = messaging.CModuleTemplateMsg()

Use the suffix OutMsgPayload in UpdateStateImpl to receive a writable view of the payload:

@staticmethod
def UpdateStateImpl(dataOutMsgPayload, ...):
    dataOutMsgPayload.dataVector[0] = 42.0

The Basilisk header timestamp is written automatically after UpdateStateImpl returns. You do not need to call .write().

For multiple output messages, assign a list or dict to the attribute:

self.dataOutMsg = [messaging.CModuleTemplateMsg(),
                   messaging.CModuleTemplateMsg()]

Inside UpdateStateImpl, access them by integer index:

dataOutMsgPayload[0].dataVector[0] = v
dataOutMsgPayload[1].dataVector[0] = v * 2.0

Or by string key when using a dict:

self.dataOutMsg = {'pos': messaging.CModuleTemplateMsg(),
                   'neg': messaging.CModuleTemplateMsg()}

# inside UpdateStateImpl:
dataOutMsgPayload['pos'].dataVector[0] =  v
dataOutMsgPayload['neg'].dataVector[0] = -v

Input Messages (Readers)

Declare readers as messaging.<Type>MsgReader() attributes:

self.dataInMsg = messaging.CModuleTemplateMsgReader()

Subscribe before InitializeSimulation:

mod.dataInMsg.subscribeTo(src.dataOutMsg)

The payload is refreshed from the source every tick by the C++ scheduler. Access it in UpdateStateImpl via the InMsgPayload parameter.

If a reader may be unlinked, declare the corresponding IsLinked parameter and guard reads behind it:

@staticmethod
def UpdateStateImpl(dataInMsgPayload, dataInMsgIsLinked, dataOutMsgPayload):
    if dataInMsgIsLinked:
        dataOutMsgPayload.dataVector[0] = dataInMsgPayload.dataVector[0]

If <name>InMsgPayload is declared without a corresponding <name>InMsgIsLinked, Reset raises RuntimeError when the reader is not subscribed. This is an intentional early failure rather than a silent wrong-data scenario.

List and dict attributes work the same way as for output messages. For a list reader, IsLinked is a uint8[:] array indexed by position (dataInMsgIsLinked[0], dataInMsgIsLinked[1], …). For a dict reader, IsLinked is a Record keyed by the same names as the payload, so both can be accessed identically:

if dataInMsgIsLinked['a']:
    val = dataInMsgPayload['a'].dataVector[0]

Persistent Memory

self.memory is a namespace that holds persistent module state across ticks. Fields are defined as attributes in __init__:

self.memory.step   = 0                 # scalar integer (→ int64)
self.memory.gain   = 1.5              # scalar float   (→ float64)
self.memory.bias   = [0.1, 0.2, 0.3]  # 1-D array (converted to numpy)
self.memory.matrix = np.eye(3)         # 2-D array

At Reset time the fields are packed into a single structured numpy array whose pointer is handed to the C++ scheduler. After Reset, the fields are accessible from Python via the same self.memory.<name> syntax and any assignment is forwarded directly to the underlying C buffer - no re-compilation is needed.

Rules:

• All fields must be defined before Reset is called (i.e., in __init__ or before InitializeSimulation).

• New fields cannot be added after Reset.

• The numpy dtype of each field is inferred from the value assigned in __init__: a plain 0 becomes int64, a plain 0.0 becomes float64. Use an explicit numpy type (e.g. np.int32(0)) only when a specific width is required.

• Inside UpdateStateImpl, memory behaves like a Numba Record. Array fields are accessed with standard bracket indexing: memory.bias[1], memory.matrix[i, j].

Logging with bskLogger

Add bskLogger as a parameter to receive a Numba-compatible logging proxy:

@staticmethod
def UpdateStateImpl(dataOutMsgPayload, bskLogger, memory):
    bskLogger.bskLog(bskLogging.BSK_INFORMATION, "tick")
    bskLogger.bskLog1(bskLogging.BSK_WARNING, "step:", memory.step)

The proxy respects the log level set on self.bskLogger at Reset time: messages below the current threshold are suppressed. Output goes to stdout via print; ANSI colour codes match the rest of Basilisk’s log output.

Available methods:

Method	Prints
bskLog(level, msg)	[TAG] msg
bskLog1(level, msg, v0)	[TAG] msg v0
bskLog2(level, msg, v0, v1)	[TAG] msg v0 v1
bskLog3(level, msg, v0, v1, v2)	[TAG] msg v0 v1 v2

Level constants (bskLogging.BSK_DEBUG, BSK_INFORMATION, BSK_WARNING, BSK_ERROR) work inside UpdateStateImpl because Numba resolves module-level integer attributes at compile time.

msg must be a string literal - variables of string type are not supported in nopython mode. Values v0–v2 can be any numeric type that Numba supports (int32, float64, etc.).

Tip

Plain print("label:", value) also works in UpdateStateImpl with no additional setup.

Per-module Random Number Generator (rng)

Add rng as a parameter to receive a per-module xoshiro256++ PRNG that is private to this module instance and seeded from self.RNGSeed at Reset time:

@staticmethod
def UpdateStateImpl(dataOutMsgPayload, rng, memory):
    noise = rng.standard_normal(3)     # independent per-module stream
    dataOutMsgPayload.dataVector[:3] = memory.signal + noise

The generator state persists across ticks, so each call advances the sequence deterministically. self.RNGSeed is inherited from SysModel and defaults to 0x1badcad1; set it before InitializeSimulation() for reproducible Monte Carlo runs:

sNav.RNGSeed = 12345

Key properties:

• Per-module isolation: each instance owns an independent generator — two modules with the same seed produce the same sequence independently, and two modules with different seeds never interfere.

• Re-seeded on Reset: calling Reset again seeds a fresh generator from the current self.RNGSeed, resetting the sequence.

Available methods:

Method	Description
standard_normal(n)	Return n independent standard-normal samples as float64[n]
random()	Uniform float64 in [0, 1)
integers(low, high)	Uniform int64 in [low, high)
uniform(low, high)	Uniform float64 in [low, high)

What is Allowed in UpdateStateImpl

UpdateStateImpl is compiled by Numba in nopython mode. The following is a summary of what is and is not supported.

Allowed

• NumPy array operations on pre-allocated arrays (element-wise arithmetic, [i] / [i, j] indexing, array.copy(), np.zeros(n), np.dot, reductions, etc.). See the Numba NumPy support list for the full set.

• Explicit numpy scalars such as np.int32(1), np.float64(3.14). Use these when a specific width is required (e.g. to match a memory field declared as np.int32); otherwise plain literals are fine.

• Control flow: if/elif/else, for over range(n), while, break, continue, return.

• Math functions: math.sqrt, math.sin, etc. (import math at module level, use inside the function normally).

• Calling other @nb.njit functions - factor complex logic into helper functions decorated with @numba.njit.

• ``print(label, value, …)`` - variadic print to stdout works and is the simplest way to emit debug output.

• ``bskLogger`` methods as described in Logging with bskLogger.

• Reading from memory Record fields and writing to them.

• Reading from InMsgPayload Record fields and writing to OutMsgPayload Record fields.

Not Allowed

• Python objects of any kind: no lists, dicts, sets, tuples, classes, generators, or iterators. Use numpy arrays for all collections.

• String variables - string literals in print and bskLogger calls work, but you cannot assign a string to a variable or pass one as a function argument (other than via the bskLog* API).

• f-strings and %-formatting - these require Python object creation. Use bskLog1/bskLog2 to attach numeric values to log messages.

• Exceptions - raise, try/except/finally are not supported. Use conditional logic and bskLogger to handle error conditions.

• Calling Python functions that are not themselves @nb.njit-compiled. Standard library functions (os, sys, etc.) are unavailable.

• Dynamic memory allocation with arbitrary Python types. NumPy array creation with a constant size (np.zeros(3)) is fine; allocation whose size depends on a runtime variable may require care.

• Global mutable state - do not capture mutable Python objects in the closure of UpdateStateImpl. Use memory for persistent state.

• Messaging API - do not call .write(), .read(), or .subscribeTo() inside UpdateStateImpl. Payload structs are passed as already-refreshed Records; writing back through them is sufficient.

Re-use and Re-initialisation

Re-wiring messages mid-simulation (calling subscribeTo without a Reset) is fully supported. The C++ scheduler refreshes the payload pointer dynamically every tick via registered reader slots.

Calling ``Reset`` a second time re-compiles the cfunc and re-initialises memory from the current buffer state. The buffer is always the source of truth: Python-side writes (self.memory.gain = 2.0) and cfunc-internal writes (memory.step += 1 inside UpdateStateImpl) are treated identically; both persist through a subsequent Reset. The next run always continues from whatever value was last written, by any means.

RigidBodyKinematicsNumba

Basilisk.utilities.RigidBodyKinematicsNumba provides @nb.njit-compiled versions of every function in Basilisk.utilities.RigidBodyKinematics, making the full rigid-body kinematics library callable from within compiled UpdateStateImpl kernels.

Import individual functions at module level (outside UpdateStateImpl):

from Basilisk.utilities.RigidBodyKinematicsNumba import MRP2C, addMRP, subMRP

class MyModule(NumbaModel):
    @staticmethod
    def UpdateStateImpl(attNavInMsgPayload, attGuidOutMsgPayload, memory):
        C_BN = MRP2C(attNavInMsgPayload.sigma_BN)
        sigma_BR = subMRP(attNavInMsgPayload.sigma_BN, memory.sigma_RN)
        attGuidOutMsgPayload.sigma_BR[:3] = sigma_BR

All kinematics functions (MRP2C, C2MRP, addMRP, subMRP, EP2C, C2EP, PRV2C, and more) are available. Because they are decorated with @nb.njit(inline='always'), calls inside UpdateStateImpl compile inline with no per-call overhead.

Note

RigidBodyKinematicsNumba requires scipy to be installed for BLAS-backed np.linalg operations inside Numba nopython mode.

Continuous-Time States (StatefulNumbaModel)

StatefulNumbaModel (from Basilisk.simulation.StatefulNumbaModel) combines NumbaModel with StatefulSysModel to add continuous-time states that are integrated by a MJScene dynamics scene. Use it when a Numba module needs to register differentiable states (position, velocity, etc.) for integration alongside MuJoCo dynamics.

Defining a StatefulNumbaModel

Subclass StatefulNumbaModel, implement registerStates to create StateData objects, and write the derivatives in UpdateStateImpl:

from Basilisk.simulation.StatefulNumbaModel import StatefulNumbaModel

class HarmonicOscillator(StatefulNumbaModel):

    def registerStates(self, registerer):
        self.posState = registerer.registerState(3, 1, "pos")
        self.velState = registerer.registerState(3, 1, "vel")

    def Reset(self, CurrentSimNanos=0):
        self.memory.omega = 2.0
        super().Reset(CurrentSimNanos)   # compiles the cfunc

    @staticmethod
    def UpdateStateImpl(posState, posStateDeriv,
                        velState, velStateDeriv, memory):
        posStateDeriv[:, 0] = velState[:, 0]
        velStateDeriv[:, 0] = -memory.omega**2 * posState[:, 0]

registerStates is called by MJScene before Reset(). Always call super().Reset(CurrentSimNanos) in your Reset override so that the cfunc is compiled after the state objects exist.

Additional UpdateStateImpl parameters

All NumbaModel parameters are supported, plus:

Parameter name pattern	Type inside UpdateStateImpl	What it provides
<name>State	float64[:, :] (Fortran-order)	Read-only current state value; resolved from self.<name>State
<name>StateDeriv	float64[:, :] (Fortran-order)	Writable state derivative; resolved from the same StateData as <name>State
<name>StateDiffusionN	float64[:, :] (Fortran-order)	Writable diffusion matrix for the N-th noise source (N ≥ 0); requires self.<name>State.setNumNoiseSources(N+1) in registerStates()

States are registered as nRow × nCol matrices with a column-major (Fortran-order) backing store, so arr[i, j] equals matrix element M(i, j). Even 1-D states are exposed as 2-D arrays — a 3×1 column vector is accessed as posState[i, 0].

Performance Notes

• Compilation happens once at Reset time and is cached to disk by Numba (keyed on the bytecode of UpdateStateImpl). Subsequent runs with the same code load the cached binary instantly.

• The cfunc is called from C++ with zero Python overhead per tick. Execution speed is comparable to a hand-written C extension for the same algorithm.

• Keep UpdateStateImpl free of Python-level function calls; every call that Numba cannot inline adds compilation complexity and may defeat caching.