archimedes.compile¶

archimedes.compile(func: Callable | None = None, *, static_argnums: int | Sequence[int] | None = None, static_argnames: str | Sequence[str] | None = None, return_names: str | Sequence[str] | None = None, jit: bool = False, kind: str = 'MX', name: str | None = None) → Callable¶

Create a “compiled” function from a Python function.

Transforms a standard NumPy-based Python function into an efficient computational graph that can be executed in C++ and supports automatic differentiation, code generation, and other advanced capabilities. This decorator is the primary entry point for converting Python functions into Archimedes’ symbolic computation system.

Parameters:

  • func (callable, optional) – A Python function to be evaluated symbolically.

  • static_argnums (int or sequence of int, optional) – Indices of arguments to treat as static (constant) in the function. Static arguments aren’t converted to symbolic variables but are passed directly to the function during tracing.

  • static_argnames (str or sequence of str, optional) – Names of arguments to treat as static (constant) in the function. Alternative to static_argnums when using keyword arguments.

  • return_names (str or sequence of str, optional) – Names of the return values of the function. If not specified, the return values will be named y0, y1, etc. This does not need to be specified, but is recommended for debugging and C code generation.

  • jit (bool, default=False) – Whether to compile the function with JIT for additional performance. Not fully implemented in the current version.

  • kind (str, default="MX") –

    The type of the symbolic variables. Options:

    • SX: Trace the function with scalar-valued symbolic type

    • MX: Trace the function with array-valued symbolic type

  • name (str, optional) – The name of the function. If None, taken from the function name. Required if the function is a lambda function.

Returns:

A compiled function that can be called with either symbolic or numeric arguments, while maintaining the same function signature.

Return type:

FunctionCache

Notes

When to use this function:

  • To accelerate numerical code by converting it to C++

  • To enable automatic differentiation of your functions

  • To generate C code from your functions

  • When embedding functions in optimization problems or ODE solvers

Choosing a symbolic type:

  • This option determines the type of symbolic variables used to construct the

    computational graph. The choice determines the efficiency and type of supported operations.

  • "SX" produces scalar symbolic arrays, meaning that every entry in the array

    has its own scalar symbol. This can produce highly efficient code, but is limited to a subset of operations. For example, "SX" symbolics don’t support interpolation with lookup tables.

  • "MX" symbolics are array-valued, meaning that the entire array is represented

    by a single symbol. This allows for embedding more general operations like interpolation, ODE solves, and optimization solves into the computational graph, but may not be as fast as "SX" for functions that are dominated by scalar operations.

  • The current default is "MX" and the current recommendation is to use "MX"

    symbolics unless you want to do targeted performance optimizations and feel comfortable with the symbolic array concepts.

When a compiled function is called, Archimedes:

  1. Replaces arguments with symbolic variables of the same shape and dtype

  2. Traces the execution of your function with these symbolic arguments

  3. Creates a computational graph representing all operations

  4. Caches this graph based on input shapes/dtypes and static arguments

  5. Evaluates the graph with the provided numeric inputs

The function is only traced once for each unique combination of argument shapes, dtypes, and static argument values. Subsequent calls with the same shapes reuse the cached graph, improving performance.

Static arguments:

Static arguments aren’t converted to symbolic variables. This is useful for:

  • Configuration flags that affect control flow

  • Constants that shouldn’t be differentiated through

  • Values that would be inefficient to recalculate online

Examples

Basic usage as a decorator:

>>> import numpy as np
>>> import archimedes as arc
>>>
>>> @arc.compile
... def rotate(x, theta):
...     R = np.array([
...         [np.cos(theta), -np.sin(theta)],
...         [np.sin(theta), np.cos(theta)],
...     ], like=x)
...     return R @ x
>>>
>>> x = np.array([1.0, 0.0])
>>> rotate(x, 0.5)

Using static arguments that modify the function behavior:

>>> @arc.compile(static_argnames=("use_boundary_conditions",))
... def solve_system(A, b, use_boundary_conditions=True):
...     if use_boundary_conditions:
...         b[[0, -1]] = 0.0  # Apply boundary conditions
...     return np.linalg.solve(A, b)

Different symbolic types:

>>> # Simple mathematical function - use SX for efficiency
>>> @arc.compile(kind="SX")
... def norm(x):
...     return np.sqrt(np.sum(x**2))
>>>
>>> # Function with interpolation - requires MX
>>> @arc.compile(kind="MX", static_argnames=("xp", "fp"))
... def interpolate(x, xp, fp):
...     return np.interp(x, xp, fp)

See also

grad

Compute gradients of compiled functions

jac

Compute Jacobians of compiled functions

codegen

Generate C code from compiled functions