wrenfold

Introduction

wrenfold is a framework for converting symbolic mathematical expressions (written in python) into generated code in compiled languages (C++, Rust). It aims to bridge the gap between prototyping of functions in expressive symbolic form, and performant production code. wrenfold is particularly relevant to domains where numerical optimization is employed to solve differentiable objective functions, such as robotics or computer vision.

Using wrenfold, mathematical functions can be expressed and composed succinctly in python:

from wrenfold import code_generation, sym
from wrenfold.type_annotations import Vector3

def angular_distance(a: Vector3, b: Vector3):
    """
    A simple example function: We compute the angle between two vectors. The angle is returned, and
    the Jacobian with respect to `a` is passed as an output argument. This might be a cost in an
    optimization, for instance.
    """
    dot = (a.T * b)[0]
    cos_theta = dot / (a.norm() * b.norm())
    theta = sym.acos(cos_theta)
    theta_D_a = sym.jacobian([theta], a)

    # Our generated function will return `theta`, and pass `theta_D_a` as an output arg.
    return (
        code_generation.ReturnValue(theta),
        code_generation.OutputArg(theta_D_a, "theta_D_a"),
    )

And corresponding compilable code can be obtained easily:

# CppGenerator can be swapped out for RustGenerator to obtain Rust. You can implement your own
# custom generator to target a new language - or override methods on the provided generators in
# order to customize the output code to your liking.
cpp = code_generation.generate_function(angular_distance, code_generation.CppGenerator())
print(cpp)

template <typename Scalar, typename T0, typename T1, typename T2>
Scalar angular_distance(const T0& a, const T1& b, T2&& theta_D_a) {
  auto _a = wf::make_input_span<3, 1>(a);
  auto _b = wf::make_input_span<3, 1>(b);
  auto _theta_D_a = wf::make_output_span<1, 3>(theta_D_a);

  const Scalar v007 = _b(2, 0);
  const Scalar v006 = _a(2, 0);
  const Scalar v004 = _b(1, 0);

  // ... Output code is truncated for brevity.

  const Scalar v009 = v000 * v001 + v003 * v004 + v006 * v007;
  const Scalar v021 = v001 * v001 + v004 * v004 + v007 * v007;

  // ...

  _theta_D_a(0, 0) = (v000 * v072 + v001 * v017) * v073;
  _theta_D_a(0, 1) = (v003 * v072 + v004 * v017) * v073;
  _theta_D_a(0, 2) = (v006 * v072 + v007 * v017) * v073;
  return std::acos(v009 * v017 * v023);
}

wrenfold draws inspiration from SymForce, but differs in a few key ways:

• Improved flexibility: Symbolic expressions can include conditional logic. This enables a broader range of functions to be generated.

• Ease of integration: wrenfold aims to make it straightforward to customize the code-generation step to suit your project. For example, you can use existing types in your codebase in generated method signatures.

• Faster code generation: Faster code generation translates to quicker iteration on experiments. The generation cost should ideally be negligible compared to compile time for the code itself.

• Narrower scope: wrenfold does not implement a numerical optimizer. Rather we aim to make it simple to integrate generated code into your project’s existing preferred optimizer (see the extended examples). It should be relatively straightforward to use wrenfold functions with GTSAM, Ceres, the SymForce optimizer, or your own custom implementation.

wrenfold is primarily written in C++, and exposes a python API via pybind11. It can presently generate code in C++17 and Rust.

Motivation

Why use symbolic code generation for mathematical functions? The SymForce paper outlines some of the rationale. In our opinion, the two main arguments are:

• Faster iteration:

  • Functions can be written quickly and expressively in python, enabling rapid prototyping. Over time, users acquire a library of composable expressions that can be combined easily to form new symbolic functions.

  • Derivatives are obtained automatically, without spending time debugging manually chain-ruled Jacobians.

• Improved runtime performance:

  • The performance of generated methods is often competitive with handwritten implementations, and can meaningfully exceed results obtained with runtime auto-diff.

  • Generated methods are fast enough to deploy on a production robot, enabling a quicker pipeline from offline prototyping to real-world testing.

  • A prudent caveat for any performance related claim: Your mileage may vary depending on expression complexity and the degree of effort exerted in optimizing different implementations.

Getting started

1. Begin with the quick start guide.

2. Peruse the user guide.

3. Check out examples in the repository. There are also additional examples that demonstrate integration of generated code into third-party optimizers.

Goals

• Enable fast iteration time. Symbolic operations and code generation must be quick enough to allow rapid experimentation with complex expressions.

• Generate reasonably performant code.

• Emphasize ease of integration of generated code. This manifests in three ways:

  • We provide a straightforward mechanism for incorporating existing types from the target code-base in generated functions.

  • The transpilation step emits a general purpose abstract syntax tree (AST) that can be converted to a new language with relative ease.

  • Generated code can invoke external handwritten functions.

To get a sense of what features are in development, refer to the Issue Tracker.

Non-goals

• Provide a full-featured computer algebra system (CAS) - realistically this is too ambitious. Instead, we implement mathematical operations on a need-to-have basis. wrenfold expressions can be converted to and from SymPy in order to unlock additional functionality.

• Be a one-stop shop for optimization. In most cases we expect the user to have an existing optimizer. Instead, it should be easy to customize generated code to better integrate with existing solutions.