This example requires external dependencies and cannot be run in the browser.

Standalone JIT

Compiling a hand-written NumPy function to native Rust, and getting its exact Jacobian with it.

What the JIT does

FastSim traces a Python function once, records the operations it performs, and lowers them to the scalar SSA form the engine uses internally. From then on the function runs as native Rust — the Python callable is not touched again.

Blocks that take a custom function do this automatically. fastsim.jit is the same machinery on its own, for any function MATHINLINE0ENDMATH:

jit(f) — the traced function, callable like the original.
jacobian(f) — the exact MATHINLINE1ENDMATH, differentiated from the recorded operations rather than approximated by finite differences.

Because the traced form is IR, a model built on such a function is not stuck in Python: it lowers to C and exports to an FMU like any native block.

jit() and jacobian() produce the same IR that Simulation.to_c() and Simulation.to_fmu() consume.

The Function

The Robertson kinetics problem — three species, rate constants nine orders of magnitude apart:

MATHDISPLAY0ENDMATH

with MATHINLINE1ENDMATH, MATHINLINE2ENDMATH, MATHINLINE3ENDMATH.

Python

System Parameters

Python

The Right-Hand Side

Written by hand, in the way you would write it anyway.

Python

Compiling It

jit returns a callable with the same signature. The trace happens on the first call, or eagerly if you pass n_x.

Python

The tape is specialised to the length it was traced with. A call with a different length re-traces:

Python

Jacobians

jacobian differentiates the traced operations, so one pass yields the whole matrix — no step size to pick, and no extra evaluations of MATHINLINE0ENDMATH.

Python

Verification

The Robertson Jacobian can be written down by hand:

MATHDISPLAY0ENDMATH

Python

Using It in a Model

The same function as the right-hand side of an ODE block. Robertson is stiff, so it wants an implicit solver — GEAR52A here.

Python

Results

Species 2 stays six orders of magnitude below the others, so it is scaled up to be visible at all — the usual way this solution is plotted.

Python

Mass is conserved, which is the easiest thing to check on this system: MATHINLINE0ENDMATH for all time.

Python

What Traces

Everything the tracer sees has to be an operation on the values, not a decision about them. When it cannot trace something it says so, and says what to write instead.

Python

A Python if on a traced value is the one people hit first. Rewrite it as np.where and it traces.