Quick start guide

Installing the python package

Python wheels are available on PyPi. Install the latest version of wrenfold with:

pip install wrenfold

wrenfold is also available as a conda-forge package:

conda install -c conda-forge wrenfold

Alternatively, python wheels may also be obtained from the GitHub Releases Page. Select the whl file appropriate to your OS and python version. For example, for python 3.10 on arm64 OSX you would download and install wrenfold-0.1.0-cp310-cp310-macosx_11_0_arm64.whl:

pip install wrenfold-0.1.0-cp310-cp310-macosx_11_0_arm64.whl

Validating the install

After installation, open a python REPL and test that wrenfold can be imported:

import wrenfold
print(wrenfold.__version__)  # prints: 0.1.0

from wrenfold import sym

x, y = sym.symbols('x, y')
f = sym.cos(x * y)
g = f.diff(x)
print(g)  # prints: -y * sin(x * y)

Generating your first function

To illustrate the wrenfold workflow, we will generate a small example function and then invoke the generated code in C++ and Rust. To keep the code brief, we will choose something very simple for our example - the Rosenbrock function:

\[f\left(x, y\right) = \left(a - x\right)^2 + b\cdot\left(y - x^2\right)^2\]

First we express the function as a python function that manipulates wrenfold types:

from wrenfold import code_generation, sym, type_annotations


def rosenbrock(
    xy: type_annotations.Vector2,
    a: type_annotations.FloatScalar,
    b: type_annotations.FloatScalar,
):
    """Evaluates the Rosenbrock function and its first derivative wrt `x` and `y`."""
    x, y = xy
    f = (a - x) ** 2 + b * (y - x ** 2) ** 2
    return (code_generation.ReturnValue(f),
            code_generation.OutputArg(sym.jacobian([f], xy), name="f_D_xy"))

The argument type annotations let wrenfold know what dimensions and numerical types to expect for the input arguments to \(f\). Our python function returns two outputs:

  1. The first is the value of the Rosenbrock function, \(f\). We specify that this will be the return value of the generated C++ function.

  2. The second is the Jacobian of \(f\), taken with respect to \(\left[x, y\right]\), which will be an output argument of the generated C++ function.

To generate some actual code, we run:

# Generate the function as C++, and apply some boilerplate (imports and namespace).
cpp = code_generation.generate_function(rosenbrock, generator=code_generation.CppGenerator())
cpp = code_generation.CppGenerator.apply_preamble(cpp, namespace="gen")
print(cpp)

Which produces:

// Machine generated code.
#pragma once
#include <cmath>
#include <cstdint>

#include <wrenfold/span.h>

namespace gen {

template <typename Scalar, typename T0, typename T3>
Scalar rosenbrock(const T0& xy, const Scalar a, const Scalar b, T3&& f_D_xy) {
  auto _xy = wf::make_input_span<2, 1>(xy);
  auto _f_D_xy = wf::make_output_span<1, 2>(f_D_xy);

  // Operation counts:
  // add: 5
  // multiply: 7
  // negate: 2
  // total: 14

  const Scalar v003 = _xy(0, 0);
  const Scalar v041 = -v003;
  const Scalar v001 = _xy(1, 0);
  const Scalar v006 = v001 + v003 * v041;
  const Scalar v000 = b;
  const Scalar v036 = v000 * v006;
  const Scalar v040 = static_cast<Scalar>(2) * v036;
  const Scalar v008 = a;
  const Scalar v010 = v008 + v041;
  _f_D_xy(0, 0) = static_cast<Scalar>(2) * (v003 + -(v008 + v003 * v040));
  _f_D_xy(0, 1) = v040;
  return v006 * v036 + v010 * v010;
}

}  // namespace gen

A couple of observations about this output:

  • The vector-valued input and output arguments were mapped to generic types. The generated function constructs n-dimensional spans in order to read/write to the vectors xy and f_D_xy.

  • Terms that appear in both f and f_D_xy were extracted and shared between both outputs.

  • Our generated code depends on the wrenfold runtime, a small header-only library that provides the wf::span type.

Tip

The runtime is also installed by the python wheel. It can be found at <VENV ROOT>/include/site/python3.XX/wrenfold.

Calling the generated C++

Next, we will create a simple C++ program that evaluates our generated function. wrenfold functions can be made to work with any dense matrix representation. In this tutorial we will use Eigen, as it is one of the most popular choices and wrenfold supports it out of the box.

