Instantly share code, notes, and snippets.

Embed
What would you like to do?
ODE example
// Build with: clang++ -Wall -I . -isystem lib/eigen_3.3.3 -isystem lib/boost_1.64.0 -isystem lib/cvodes_2.9.0/include -std=c++1y -DBOOST_RESULT_OF_USE_TR1 -DBOOST_NO_DECLTYPE -DBOOST_DISABLE_ASSERTS -Wno-unused-function -Wno-uninitialized -DNO_FPRINTF_OUTPUT -pipe test2.cpp -O3 -o test2 lib/cvodes_2.9.0/lib/libsundials_nvecserial.a lib/cvodes_2.9.0/lib/libsundials_cvodes.a
#include <iostream>
#include "stan/math.hpp"
using namespace stan::math;
// This is a 1D diffusion PDE with a spatially varying diffusion coefficient.
// So if the domain is like this:
// x--o--x--o--x--o--x
// Then diffusion coefficients are defined at the xs and the state is at the o
// values. This is an N = 3 system. There are N + 1 parameters.
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__);
}
};
int main(int argc, char **argv) {
std::ostream pstream(std::cout.rdbuf());
int N = 20;
double dx = 1.0 / N;
for(int j = 0; j < 100; j++) {
std::vector<var> D;
std::vector<var> y0;
std::vector<double> ts;
std::vector<double> x_r = { dx };
std::vector<int> x_i;
for(int i = 0; i < N; i++) {
y0.push_back(0.0);
D.push_back(1.0e-3);
}
D.push_back(1.0e-3);
// Start with a delta in the leftmost point of space
y0[0] = 100.0;
//for(int i = 0; i < 10; i++)
// ts.push_back((i + 1) * 1.0);
ts.push_back(10.0);
std::vector<std::vector<var> > y = integrate_ode_bdf(ode_functor(),
y0,
0.0,
ts,
D,
x_r,
x_i,
&pstream);
// Make our output just the sum of everything
var sum = 0.0;
for(int i = 0; i < ts.size(); i++) {
for(int j = 0; j < y[i].size(); j++) {
sum += y[i][j];
}
}
sum.grad();
for(int i = 0; i < D.size(); i++)
std::cout << D[i].adj() << " ";
std::cout << std::endl;
//print_stack(pstream);
set_zero_all_adjoints();
recover_memory(); // Not really haha!
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment