Introduction
As a programming exercise I recently wrote this automatic differentiation library based on forward accumulation. It should offer a simple way to create mathematical functions whose directional derivatives can be computed. For example:#include <iostream>
#include "autodiff.h"
int main()
{
mx::Var x, y;
mx::SField f = mx::sqrt(mx::pow(x,2)+mx::pow(y,2));
f.setVariables({&x, &y});
std::cout << f({3.0, 4.0}) << std::endl; //prints 5
std::cout << f.derivative({-1.0, 0.0}, x) << std::endl; //prints -1
return 0;
}
It basically works as follows:
The class SField (scalar field) holds a vector of Var pointers and a pointer to the top node of an expression tree consisting of Nodes:
ConstNode's just hold a constant value. They are created when anSFieldis initialized from a double.VarNodes hold the value of the variable objects (Var) they correspond to.UnOpandBinOphold pointers to basic mathematical operations (BaseUnOpandBaseBinOp) like addition and multiplication but also exponential and trigonometric functions. They also point to one or two input nodes which can be any of the node types in this list.
The objects of type BaseUnOp and BaseBinOp contain function pointers to the (hard coded) partial derivatives of the elementary operations. From this, the derivative can be computed by repeatedly applying the chain rule while traversing the expression tree. The base operation objects also contain the methods for combining two SField's (i.e. connecting the expression graphs).
The setVariable methods are necessary to specify in which order the variables shall be assigned when evaluating the function or its derivative.
Code
autodiff.h#ifndef AUTODIFF_H_INCLUDED
#define AUTODIFF_H_INCLUDED
#include <memory>
#include <vector>
namespace mx
{
class SField;
class Var;
namespace IMPLEMENTATION
{
class Node;
class VarNode;
class ConstNode;
class UnOp;
class BinOp;
class BaseUnOp;
class BaseBinOp;
}
class SField
{
friend IMPLEMENTATION::BaseUnOp;
friend IMPLEMENTATION::BaseBinOp;
private:
std::shared_ptr<IMPLEMENTATION::Node> start_Node;
std::vector<Var*> var_ptrs;
SField(std::shared_ptr<IMPLEMENTATION::Node> head_node) : start_Node(head_node) {}
SField(std::shared_ptr<IMPLEMENTATION::Node> head_node, const std::vector<Var*>& Var_ptrs)
: start_Node(head_node), var_ptrs(Var_ptrs) {}
public:
SField(const double& const_val);
SField() : SField(0) {}
SField(Var& x);
SField(std::vector<Var>& Var_ptrs) : SField() {setVariables(Var_ptrs);}
void setVariables(const std::vector<Var*>& Var_ptrs) {this->var_ptrs = Var_ptrs;}
void setVariables(std::vector<Var>& vars);
double operator()(const std::vector<double>&) const;
double operator()(const double&) const;
double derivative(const std::vector<double>&, const std::vector<double>&) const;
double derivative(const std::vector<double>&, const Var&) const;
double derivative(const double&, const double&) const;
double valueAtLastPosition() const;
SField& operator+=(const SField&);
SField& operator*=(const SField&);
SField& operator-=(const SField&);
SField& operator/=(const SField&);
};
class Var
{
friend SField;
private:
std::shared_ptr<IMPLEMENTATION::VarNode> vn;
public:
Var();
};
namespace IMPLEMENTATION
{
class Node
{
public:
mutable double value;
virtual double eval() const = 0;
virtual double derive() const = 0;
};
class VarNode : public Node
{
public:
double value_dot;
VarNode() = default;
double eval() const {return value;}
double derive() const {return value_dot;}
};
class ConstNode : public Node
{
public:
ConstNode(const double& const_val) {value = const_val;}
ConstNode() : ConstNode(0.0) {}
double eval() const {return value;}
double derive() const {return 0.0;}
};
class UnOp : public Node
{
private:
const BaseUnOp& f;
std::shared_ptr<Node> g;
public:
UnOp() = delete;
UnOp(const BaseUnOp& operation, std::shared_ptr<Node> input);
double eval() const;
double derive() const;
};
class BinOp : public Node
{
private:
const BaseBinOp& f;
std::shared_ptr<Node> g;
std::shared_ptr<Node> h;
public:
BinOp() = delete;
BinOp(const BaseBinOp& operation, std::shared_ptr<Node> input1, std::shared_ptr<Node> input2);
double eval() const;
double derive() const;
};
class BaseUnOp
{
typedef double fn(const double&);
private:
const fn* f;
public:
BaseUnOp() = delete;
BaseUnOp(fn* fun, fn* deriv) : f(fun), dfdx(deriv){}
double operator()(const double& x) const {return f(x);}
SField operator()(const SField& f) const;
const fn* dfdx;
};
class BaseBinOp
{
typedef double fn(const double&, const double&);
private:
const fn* f;
public:
BaseBinOp() = delete;
BaseBinOp(fn* fun, fn* x_deriv, fn* y_deriv) : f(fun), dfdx(x_deriv), dfdy(y_deriv){}
double operator()(const double& x, const double& y) const {return f(x,y);}
SField operator()(const SField&, const SField&) const;
const fn* dfdx;
const fn* dfdy;
};
///////////////////////////////////////////////////////////////////////////////////////////////
//LIST OF BASE OPERATIONS//////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
constexpr double b_add (const double& x, const double& y);
constexpr double b_add_del_x (const double& x, const double& y);
constexpr double b_add_del_y (const double& x, const double& y);
constexpr double b_multiply (const double& x, const double& y);
constexpr double b_multiply_del_x (const double& x, const double& y);
constexpr double b_multiply_del_y (const double& x, const double& y);
constexpr double b_subtract (const double& x, const double& y);
constexpr double b_subtract_del_x (const double& x, const double& y);
constexpr double b_subtract_del_y (const double& x, const double& y);
constexpr double b_divide (const double& x, const double& y);
constexpr double b_divide_del_x (const double& x, const double& y);
constexpr double b_divide_del_y (const double& x, const double& y);
double b_power (const double& x, const double& y);
double b_pow_del_x (const double& x, const double& y);
double b_pow_del_y (const double& x, const double& y);
double b_sqrt (const double& x);
double b_sqrt_del (const double& x);
constexpr double b_negate (const double& x);
constexpr double b_negate_del (const double& x);
double b_sin (const double& x);
double b_cos (const double& x);
double b_cos_del (const double& x);
double b_tan (const double& x);
double b_tan_del (const double& x);
double b_exp (const double& x);
double b_log (const double& x);
double b_log_del (const double& x);
}
///////////////////////////////////////////////////////////////////////////////////////////////
//GLOBAL BASE OPERATION OBJECTS////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
extern const IMPLEMENTATION::BaseBinOp add;
extern const IMPLEMENTATION::BaseBinOp multiply;
extern const IMPLEMENTATION::BaseBinOp subtract;
extern const IMPLEMENTATION::BaseBinOp divide;
extern const IMPLEMENTATION::BaseBinOp pow;
extern const IMPLEMENTATION::BaseUnOp sqrt;
extern const IMPLEMENTATION::BaseUnOp negation;
extern const IMPLEMENTATION::BaseUnOp sin;
extern const IMPLEMENTATION::BaseUnOp cos;
extern const IMPLEMENTATION::BaseUnOp tan;
extern const IMPLEMENTATION::BaseUnOp exp;
extern const IMPLEMENTATION::BaseUnOp log;
//OPERATOR OVERLOADING
SField operator+(const SField& f, const SField& g);
SField operator*(const SField& f, const SField& g);
SField operator-(const SField& f, const SField& g);
SField operator/(const SField& f, const SField& g);
SField operator-(const SField& f);
}
#endif // AUTODIFF_H_INCLUDED
autodiff.cpp
#include "autodiff.h"
#include <cmath>
#include <cassert>
namespace mx
{
///////////////////////////////////////////////////////////////////////////////////////////////
//SFIELD///////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
SField::SField(const double& const_val) : start_Node(new IMPLEMENTATION::ConstNode(const_val)) {}
SField::SField(Var& x) : start_Node(x.vn)
{
var_ptrs.push_back(&x);
}
void SField::setVariables(std::vector<Var>& vars)
{
var_ptrs.clear();
auto end = vars.end();
for (auto it = vars.begin(); it != end; ++it)
{
var_ptrs.push_back(&(*it));
}
}
double SField::operator()(const std::vector<double>& position) const
{
unsigned N = position.size();
assert(N == var_ptrs.size());
for (unsigned i = 0; i != N; ++i)
{
var_ptrs[i]->vn->value = position[i];
}
return start_Node->eval();
}
double SField::operator()(const double& position) const
{
return this->operator()(std::vector<double>{position});
}
double SField::derivative(const std::vector<double>& position, const std::vector<double>& direction) const
{
unsigned N = var_ptrs.size();
assert(N == position.size());
assert(N == direction.size());
for (unsigned i = 0; i != N; ++i)
{
var_ptrs[i]->vn->value = position[i];
var_ptrs[i]->vn->value_dot = direction[i];
}
return start_Node->derive();
}
double SField::derivative(const std::vector<double>& position, const Var& var_ptr) const
{
unsigned N = var_ptrs.size();
assert(N == position.size());
for (unsigned i = 0; i != N; ++i)
{
var_ptrs[i]->vn->value = position[i];
var_ptrs[i]->vn->value_dot = 0.0;
}
var_ptr.vn->value_dot = 1.0;
return start_Node->derive();
}
double SField::derivative(const double& position, const double& direction) const
{
return this->derivative(std::vector<double>{position},std::vector<double>{direction});
}
double SField::valueAtLastPosition() const
{
return start_Node->value;
}
SField& SField::operator+=(const SField& g) {return *this = *this+g;}
SField& SField::operator*=(const SField& g) {return *this = *this*g;}
SField& SField::operator-=(const SField& g) {return *this = *this-g;}
SField& SField::operator/=(const SField& g) {return *this = *this/g;}
///////////////////////////////////////////////////////////////////////////////////////////////
//VAR//////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
Var::Var() : vn(new IMPLEMENTATION::VarNode()) {}
namespace IMPLEMENTATION
{
///////////////////////////////////////////////////////////////////////////////////////////////
//NODE/////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
//CONSTNODE////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
//UNOP/////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
UnOp::UnOp(const BaseUnOp& operation, std::shared_ptr<Node> input) : f(operation), g(input)
{
value = operation(input->value);
}
double UnOp::eval() const
{
return f(g->eval());
}
double UnOp::derive() const
{
const double& g_dot = g->derive();
value = f(g->value);
return f.dfdx(g->value)*g_dot;
}
///////////////////////////////////////////////////////////////////////////////////////////////
//BINOP////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
BinOp::BinOp(const BaseBinOp& operation, std::shared_ptr<Node> input1, std::shared_ptr<Node> input2): f(operation), g(input1), h(input2)
{
value = operation(input1->value, input2->value);
}
double BinOp::eval() const
{
return f(g->eval(),h->eval());
}
double BinOp::derive() const
{
const double& g_dot = g->derive();
const double& h_dot = h->derive();
value = f(g->value, h->value);
return ((g_dot == 0.0) ? 0.0 : f.dfdx(g->value, h->value)*g_dot)
+((h_dot == 0.0) ? 0.0 : f.dfdy(g->value, h->value)*h_dot);
}
///////////////////////////////////////////////////////////////////////////////////////////////
//BASEUNOP/////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
SField BaseUnOp::operator()(const SField& f) const
{
std::shared_ptr<UnOp> op_ptr(new UnOp(*this, f.start_Node));
return SField(op_ptr, f.var_ptrs);
}
///////////////////////////////////////////////////////////////////////////////////////////////
//BASEBINOP////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
SField BaseBinOp::operator()(const SField& f, const SField& g) const
{
std::shared_ptr<BinOp> op_ptr(new BinOp(*this, f.start_Node, g.start_Node));
return SField(op_ptr, f.var_ptrs);
}
///////////////////////////////////////////////////////////////////////////////////////////////
//LIST OF BASE OPERATIONS//////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
constexpr double b_add (const double& x, const double& y) {return x+y;}
constexpr double b_add_del_x (const double& x, const double& y) {return 1.0;}
constexpr double b_add_del_y (const double& x, const double& y) {return 1.0;}
constexpr double b_multiply (const double& x, const double& y) {return x*y;}
constexpr double b_multiply_del_x (const double& x, const double& y) {return y;}
constexpr double b_multiply_del_y (const double& x, const double& y) {return x;}
constexpr double b_subtract (const double& x, const double& y) {return x-y;}
constexpr double b_subtract_del_x (const double& x, const double& y) {return 1.0;}
constexpr double b_subtract_del_y (const double& x, const double& y) {return -1.0;}
constexpr double b_divide (const double& x, const double& y) {return x/y;}
constexpr double b_divide_del_x (const double& x, const double& y) {return 1.0/y;}
constexpr double b_divide_del_y (const double& x, const double& y) {return -x/(y*y);}
double b_pow (const double& x, const double& y) {return std::pow(x,y);}
double b_pow_del_x (const double& x, const double& y) {return y*std::pow(x,y-1);}
double b_pow_del_y (const double& x, const double& y) {return (x == 0) ? 0 : std::pow(x,y)*std::log(x);}
double b_sqrt (const double& x) {return std::sqrt(x);}
double b_sqrt_del (const double& x) {return 0.5/std::sqrt(x);}
constexpr double b_negate (const double& x) {return -x;}
constexpr double b_negate_del (const double& x) {return -1.0;}
double b_sin (const double& x) {return std::sin(x);}
double b_cos (const double& x) {return std::cos(x);}
double b_cos_del (const double& x) {return -std::sin(x);}
double b_tan (const double& x) {return std::tan(x);}
double b_tan_del (const double& x) {double t = std::tan(x); return 1.0+t*t;}
double b_exp (const double& x) {return std::exp(x);}
double b_log (const double& x) {return std::log(x);}
double b_log_del (const double& x) {return 1.0/x;}
}
///////////////////////////////////////////////////////////////////////////////////////////////
//GLOBAL BASE OPERATION OBJECTS////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
const IMPLEMENTATION::BaseBinOp add(IMPLEMENTATION::b_add, IMPLEMENTATION::b_add_del_x, IMPLEMENTATION::b_add_del_y);
const IMPLEMENTATION::BaseBinOp multiply(IMPLEMENTATION::b_multiply, IMPLEMENTATION::b_multiply_del_x, IMPLEMENTATION::b_multiply_del_y);
const IMPLEMENTATION::BaseBinOp subtract(IMPLEMENTATION::b_subtract, IMPLEMENTATION::b_subtract_del_x, IMPLEMENTATION::b_subtract_del_y);
const IMPLEMENTATION::BaseBinOp divide(IMPLEMENTATION::b_divide, IMPLEMENTATION::b_divide_del_x, IMPLEMENTATION::b_divide_del_y);
const IMPLEMENTATION::BaseBinOp pow(IMPLEMENTATION::b_pow, IMPLEMENTATION::b_pow_del_x, IMPLEMENTATION::b_pow_del_y);
const IMPLEMENTATION::BaseUnOp sqrt(IMPLEMENTATION::b_sqrt, IMPLEMENTATION::b_sqrt_del);
const IMPLEMENTATION::BaseUnOp negation(IMPLEMENTATION::b_negate, IMPLEMENTATION::b_negate_del);
const IMPLEMENTATION::BaseUnOp sin(IMPLEMENTATION::b_sin, IMPLEMENTATION::b_cos);
const IMPLEMENTATION::BaseUnOp cos(IMPLEMENTATION::b_cos, IMPLEMENTATION::b_cos_del);
const IMPLEMENTATION::BaseUnOp tan(IMPLEMENTATION::b_tan, IMPLEMENTATION::b_tan_del);
const IMPLEMENTATION::BaseUnOp exp(IMPLEMENTATION::b_exp, IMPLEMENTATION::b_exp);
const IMPLEMENTATION::BaseUnOp log(IMPLEMENTATION::b_log, IMPLEMENTATION::b_log_del);
//OPERATOR OVERLOADS
SField operator+(const SField& f, const SField& g) {return add(f,g);}
SField operator*(const SField& f, const SField& g) {return multiply(f,g);}
SField operator-(const SField& f, const SField& g) {return subtract(f,g);}
SField operator/(const SField& f, const SField& g) {return divide(f,g);}
SField operator-(const SField& f) {return negation(f);}
}
Questions
I would be thankful if someone could comment on whether or not I was able to achieve the following goals (and of course if not, what I can improve):- Code structure, style and readability: Is the code comprehensible? Are there many bad practices?
- User Interface I: I tried to make adequate use of data encapsulation, in the sense that the user of the library has only access to those features he should, neither more nor less.
- User Interface II: The library should be user-friendly. I think I achieved that overall (probably not that hard with only two classes that the user has access to), but I wonder if there is a better way to set the order of the variables. At the moment one can either use
setVariableor already start with the corresponding constructor and afterwards only use compound operators (e.g.+=) to change theSField. This is allowed because (e.g.) the sum of twoSFieldcontains only theVarpointers of the left operand. constcorrectness and appropriate use of other keywords likemutable,constexpr,friend,...- Performance: Although I could avoid many unnecessary
evalcalls (in comparison to a prior version of this library) by caching thevaluein each node, I guess there are still ways to make the computation of the derivative more efficient.
Since performance requirements always depend on the application, I want to explain what I intend to do with the library:
I already used it to build a solver for arbitrary Hamiltonian systems and in particular computed the dynamics of a three-body system. I want to expand the library by vector fields and matrices and do more complicated physics simulations. As long as I don't want to simulate tens of thousands of particles I guess this should be possible even for real-time simulations.
I also plan to create a neural network class with the gradient descent being based on this library. I think here I will need at least \$10^4-10^5\$ variables (weights) already for simpler applications (e.g. training the MNIST set).
Do you think it is reasonable to use this library for those purposes? Of course, I could use a professional AD library, but since all I do is intended entirely for learning purposes and fun I think it is better to reuse my own code whenever possible.
Since this is my first question on SE, please don't hesitate to point out flaws in the question itself (too long, too broad, missing important details, inappropriate tags etc.).
