Adams Moulton ODE solver in Stan useful?

General
11

Here’s another example, related to the orbital dynamics that @s.maskell mentioned here. In the ODE system(2body), I’m using the default adams/bdf tolerance, and Adams is faster than the other two. Then I changed the ODE system to a trivial linear one defined by identity matrix, and use a “very low” tolerance, and this time adams is still faster. The example is intended to demonstrate that without involving the nature of ODE system, talking about “high” & “low” tolerance is meaningless, let alone if an integrator “is useful”.

class two_body {
 public:
  template <typename T0, typename T1, typename T2>
  inline std::vector<stan::return_type_t<T1, T2>>
  operator()(const T0& t_in, const std::vector<T1>& y,
             const std::vector<T2>& theta, const std::vector<double>& x,
             const std::vector<int>& x_int, std::ostream* msgs) const {
    if (y.size() != 4)
      throw std::domain_error("Functor called with inconsistent state");

    T1 r = sqrt(y[0]*y[0] + y[2]*y[2]);
    std::vector<stan::return_type_t<T1, T2>> dydt(4);
    dydt[0] = y[1];
    dydt[1] = -y[0]/pow(r, 3.0);
    dydt[2] = y[3];
    dydt[3] = -y[2]/pow(r, 3.0);
    return dydt;
  }
};

TEST(demo, two_body) {
  double t0 = 0.0;
  std::vector<double> ts{1.0, 100.0};
  std::vector<double> theta;
  std::vector<double> x_r;
  std::vector<int> x_i;

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_rk45(two_body(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-10, 1.e-10, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "rk45 timed: " << duration.count() << "\n";
  }

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_bdf(two_body(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-10, 1.e-10, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "bdf timed: " << duration.count() << "\n";
  }

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_adams(two_body(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-10, 1.e-10, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "adams timed: " << duration.count() << "\n";
  }
}

class degen {
 public:
  template <typename T0, typename T1, typename T2>
  inline std::vector<stan::return_type_t<T1, T2>>
  operator()(const T0& t_in, const std::vector<T1>& y,
             const std::vector<T2>& theta, const std::vector<double>& x,
             const std::vector<int>& x_int, std::ostream* msgs) const {
    if (y.size() != 4)
      throw std::domain_error("Functor called with inconsistent state");

    T1 r = sqrt(y[0]*y[0] + y[2]*y[2]);
    std::vector<stan::return_type_t<T1, T2>> dydt(4);
    dydt[0] = -y[0];
    dydt[1] = -y[1];
    dydt[2] = -y[2];
    dydt[3] = -y[3];
    return dydt;
  }
};

TEST(demo, rk45_degen) {
  double t0 = 0.0;
  std::vector<double> ts{1.0, 100.0};
  std::vector<double> theta;
  std::vector<double> x_r;
  std::vector<int> x_i;

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_rk45(degen(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-14, 1.e-16, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "rk45 timed: " << duration.count() << "\n";
  }

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_bdf(degen(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-14, 1.e-16, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "bdf timed: " << duration.count() << "\n";
  }

  {
    std::vector<stan::math::var> y0{2, 0, 0, 0.5};
    auto start = high_resolution_clock::now(); 
    std::vector<std::vector<stan::math::var>> ys = stan::math::integrate_ode_adams(degen(), y0, t0, ts, theta, x_r, x_i, 0, 1.e-14, 1.e-16, 100000);
    auto stop = high_resolution_clock::now(); 
    auto duration = duration_cast<microseconds>(stop - start);
    std::cout << "adams timed: " << duration.count() << "\n";
  }
}

output

mpirun -np 1 test/unit/math/rev/integrate_ode_adams_2_perf_test --gtest_output="xml:test/unit/math/rev/integrate_ode_adams_2_perf_test.xml"
Running main() from lib/gtest_1.8.1/src/gtest_main.cc
[==========] Running 2 tests from 1 test case.
[----------] Global test environment set-up.
[----------] 2 tests from demo
[ RUN      ] demo.two_body
rk45 timed: 12119
bdf timed: 16825
adams timed: 10083
[       OK ] demo.two_body (39 ms)
[ RUN      ] demo.rk45_degen
rk45 timed: 5143
bdf timed: 5859
adams timed: 2391
[       OK ] demo.rk45_degen (13 ms)
[----------] 2 tests from demo (52 ms total)

[----------] Global test environment tear-down
[==========] 2 tests from 1 test case ran. (53 ms total)
[  PASSED  ] 2 tests.