Our C++ file looks like:

#include <iostream>

// Defining `WF_SPAN_EIGEN_SUPPORT` will enable automatic conversion of Eigen types to
// wrenfold spans.
#define WF_SPAN_EIGEN_SUPPORT
#include <wrenfold/span.h>

// We assume our generated function was saved in `rosenbrock.h`:
#include "rosenbrock.h"

int main() {
  constexpr double a = 2.0;
  constexpr double b = 10.0;

  // The global minimum is (a, a^2):
  const Eigen::Vector2d xy_minimum{a, a * a};

  // Note that we can pass RowVector2d directly for the `f_D_xy` argument:
  Eigen::RowVector2d f_D_xy;
  const double f = gen::rosenbrock(xy_minimum, a, b, f_D_xy);

  // A crude test: at the minimum f should be zero, and f_D_xy = [0, 0]
  std::cout << "f = " << f << "\n";
  std::cout << "f_D_xy = [" << f_D_xy.x() << ", " << f_D_xy.y() << "]" << std::endl;
  return 0;
}

To compile it we need to provide include paths for Eigen and the wrenfold runtime headers. You may need to adjust these paths for your system.

g++ -std=c++17 -I/usr/local/include/eigen3 -I<WRENFOLD REPO>/components/runtime main.cpp
./a.out

And sure enough our test program outputs:

f = 0
f_D_xy = [0, 0]

Tip

The wrenfold repo contains a more complete version of this example.

Calling generated Rust

By swapping out CppGenerator for RustGenerator, we can emit rust code instead:

rust = code_generation.generate_function(rosenbrock, generator=code_generation.RustGenerator())
rust = code_generation.RustGenerator.apply_preamble(rust)
print(rust)

Which produces:

//! Machine generated code.
#![cfg_attr(rustfmt, rustfmt_skip)]

#[inline]
#[allow(non_snake_case, clippy::unused_unit, clippy::collapsible_else_if, clippy::needless_late_init, unused_variables)]
pub fn rosenbrock<T0, T3, >(xy: &T0, a: f64, b: f64, f_D_xy: &mut T3) -> f64
where
  T0: wrenfold_traits::Span2D<2, 1, ValueType = f64>,
  T3: wrenfold_traits::OutputSpan2D<1, 2, ValueType = f64>,
{
  // Operation counts:
  // add: 5
  // multiply: 7
  // negate: 2
  // total: 14

  let v003: f64 = xy.get(0, 0);
  let v041: f64 = -v003;
  let v001: f64 = xy.get(1, 0);
  let v006: f64 = v001 + v003 * v041;
  let v000: f64 = b;
  let v036: f64 = v000 * v006;
  let v040: f64 = (2i64) as f64 * v036;
  let v008: f64 = a;
  let v010: f64 = v008 + v041;
  f_D_xy.set(0, 0, (2i64) as f64 * (v003 + -(v008 + v003 * v040)));
  f_D_xy.set(0, 1, v040);
  v006 * v036 + v010 * v010
}

Like the C++ example, the generated rust code has a small runtime dependency. In this instance, it is a set of traits for constraining the input and output arguments. These traits are found in the wrenfold-traits crate. The wrenfold-traits crate includes default implementations for use with nalgebra.

To test our generated function, we will create a small test crate with the following Cargo.toml:

[package]
name = "rosenbrock_test"
version = "0.1.0"
edition = "2021"
publish = false

[dependencies]
wrenfold-traits = { version = "0.1.0", features = ["nalgebra"] }
nalgebra = { version = "0.33" }

[workspace]

Tip

Note that we enabled features = ["nalgebra"] for the wrenfold-traits crate.

And the following lib.rs:

#[cfg(test)]
mod generated; //  We assume the generated output was written to `generated.rs`.

#[cfg(test)]
mod tests {
    pub use super::*;
    use nalgebra as na;

    #[test]
    fn test_rosenbrock() {
        let a = 2.0;
        let b = 10.0;
        let xy_min = na::Vector2::new(a, a * a);

        let mut f_D_xy = na::SMatrix::<f64, 1, 2>::zeros();
        let f = generated::rosenbrock(&xy_min, a, b, &mut f_D_xy);

        assert_eq!(0.0, f);
        assert_eq!(0.0, f_D_xy[0]);
        assert_eq!(0.0, f_D_xy[1]);
    }
}