archimedes.compile¶

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.
buffered (bool, optional) – Optimize for evaluation in a Python runtime (see notes). Default is False.

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:

Replaces arguments with symbolic variables of the same shape and dtype
Traces the execution of your function with these symbolic arguments
Creates a computational graph representing all operations
Caches this graph based on input shapes/dtypes and static arguments
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

Buffered compilation:

By default, compile supports flexible function structure, including tree-structured arguments and returns, typical Python signatures, and static arguments. The same function can also be called with different argument shapes and data types, which will cache different computational graphs for repeated execution. However, this flexibility comes with some Python overhead at runtime, and can mean that “compiled” functions perform worse than their pure-NumPy counterparts.

The buffered flag optimizes for Python runtime execution by disallowing keyword args, static args, and repeated execution with different shapes/dtypes. In buffered mode, all arguments and returns must be numeric arrays. This removes most of the Python overhead and is suitable for “production” use where a single function will be evaluated many times with the same argument types. Additionally, this mode pre-allocates memory for faster execution, although for safety the arguments and returns are still copies rather than views into the buffer.

The buffered mode should only be used for top-level function calls, not embedded within larger compiled computations. For functions that are embedded within larger compiled functions, the “runtime overhead” is not applicable and there is no advantage to buffered mode.

Note that even though you can only pass flat arrays as arguments to buffered functions, these can still be “unraveled” into structured data types as usual.

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)