# ruff: noqa: N802, N803, N806, N815, N816
import os
import pprint
import numpy as np
from scipy import signal
# Simple utilities for displaying generated code in the notebook
from utils import cleanup, display_text
import archimedes as arc
from archimedes import struct
THEME = os.environ.get("ARCHIMEDES_THEME", "dark")
arc.theme.set_theme(THEME)
cleanup() # Clean up any previous generated code
Structured data types¶
So far we have been working with a function that operates with regular arrays:
# Design a simple IIR filter with SciPy
dt = 0.01 # Sampling time [seconds]
Wn = 10 # Cutoff frequency [Hz]
order = 4
b, a = signal.butter(order, Wn, "low", analog=False, fs=1 / dt)
@arc.compile(return_names=("u_hist", "y_hist"))
def iir_filter_flat(u, b, a, u_prev, y_prev):
# Update input history
u_prev[1:] = u_prev[:-1]
u_prev[0] = u
# Compute output using the direct II transposed structure
y = (np.dot(b, u_prev) - np.dot(a[1:], y_prev[: len(a) - 1])) / a[0]
# Update output history
y_prev[1:] = y_prev[:-1]
y_prev[0] = y
return u_prev, y_prev
However, this approach of passing each individual input, output, or state component as an arg doesn’t scale well to more complex algorithms. Additionally, for stateful functions it requires manual state management:
// Copy output arrays back to inputs
for (int j = 0; j < n; j++) {
arg.u_prev[j] = res.u_hist[j];
arg.y_prev[j] = res.y_hist[j];
}
arg.u_prev[n] = res.u_hist[n];
This single filter uses five arguments; even if we just added two extra IIR filters, we might need as many as 15 arguments to keep track of the filter histories, coefficients, etc. This quickly becomes difficult to understand and maintain.
One solution is to use hierarchical data structures to organize the arguments and returns.
We call these “PyTrees”, using terminology borrowed from JAX.
These let you organize state and parameters in dataclass-like containers which can be mapped directly to a C struct.
They can even be nested inside of one another to arbitrary depth.
See the overview of PyTrees or the documentation for the @pytree_node decorator for more details.
Multiple filters¶
For example, let’s say we have an algorithm that is composed of three discrete transfer functions, implemented as IIR filters:
|--> G --> y_g
u --> F --> y_f --|
|--> H --> y_h
Instead of maintaining six arrays (one each for u and y histories, across three filters), plus inputs and outputs, we can organize these into a logical hierarchical structure.
First, let’s rewrite the basic IIR filter to work with a structured state:
@struct.pytree_node
class FilterState:
u_prev: np.ndarray
y_prev: np.ndarray
def iir_filter(
x: FilterState, u: float, b: np.ndarray, a: np.ndarray
) -> tuple[FilterState, float]:
u_prev, y_prev = x.u_prev, x.y_prev
# Update input history
u_prev[1:] = u_prev[:-1]
u_prev[0] = u
# Compute output using the direct II transposed structure
y = (np.dot(b, u_prev) - np.dot(a[1:], y_prev[: len(a) - 1])) / a[0]
# Update output history
y_prev[1:] = y_prev[:-1]
y_prev[0] = y
return FilterState(u_prev, y_prev), y
Next we can combine three such filters into our “compound” filter with transfer functions F, G, and H:
@struct.pytree_node
class CompoundState:
x_f: FilterState
x_g: FilterState
x_h: FilterState
@struct.pytree_node
class CompoundOutput:
y_f: float
y_g: float
y_h: float
@arc.compile(return_names=("state_new", "y"))
def compound_filter(
state: CompoundState, u: float
) -> tuple[CompoundState, CompoundOutput]:
# For simplicity, just re-use the coefficients from the low-pass filter
# we've already designed
x_f, y_f = iir_filter(state.x_f, u, b, a)
x_g, y_g = iir_filter(state.x_g, y_f, b, a)
x_h, y_h = iir_filter(state.x_h, y_f, b, a)
return CompoundState(x_f, x_g, x_h), CompoundOutput(y_f, y_g, y_h)
x0 = CompoundState(
x_f=FilterState(u_prev=np.zeros(len(b)), y_prev=np.zeros(len(a) - 1)),
x_g=FilterState(u_prev=np.zeros(len(b)), y_prev=np.zeros(len(a) - 1)),
x_h=FilterState(u_prev=np.zeros(len(b)), y_prev=np.zeros(len(a) - 1)),
)
u0 = 1.0
args = (x0, u0)
x, y = compound_filter(*args)
pprint.pprint(y)
CompoundOutput(y_f=array(0.00482434),
y_g=array(2.32742888e-05),
y_h=array(2.32742888e-05))
If you’ve read the documentation on PyTrees and/or worked with them a bit, you can probably already see where this is going - a pytree_node behaves a lot like a C struct!
This makes it straightforward to generate code that maps directly back to your Python data structures.
arc.codegen(compound_filter, args)
with open("compound_filter.h", "r") as f:
c_code = f.read()
display_text(c_code)
#ifndef COMPOUND_FILTER_H
#define COMPOUND_FILTER_H
#include "compound_filter_kernel.h"
#ifdef __cplusplus
extern "C" {
#endif
typedef struct {
float u_prev[5];
float y_prev[4];
} filter_state_t;
typedef struct {
filter_state_t x_f;
filter_state_t x_g;
filter_state_t x_h;
} compound_state_t;
typedef struct {
float y_f;
float y_g;
float y_h;
} compound_output_t;
// Input arguments struct
typedef struct {
compound_state_t state;
float u;
} compound_filter_arg_t;
// Output results struct
typedef struct {
compound_state_t state_new;
compound_output_t y;
} compound_filter_res_t;
// Workspace struct
typedef struct {
long int iw[compound_filter_SZ_IW];
float w[compound_filter_SZ_W];
} compound_filter_work_t;
// Runtime API
int compound_filter_init(compound_filter_arg_t* arg, compound_filter_res_t* res, compound_filter_work_t* work);
int compound_filter_step(compound_filter_arg_t* arg, compound_filter_res_t* res, compound_filter_work_t* work);
#ifdef __cplusplus
}
#endif
#endif // COMPOUND_FILTER_H
We see that the FilterState struct is automatically translated to filter_state_t, and the CompoundState and CompoundOutput structs are likewise translated to compound_state_t and compound_output_t.
These can be nested in one another just as in the Python code, so for example a compound_state_t is composed of three filter_state_ts
State management¶
As before, the _step function will not copy the output state in res back into arg; it is still implemented as a “pure” function.
As a result, you do still need to manually manage the state by copying the data from res back the arg, but now this is much easier by virtue of C’s “direct assignment” semantics for structs.
Since all of the data in our structs is either a scalar, a statically-allocated array, or another struct (or statically-allocated array of structs) with the same property, copying all stateful data is now a one-liner, regardless of how complex the state is:
filter_arg.state = filter_res.state_new;
Using the generated code¶
Code using the generated API will now look like the following:
// Initialize as before
compound_filter_arg_t filter_arg;
compound_filter_res_t filter_res;
compound_filter_work_t filter_work;
compound_filter_init(&filter_arg, &filter_res, &filter_work);
// Set up the inputs to the function
filter_arg.u = read_sensor();
// Access nested states if necessary
check_state(&filter_arg.state.x_f);
// Evaluate the function numerically
compound_filter_step(&filter_arg, &filter_res, &filter_work);
// Copy state back to the argument
filter_arg.state = filter_res.state_new;
// Do something with the outputs
handle_outputs(filter_res.y.y_g, filter_res.y.y_h);
When you call the _step function, all that happens is that the pointers to the underlying data are “marshalled” into a pointer array and passed to CasADi, amounting to minimal overhead compared to the gain in flexibility and readability for more complex functions.
With this approach you can create maintainable and scalable data types that match how you think about parameters, plant models, and control algorithms.
Supported data types¶
We’ve already seen that codegen is supported for scalars (floats), NumPy arrays (up to 2D), and pytree_node structured data types.
Similar to PyTrees, Dicts and NamedTuples have named fields containing heterogeneous data; these can also be mapped to C struct types.
For instance,
class Point(NamedTuple):
x: float
y: float
state = {"pos": np.zeros(3), "vel": np.zeros(3)}
would generate the following C code:
typedef struct {
float x;
float y;
} point_t;
typedef struct {
float pos[3];
float vel[3];
} state_t;
Lists and tuples map to arrays, provided their data is “homogeneous”, meaning that all items have the same type.
That is, a list of three NamedTuple Point objects in Python will generate a C array like point_t points[3], but a list containing one Point and a dict will raise an error (since this cannot be mapped to a C array).
For containers with heterogeneous data types, use a NamedTuple, dict, or (recommended) pytree_node.
The full set of supported conversions is summarized in this table:
Python type |
C type |
|---|---|
|
|
|
|
|
|
|
|
|
|
These can compose with one another; you can have a dict whose entries are lists of NamedTuples: this will generate a struct containing arrays of structs.
This flexibility lets you create composite data structures in Python to represent complex input, output, and state and get a predictable C API that maps logically to the Python.
Summary¶
In this final part of the hardware deployment tutorial, we’ve finally put the pieces together to illustrate a scalable workflow for converting compatible Python functions to embedded-friendly C code.
This code generation provides a convenient workflow for writing and evaluating high-level logic in Python and then rapidly deploying to a range of C environments.
Where to Go From Here¶
As mentioned, this code generation and hardware deployment workflow is an early proof of concept, but it is central to the Archimedes roadmap.
If you’re using Archimedes for hardware control and have questions, comments, or requests, please don’t hesitate to reach out!
cleanup()