Created
September 9, 2017 18:30
-
-
Save bbbales2/e3f790a3f5496c0115cd5c64d8daaa3a to your computer and use it in GitHub Desktop.
File for benchmarking different ways to compute Jacobians
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Gotta add -fopenmp to make this compile