Updated ODE function docs

Author: bbbales2

stan/math/prim/fun.hpp

@@ -213,7 +213,6 @@
 #include <stan/math/prim/fun/multiply_lower_tri_self_transpose.hpp>
 #include <stan/math/prim/fun/norm.hpp>
 #include <stan/math/prim/fun/num_elements.hpp>
-#include <stan/math/prim/fun/ode_store_sensitivities.hpp>
 #include <stan/math/prim/fun/offset_multiplier_constrain.hpp>
 #include <stan/math/prim/fun/offset_multiplier_free.hpp>
 #include <stan/math/prim/fun/one_hot_array.hpp>

stan/math/prim/fun/ode_store_sensitivities.hpp

@@ -12,6 +12,7 @@
 #include <stan/math/prim/functor/integrate_ode_rk45.hpp>
 #include <stan/math/prim/functor/integrate_ode_std_vector_interface_adapter.hpp>
 #include <stan/math/prim/functor/ode_rk45.hpp>
+#include <stan/math/prim/functor/ode_store_sensitivities.hpp>
 #include <stan/math/prim/functor/map_rect.hpp>
 #include <stan/math/prim/functor/map_rect_combine.hpp>
 #include <stan/math/prim/functor/map_rect_concurrent.hpp>

stan/math/prim/functor.hpp

