gradgen 0.4.0a2

Symbolic differentiation and Rust code generation library.

Gradgen

What it does

Gradgen is a Python library for symbolic differentiation and (embedded) Rust code generation.

Code generation example

See demo

Here is an example where we will define the function

$$f(x, u) = \Vert x \Vert_2^2 + u \sin(x_1) + x_2 x_3,$$

for a three-dimensional input $x$ and scalar $u$.

The goal is to generate Rust code for the functions $f$, $Jf$ (the Jacobian matrix of $f$).

Furthermore, we want to generate a Rust function that computes simultaneous $f$ and $\nabla_x f$. This often is computationally more efficient compared to computing $f(x, u)$ and $\nabla_x f(x, u)$ in separate functions (look for FunctionBundle below).

from gradgen import CodeGenerationBuilder, Function, RustBackendConfig, SXVector, sin

# Define the symbolic inputs.
x = SXVector.sym("x", 3)
u = SXVector.sym("u", 1)

# Build a simple scalar-valued function of x and u
# f(x, u) = ||x||_2^2 + u_1 * sin(x_1) + x_2 * x_3
f_expr = x.norm2sq() + u[0] * sin(x[0]) + x[1] * x[2]

# Define a Function object
f = Function(
    "energy",
    [x, u],
    [f_expr],
    input_names=["x", "u"],
    output_names=["energy"],
)

# (Optional) Evaluate f in Python
x_value = [1.0, 2.0, -0.5]
u_value = [3.0]
print("f(x, u) =", f(x_value, u_value))

# Generate code
project = (
    CodeGenerationBuilder()
    .with_backend_config(
        RustBackendConfig()
        .with_crate_name("my_kernel")
        .with_backend_mode("no_std")
        .with_scalar_type("f64")
    )
    .for_function(f)
        .add_primal()
        .add_jacobian()
        .add_joint(
            FunctionBundle()
            .add_f()
            .add_jf(wrt=0)
        )
        .with_simplification("medium")
        .done()
    .build(Path(__file__).resolve().parent / "codegen_kernel")
)

See the demos and this more complete tutorial.

Special case: optimal control

See demo

In applications such as optimal control, the generated code can become too large very easily. However, the problem structure can be exploited to generate code with complexity that doesn't increase with the prediction horizon.

Instead of completely unrolled code, Gradgen exploits the problem structure to create high-performance, human-readable embeddable Rust code.

See the demos and this more complete tutorial for details.

Unique features

• Truly embdedable safe Rust code with optional #[no_std], no dynamic memory allocation, no panic!s
• Specialised code generation tools for optimal control problems (demo)
• Very efficient code generation thanks to modular code generation using map, zip, repeat, and chain high-order functions.
• Supports both single (f32) and double (f64) precision arithmetic. You can use libm or your own library for numerical operations (e.g., micromath)

Where to go next?

See the demos and this more complete tutorial for details.

Show us some love!

If you find Gradgen useful, give us a star on GitHub!