File for benchmarking different ways to compute Jacobians

test_jacobian.cpp

#include <iostream>
#include "stan/math.hpp"
#include "stan/math/fwd/mat.hpp"

using namespace stan::math;

template<typename T1, typename T2>
std::vector<typename stan::return_type<T1, T2>::type > ode(const double& t,
                   const std::vector<T1>& u,
                   const std::vector<T2>& D,
                   const std::vector<double>& x_r,
                   const std::vector<int>& x_i,
                   std::ostream* pstream) {
  int N = u.size();
  std::vector<typename stan::return_type<T1, T2>::type > dudt(N, 0.0);
  double dx = x_r[0];

  dudt[0] = (D[1] * (u[1] - u[0]) - D[0] * (u[0] - 1.0)) / (dx * dx);
  dudt[N - 1] = (D[N] * (1.0 - u[N - 1]) - D[N - 1] * (u[N - 1] - u[N - 2])) / (dx * dx);

  for(int i = 1; i < N - 1; i++) {
    dudt[i] = (D[i + 1] * (u[i + 1] - u[i]) - D[i] * (u[i] - u[i - 1])) / (dx * dx);
  }

  return dudt;
}

struct ode_functor {
  template<typename T1, typename T2>
  std::vector<typename stan::return_type<T1, T2>::type > operator()(const double& t,
                            const std::vector<T1>& state,
                            const std::vector<T2>& theta,
                            const std::vector<double>& x_r,
                            const std::vector<int>& x_i,
                            std::ostream* pstream__) const {
    return ode(t, state, theta, x_r, x_i, pstream__);
  }
};

template<typename T_func>
inline void jacobian(double t,
                     T_func f_,
                     const std::vector<double>& y,
                     const std::vector<double>& theta,
                     const std::vector<double>& x_r,
                     const std::vector<int>& x_i,
                     Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& Jy,
                     Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& Jtheta) {
  std::ostream* pstream(NULL);
  using Eigen::Dynamic;
  using Eigen::Map;
  using Eigen::Matrix;
  using Eigen::RowVectorXd;
  using std::vector;
  Jy.resize(y.size(), y.size());
  Jtheta.resize(y.size(), theta.size());
  vector<double> grad(y.size() + theta.size());
  Map<RowVectorXd> grad_eig(&grad[0], y.size() + theta.size());
  try {
    start_nested();
    vector<var> y_var(y.begin(), y.end());
    vector<var> thetavar(theta.begin(), theta.end());
    vector<var> z_var;
    z_var.reserve(y.size() + theta.size());
    z_var.insert(z_var.end(), y_var.begin(), y_var.end());
    z_var.insert(z_var.end(), thetavar.begin(), thetavar.end());
    vector<var> dy_dt_var = f_(t, y_var, thetavar, x_r, x_i, pstream);
    for (size_t i = 0; i < dy_dt_var.size(); ++i) {
      set_zero_all_adjoints_nested();
      dy_dt_var[i].grad(z_var, grad);
      Jy.row(i) = grad_eig.leftCols(y.size());
      Jtheta.row(i) = grad_eig.rightCols(theta.size());
    }
  } catch (const std::exception& e) {
    recover_memory_nested();
    throw;
  }
  recover_memory_nested();
}

template<typename T_func>
inline void jacobian_fwd(double t,
                         T_func f_,
                         const std::vector<double>& y,
                         const std::vector<double>& theta,
                         const std::vector<double>& x_r,
                         const std::vector<int>& x_i,
                         Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& Jy,
                         Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& Jtheta) {
  std::ostream* pstream(NULL);
  using Eigen::Dynamic;
  using Eigen::Map;
  using Eigen::Matrix;
  using Eigen::RowVectorXd;
  using std::vector;
  Jy.resize(y.size(), y.size());
  Jtheta.resize(y.size(), theta.size());
  vector<fvar<double> > y_var(y.begin(), y.end());
  vector<fvar<double> > thetavar(theta.begin(), theta.end());
  
  for(int i = 0; i < y.size(); i++) {
    if(i > 0)
      y_var[i - 1].d_ = 0.0;
    y_var[i].d_ = 1.0;
    
    vector<fvar<double> > dy_dt_var = f_(t, y_var, thetavar, x_r, x_i, pstream);
    
    for(int j = 0; j < y.size(); j++)
      Jy(j, i) = dy_dt_var[j].tangent();
  }

  y_var[y.size() - 1].d_ = 0.0;

  for(int i = 0; i < theta.size(); i++) {
    if(i > 0)
      thetavar[i - 1].d_ = 0.0;
    thetavar[i].d_ = 1.0;
    
    vector<fvar<double> > dy = f_(t, y_var, thetavar, x_r, x_i, pstream);
    
    for(int j = 0; j < dy.size(); j++)
      Jtheta(j, i) = dy[j].tangent();
  }
}

int main(int argc, char **argv) {
  std::ostream pstream(std::cout.rdbuf());
  int N = 50;
  double dx = 1.0 / N;

  std::vector<double> y;
  std::vector<double> theta;
  std::vector<double> x_r = {dx};
  std::vector<int> x_i;

  for(int i = 0; i < N; i++) {
    y.push_back(i);
    theta.push_back(i);
  }

  theta.push_back(N + 1);
  
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Jy1;
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Jtheta1;

  double tmp = omp_get_wtime();
  for(int i = 0; i < 10000; i++) {
    jacobian(0.0, ode_functor{}, y, theta, x_r, x_i, Jy1, Jtheta1);
  }
  std::cout << "Reverse mode: " << (omp_get_wtime() - tmp) / 10000 << std::endl;

  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Jy2;
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Jtheta2;

  tmp = omp_get_wtime();
  for(int i = 0; i < 10000; i++) {
    jacobian_fwd(0.0, ode_functor{}, y, theta, x_r, x_i, Jy2, Jtheta2);
  }
  std::cout << "Forward mode: " << (omp_get_wtime() - tmp) / 10000 << std::endl;

  std::cout << "Jy error: " << (Jy1 - Jy2).array().abs().maxCoeff() << std::endl;
  std::cout << "Jtheta error: " << (Jtheta1 - Jtheta2).array().abs().maxCoeff() << std::endl;
}

Gotta add -fopenmp to make this compile