@@ -11,52 +11,8 @@ namespace stan {
 namespace math {
 
 /**
- * Return the solutions for the specified system of ordinary
- * differential equations given the specified initial state,
- * initial times, times of desired solution, and parameters and
- * data, writing error and warning messages to the specified
- * stream.
- *
- * <b>Warning:</b> If the system of equations is stiff, roughly
- * defined by having varying time scales across dimensions, then
- * this solver is likely to be slow.
- *
- * This function is templated to allow the initial times to be
- * either data or autodiff variables and the parameters to be data
- * or autodiff variables.  The autodiff-based implementation for
- * reverse-mode are defined in namespace <code>stan::math</code>
- * and may be invoked via argument-dependent lookup by including
- * their headers.
- *
- * This function uses the <a
- * href="http://en.wikipedia.org/wiki/Dormand–Prince_method">Dormand-Prince
- * method</a> as implemented in Boost's <code>
- * boost::numeric::odeint::runge_kutta_dopri5</code> integrator.
- *
- * @tparam F type of ODE system function.
- * @tparam T1 type of scalars for initial values.
- * @tparam T2 type of scalars for parameters.
- * @tparam T_t0 type of scalar of initial time point.
- * @tparam T_ts type of time-points where ODE solution is returned.
- * @param[in] f functor for the base ordinary differential equation.
- * @param[in] y0 initial state.
- * @param[in] t0 initial time.
- * @param[in] ts times of the desired solutions, in strictly
- * increasing order, all greater than the initial time.
- * @param[in] theta parameter vector for the ODE.
- * @param[in] x continuous data vector for the ODE.
- * @param[in] x_int integer data vector for the ODE.
- * @param[out] msgs the print stream for warning messages.
- * @param[in] relative_tolerance relative tolerance parameter
- *   for Boost's ode solver. Defaults to 1e-6.
- * @param[in] absolute_tolerance absolute tolerance parameter
- *   for Boost's ode solver. Defaults to 1e-6.
- * @param[in] max_num_steps maximum number of steps to take within
- *   the Boost ode solver.
- * @return a vector of states, each state being a vector of the
- * same size as the state variable, corresponding to a time in ts.
+ * @deprecated use <code>ode_rk45</code>
  */
-
 template <typename F, typename T1, typename T_param, typename T_t0,
           typename T_ts>
 std::vector<std::vector<return_type_t<T1, T_t0, T_ts, T_param>>>

stan/math/prim/functor/integrate_ode_rk45.hpp

@@ -11,10 +11,16 @@ namespace math {
 
 namespace internal {
 
+/**
+ * Wrap a functor designed for use with integrate_ode_bdf, integrate_ode_rk45,
+ * and integrate_ode_adams to use with the new ode_bdf/ode_rk45 interfaces.
+ *
+ * The old functors took the ODE state as a std::vector. The new ones take
+ * state as an Eigen::Matrix. The adapter converts to and from these forms
+ * so that the old ODE interfaces can work.
+ */
 template <typename F>
 struct integrate_ode_std_vector_interface_adapter {
-  // NOTE: I made this copy by value because
-  //  this function might go out of scope in reverse mode
   const F f_;
 
   integrate_ode_std_vector_interface_adapter(const F& f) : f_(f) {

stan/math/prim/functor/integrate_ode_std_vector_interface_adapter.hpp

@@ -3,7 +3,7 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/functor/coupled_ode_system.hpp>
-#include <stan/math/prim/fun/ode_store_sensitivities.hpp>
+#include <stan/math/prim/functor/ode_store_sensitivities.hpp>
 #include <boost/numeric/odeint/external/eigen/eigen_algebra.hpp>
 #include <boost/numeric/odeint.hpp>
 #include <ostream>
@@ -12,20 +12,94 @@
 namespace stan {
 namespace math {
 
+/**
+ * Solve the ODE initial value problem y' = f(t, y), y(t0) = y0 at a set of
+ * times, { t1, t2, t3, ... } using the non-stiff Runge-Kutta 45 solver in Boost
+ * with a relative tolerance of 1e-10, an absolute tolerance of 1e-10, and
+ * taking a maximum of 1e8 steps.
+ *
+ * If the system of equations is stiff, <code>ode_bdf</code> will likely be
+ * faster.
+ *
+ * \p f must define an operator() with the signature as:
+ *   template<typename T_t, typename T_y, typename... T_Args>
+ *   Eigen::Matrix<stan::return_type_t<T_t, T_y, T_Args...>, Eigen::Dynamic, 1>
+ *     operator()(const T_t& t, const Eigen::Matrix<T_y, Eigen::Dynamic, 1>& y,
+ *     std::ostream* msgs, const T_Args&... args);
+ *
+ * t is the time, y is the state, msgs is a stream for error messages, and args
+ * are optional arguments passed to the ODE solve function (which are passed
+ * through to \p f without modification).
+ *
+ * @tparam F Type of ODE right hand side
+ * @tparam T_0 Type of initial time
+ * @tparam T_ts Type of output times
+ * @tparam T_Args Types of pass-through parameters
+ *
+ * @param f Right hand side of the ODE
+ * @param y0 Initial state
+ * @param t0 Initial time
+ * @param ts Times at which to solve the ODE at. All values must be sorted and
+ *   not less than t0.
+ * @param relative_tolerance Relative tolerance passed to CVODES
+ * @param absolute_tolerance Absolute tolerance passed to CVODES
+ * @param max_num_steps Upper limit on the number of integration steps to
+ *   take between each output (error if exceeded)
+ * @param[in, out] msgs the print stream for warning messages
+ * @param args Extra arguments passed unmodified through to ODE right hand side
+ * @return Solution to ODE at times \p ts
+ */
 template <typename F, typename T_initial, typename T_t0, typename T_ts,
 	  typename... Args>
 std::vector<Eigen::Matrix<stan::return_type_t<T_initial, Args...>, Eigen::Dynamic, 1>>
 ode_rk45(const F& f, const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0, double t0,
          const std::vector<double>& ts, std::ostream* msgs,
          const Args&... args) {
-  double relative_tolerance = 1e-6;
-  double absolute_tolerance = 1e-6;
-  long int max_num_steps = 1e6;
+  double relative_tolerance = 1e-10;
+  double absolute_tolerance = 1e-10;
+  long int max_num_steps = 1e8;
 
   return ode_rk45_tol(f, y0, t0, ts, relative_tolerance, absolute_tolerance,
                       max_num_steps, msgs, args...);
 }
 
+/**
+ * Solve the ODE initial value problem y' = f(t, y), y(t0) = y0 at a set of
+ * times, { t1, t2, t3, ... } using the non-stiff Runge-Kutta 45 solver in Boost
+ * with a relative tolerance of 1e-10, an absolute tolerance of 1e-10, and
+ * taking a maximum of 1e8 steps.
+ *
+ * If the system of equations is stiff, <code>ode_bdf</code> will likely be
+ * faster.
+ *
+ * \p f must define an operator() with the signature as:
+ *   template<typename T_t, typename T_y, typename... T_Args>
+ *   Eigen::Matrix<stan::return_type_t<T_t, T_y, T_Args...>, Eigen::Dynamic, 1>
+ *     operator()(const T_t& t, const Eigen::Matrix<T_y, Eigen::Dynamic, 1>& y,
+ *     std::ostream* msgs, const T_Args&... args);
+ *
+ * t is the time, y is the state, msgs is a stream for error messages, and args
+ * are optional arguments passed to the ODE solve function (which are passed
+ * through to \p f without modification).
+ *
+ * @tparam F Type of ODE right hand side
+ * @tparam T_0 Type of initial time
+ * @tparam T_ts Type of output times
+ * @tparam T_Args Types of pass-through parameters
+ *
+ * @param f Right hand side of the ODE
+ * @param y0 Initial state
+ * @param t0 Initial time
+ * @param ts Times at which to solve the ODE at. All values must be sorted and
+ *   not less than t0.
+ * @param relative_tolerance Relative tolerance passed to CVODES
+ * @param absolute_tolerance Absolute tolerance passed to CVODES
+ * @param max_num_steps Upper limit on the number of integration steps to
+ *   take between each output (error if exceeded)
+ * @param[in, out] msgs the print stream for warning messages
+ * @param args Extra arguments passed unmodified through to ODE right hand side
+ * @return Solution to ODE at times \p ts
+ */
 template <typename F, typename T_initial, typename T_t0, typename T_ts,
 	  typename... Args>
 std::vector<Eigen::Matrix<stan::return_type_t<T_initial, T_t0, T_ts, Args...>, Eigen::Dynamic, 1>>
@@ -88,6 +162,8 @@ ode_rk45_tol(const F& f, const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0_a
   bool observer_initial_recorded = false;
   size_t time_index = 0;
 
+  max_step_checker step_checker(max_num_steps);
+  
   // avoid recording of the initial state which is included by the
   // conventions of odeint in the output
   auto filtered_observer
@@ -99,6 +175,7 @@ ode_rk45_tol(const F& f, const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0_a
     y.emplace_back(ode_store_sensitivities(f, coupled_state, y0,
 					   t0, ts[time_index],
 					   msgs, args...));
+    step_checker.reset();
     time_index++;
   };
 
@@ -113,7 +190,7 @@ ode_rk45_tol(const F& f, const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0_a
 			                   vector_space_algebra>()),
       std::ref(coupled_system), initial_coupled_state,
       std::begin(ts_vec), std::end(ts_vec), step_size, filtered_observer,
-      max_step_checker(max_num_steps));
+      step_checker);
 
   return y;
 }

stan/math/prim/functor/ode_rk45.hpp

@@ -0,0 +1,40 @@
+#ifndef STAN_MATH_PRIM_FUNCTOR_ODE_STORE_SENSITIVITIES_HPP
+#define STAN_MATH_PRIM_FUNCTOR_ODE_STORE_SENSITIVITIES_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <ostream>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * Build output for state of ODE solve from the current coupled_state when
+ * all arguments are arithmetic.
+ * 
+ * @tparam F Type of ODE right hand side
+ *
+ * @param f Right hand side of the ODE
+ * @param coupled_state Current state of the coupled_ode_system
+ * @param y0 Initial state
+ * @param t0 Initial time
+ * @param ts Times at which to solve the ODE at
+ * @param[in, out] msgs the print stream for warning messages
+ * @param args Extra arguments passed unmodified through to ODE right hand side
+ * @return ODE state
+ */
+template <typename F, typename... Args,
+          typename = require_all_arithmetic_t<scalar_type_t<Args>...>>
+Eigen::VectorXd ode_store_sensitivities(const F& f,
+					const Eigen::VectorXd& coupled_state,
+					const Eigen::VectorXd& y0,
+					double t0,
+					double t,
+					std::ostream* msgs,
+					const Args&... args) {
+  return coupled_state.head(y0.size());
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/functor/ode_store_sensitivities.hpp

@@ -15,6 +15,7 @@
 #include <stan/math/rev/functor/integrate_ode_adams.hpp>
 #include <stan/math/rev/functor/integrate_ode_bdf.hpp>
 #include <stan/math/rev/functor/ode_bdf.hpp>
+#include <stan/math/rev/functor/ode_store_sensitivities.hpp>
 #include <stan/math/rev/functor/jacobian.hpp>
 #include <stan/math/rev/functor/kinsol_data.hpp>
 #include <stan/math/rev/functor/kinsol_solve.hpp>