Feature/adjoint odes

stan/math/rev/functor/cvodes_integrator_adjoint.hpp

@@ -38,18 +38,14 @@ class cvodes_integrator_adjoint_vari : public vari_base {
 
   static constexpr bool is_var_ts_{is_var<T_ts>::value};
   static constexpr bool is_var_t0_{is_var<T_t0>::value};
+  static constexpr bool is_var_y0_{is_var<T_y0>::value};
   static constexpr bool is_var_y0_t0_{is_var<T_y0_t0>::value};
   static constexpr bool is_any_var_args_{
       disjunction<is_var<scalar_type_t<T_Args>>...>::value};
   static constexpr bool is_var_return_{is_var<T_Return>::value};
   static constexpr bool is_var_only_ts_{
       is_var_ts_ && !(is_var_t0_ || is_var_y0_t0_ || is_any_var_args_)};
 
-  std::tuple<arena_t<T_Args>...> local_args_tuple_;
-  std::tuple<arena_t<
-      promote_scalar_t<partials_type_t<scalar_type_t<T_Args>>, T_Args>>...>
-      value_of_args_tuple_;
-
   arena_t<std::vector<Eigen::VectorXd>> y_;
   arena_t<std::vector<T_ts>> ts_;
   arena_t<Eigen::Matrix<T_y0_t0, Eigen::Dynamic, 1>> y0_;
@@ -82,7 +78,9 @@ class cvodes_integrator_adjoint_vari : public vari_base {
    * vari class).
    */
   struct cvodes_solver : public chainable_alloc {
-    std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> y_return_;
+    const std::string function_name_str_;
+    const std::decay_t<F> f_;
+    const size_t N_;
     N_Vector nv_state_forward_;
     N_Vector nv_state_backward_;
     N_Vector nv_quad_;
@@ -92,19 +90,23 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     SUNLinearSolver LS_forward_;
     SUNMatrix A_backward_;
     SUNLinearSolver LS_backward_;
-    const size_t N_;
-    const std::string function_name_str_;
     void* cvodes_mem_;
-    const std::decay_t<F> f_;
+    std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> y_return_;
+    std::tuple<T_Args...> local_args_tuple_;
+    const std::tuple<
+        promote_scalar_t<partials_type_t<scalar_type_t<T_Args>>, T_Args>...>
+        value_of_args_tuple_;
 
     template <typename FF, typename StateFwd, typename StateBwd, typename Quad,
               typename AbsTolFwd, typename AbsTolBwd>
     cvodes_solver(const char* function_name, FF&& f, size_t N,
                   size_t num_args_vars, size_t ts_size, int solver_forward,
                   StateFwd& state_forward, StateBwd& state_backward, Quad& quad,
                   AbsTolFwd& absolute_tolerance_forward,
-                  AbsTolBwd& absolute_tolerance_backward)
+                  AbsTolBwd& absolute_tolerance_backward, const T_Args&... args)
         : chainable_alloc(),
+          f_(std::forward<FF>(f)),
+          function_name_str_(function_name),
           y_return_(ts_size),
           nv_state_forward_(N_VMake_Serial(N, state_forward.data())),
           nv_state_backward_(N_VMake_Serial(N, state_backward.data())),
@@ -122,9 +124,9 @@ class cvodes_integrator_adjoint_vari : public vari_base {
               N == 0 ? nullptr
                      : SUNDenseLinearSolver(nv_state_backward_, A_backward_)),
           N_(N),
-          function_name_str_(function_name),
           cvodes_mem_(CVodeCreate(solver_forward)),
-          f_(std::forward<FF>(f)) {
+          local_args_tuple_(deep_copy_vars(args)...),
+          value_of_args_tuple_(value_of(args)...) {
       if (cvodes_mem_ == nullptr) {
         throw std::runtime_error("CVodeCreate failed to allocate memory");
       }
@@ -202,8 +204,8 @@ class cvodes_integrator_adjoint_vari : public vari_base {
       int interpolation_polynomial, int solver_forward, int solver_backward,
       std::ostream* msgs, const T_Args&... args)
       : vari_base(),
-        local_args_tuple_(to_arena(deep_copy_vars(args))...),
-        value_of_args_tuple_(to_arena(value_of(args))...),
+        // local_args_tuple_(deep_copy_vars(args)...),
+        // value_of_args_tuple_(value_of(args)...),
         y_(ts.size()),
         ts_(ts.begin(), ts.end()),
         y0_(y0),
@@ -238,12 +240,6 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     check_finite(function_name, "initial time", t0);
     check_finite(function_name, "times", ts);
 
-    stan::math::for_each(
-        [func_name = function_name](auto&& arg) {
-          check_finite(func_name, "ode parameters and data", arg);
-        },
-        local_args_tuple_);
-
     check_nonzero_size(function_name, "times", ts);
     check_nonzero_size(function_name, "initial state", y0);
     check_sorted(function_name, "times", ts);
@@ -268,16 +264,29 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     check_positive(function_name, "num_steps_between_checkpoints",
                    num_steps_between_checkpoints_);
     // for polynomial: 1=CV_HERMITE / 2=CV_POLYNOMIAL
-    check_bounded(function_name, "interpolation_polynomial",
-                  interpolation_polynomial_, 1, 2);
+    if (interpolation_polynomial_ != 1 && interpolation_polynomial_ != 2)
+      invalid_argument(function_name, "interpolation_polynomial",
+                       interpolation_polynomial_, "",
+                       ", must be 1 for Hermite or 2 for polynomial "
+                       "interpolation of ODE solution");
     // 1=Adams=CV_ADAMS, 2=BDF=CV_BDF
-    check_bounded(function_name, "solver_forward", solver_forward_, 1, 2);
-    check_bounded(function_name, "solver_backward", solver_backward_, 1, 2);
+    if (solver_forward_ != 1 && solver_forward_ != 2)
+      invalid_argument(function_name, "solver_forward", solver_forward_, "",
+                       ", must be 1 for Adams or 2 for BDF forward solver");
+    if (solver_backward_ != 1 && solver_backward_ != 2)
+      invalid_argument(function_name, "solver_backward", solver_backward_, "",
+                       ", must be 1 for Adams or 2 for BDF backward solver");
 
     solver_ = new cvodes_solver(
         function_name, f, N_, num_args_vars_, ts_.size(), solver_forward_,
         state_forward_, state_backward_, quad_, absolute_tolerance_forward_,
-        absolute_tolerance_backward_);
+        absolute_tolerance_backward_, args...);
+
+    stan::math::for_each(
+        [func_name = function_name](auto&& arg) {
+          check_finite(func_name, "ode parameters and data", arg);
+        },
+        solver_->local_args_tuple_);
 
     check_flag_sundials(
         CVodeInit(solver_->cvodes_mem_, &cvodes_integrator_adjoint_vari::cv_rhs,
@@ -338,13 +347,12 @@ class cvodes_integrator_adjoint_vari : public vari_base {
           } else {
             check_flag_sundials(error_code, "CVodeF");
           }
-
         } else {
           int error_code
               = CVode(solver_->cvodes_mem_, t_final, solver_->nv_state_forward_,
                       &t_init, CV_NORMAL);
 
-          if (error_code == CV_TOO_MUCH_WORK) {
+          if (unlikely(error_code == CV_TOO_MUCH_WORK)) {
             throw_domain_error(solver_->function_name_str_.c_str(), "", t_final,
                                "Failed to integrate to next output time (",
                                ") in less than max_num_steps steps");
@@ -357,9 +365,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
 
       t_init = t_final;
     }
-    if (is_var_return_) {
-      ChainableStack::instance_->var_stack_.push_back(this);
-    }
+    ChainableStack::instance_->var_stack_.push_back(this);
   }
 
  private:
@@ -414,8 +420,8 @@ class cvodes_integrator_adjoint_vari : public vari_base {
               += forward_as<var>(solver_->y_return_[i].coeff(j)).adj();
         }
 
-        forward_as<var>(ts_[i])->adj_ += step_sens.dot(
-            rhs(value_of(ts_[i]), y_[i], value_of_args_tuple_));
+        adjoint_of(ts_[i]) += step_sens.dot(
+            rhs(value_of(ts_[i]), y_[i], solver_->value_of_args_tuple_));
         step_sens.setZero();
       }
 
@@ -543,8 +549,8 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     }
 
     if (is_var_t0_) {
-      forward_as<var>(t0_)->adj_ += -state_backward_.dot(
-          rhs(t_init, value_of(y0_), value_of_args_tuple_));
+      adjoint_of(t0_) += -state_backward_.dot(
+          rhs(t_init, value_of(y0_), solver_->value_of_args_tuple_));
     }
 
     // After integrating all the way back to t0, we finally have the
@@ -589,7 +595,8 @@ class cvodes_integrator_adjoint_vari : public vari_base {
    */
   inline int rhs(double t, const double* y, double*& dy_dt) const {
     const Eigen::VectorXd y_vec = Eigen::Map<const Eigen::VectorXd>(y, N_);
-    const Eigen::VectorXd dy_dt_vec = rhs(t, y_vec, value_of_args_tuple_);
+    const Eigen::VectorXd dy_dt_vec
+        = rhs(t, y_vec, solver_->value_of_args_tuple_);
     check_size_match(solver_->function_name_str_.c_str(), "dy_dt",
                      dy_dt_vec.size(), "states", N_);
     Eigen::Map<Eigen::VectorXd>(dy_dt, N_) = dy_dt_vec;
@@ -624,7 +631,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     Eigen::Matrix<var, Eigen::Dynamic, 1> y_vars(
         Eigen::Map<const Eigen::VectorXd>(NV_DATA_S(y), N_));
     Eigen::Matrix<var, Eigen::Dynamic, 1> f_y_t_vars
-        = rhs(t, y_vars, value_of_args_tuple_);
+        = rhs(t, y_vars, solver_->value_of_args_tuple_);
     check_size_match(solver_->function_name_str_.c_str(), "dy_dt",
                      f_y_t_vars.size(), "states", N_);
     f_y_t_vars.adj() = -Eigen::Map<Eigen::VectorXd>(NV_DATA_S(yB), N_);
@@ -662,9 +669,9 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     // The vars here do not live on the nested stack so must be zero'd
     // separately
     stan::math::for_each([](auto&& arg) { zero_adjoints(arg); },
-                         local_args_tuple_);
+                         solver_->local_args_tuple_);
     Eigen::Matrix<var, Eigen::Dynamic, 1> f_y_t_vars
-        = rhs(t, y_vec, local_args_tuple_);
+        = rhs(t, y_vec, solver_->local_args_tuple_);
     check_size_match(solver_->function_name_str_.c_str(), "dy_dt",
                      f_y_t_vars.size(), "states", N_);
     f_y_t_vars.adj() = -Eigen::Map<Eigen::VectorXd>(NV_DATA_S(yB), N_);
@@ -673,7 +680,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
         [&qBdot](auto&&... args) {
           accumulate_adjoints(NV_DATA_S(qBdot), args...);
         },
-        local_args_tuple_);
+        solver_->local_args_tuple_);
     return 0;
   }
 
@@ -698,7 +705,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
     Eigen::Matrix<var, Eigen::Dynamic, 1> y_var(
         Eigen::Map<const Eigen::VectorXd>(NV_DATA_S(y), N_));
     Eigen::Matrix<var, Eigen::Dynamic, 1> fy_var
-        = rhs(t, y_var, value_of_args_tuple_);
+        = rhs(t, y_var, solver_->value_of_args_tuple_);
 
     check_size_match(solver_->function_name_str_.c_str(), "dy_dt",
                      fy_var.size(), "states", N_);

stan/math/rev/functor/ode_adjoint.hpp

@@ -66,7 +66,8 @@ namespace math {
 template <typename F, typename T_y0, typename T_t0, typename T_ts,
           typename T_abs_tol_fwd, typename T_abs_tol_bwd, typename... T_Args,
           require_all_eigen_col_vector_t<T_y0, T_abs_tol_fwd,
-                                         T_abs_tol_bwd>* = nullptr>
+                                         T_abs_tol_bwd>* = nullptr,
+          require_any_not_st_arithmetic<T_y0, T_t0, T_ts, T_Args...>* = nullptr>
 auto ode_adjoint_impl(
     const char* function_name, F&& f, const T_y0& y0, const T_t0& t0,
     const std::vector<T_ts>& ts, double relative_tolerance_forward,
@@ -91,6 +92,98 @@ auto ode_adjoint_impl(
   return integrator->solution();
 }
 
+/**
+ * Solve the ODE initial value problem y' = f(t, y), y(t0) = y0 at a set of
+ * times, { t1, t2, t3, ... } using the stiff backward differentiation formula
+ * BDF solver or the non-stiff Adams solver from CVODES. The ODE system is
+ * integrated using the adjoint sensitivity approach of CVODES. This
+ * implementation handles the case of a double return type which ensures that no
+ * resources are left on the AD stack.
+ *
+ * \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_y0 Type of initial state
+ * @tparam T_t0 Type of initial time
+ * @tparam T_ts Type of output times
+ * @tparam T_Args Types of pass-through parameters
+ *
+ * @param function_name Calling function name (for printing debugging messages)
+ * @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_forward Relative tolerance for forward problem
+ * passed to CVODES
+ * @param absolute_tolerance_forward Absolute tolerance per ODE state for
+ * forward problem passed to CVODES
+ * @param relative_tolerance_backward Relative tolerance for backward problem
+ * passed to CVODES
+ * @param absolute_tolerance_backward Absolute tolerance per ODE state for
+ * backward problem passed to CVODES
+ * @param relative_tolerance_quadrature Relative tolerance for quadrature
+ * problem passed to CVODES
+ * @param absolute_tolerance_quadrature Absolute tolerance for quadrature
+ * problem passed to CVODES
+ * @param max_num_steps Upper limit on the number of integration steps to
+ *   take between each output (error if exceeded)
+ * @param num_steps_between_checkpoints Number of integrator steps after which a
+ * checkpoint is stored for the backward pass
+ * @param interpolation_polynomial type of polynomial used for interpolation
+ * @param solver_forward solver used for forward pass
+ * @param solver_backward solver used for backward pass
+ * @param[in, out] msgs the print stream for warning messages
+ * @param args Extra arguments passed unmodified through to ODE right hand side
+ * @return An `std::vector` of Eigen column vectors with scalars equal to
+ *  the least upper bound of `T_y0`, `T_t0`, `T_ts`, and the lambda's arguments.
+ *  This represents the solution to ODE at times \p ts
+ */
+template <typename F, typename T_y0, typename T_t0, typename T_ts,
+          typename T_abs_tol_fwd, typename T_abs_tol_bwd, typename... T_Args,
+          require_all_eigen_col_vector_t<T_y0, T_abs_tol_fwd,
+                                         T_abs_tol_bwd>* = nullptr,
+          require_all_st_arithmetic<T_y0, T_t0, T_ts, T_Args...>* = nullptr>
+std::vector<Eigen::VectorXd> ode_adjoint_impl(
+    const char* function_name, F&& f, const T_y0& y0, const T_t0& t0,
+    const std::vector<T_ts>& ts, double relative_tolerance_forward,
+    const T_abs_tol_fwd& absolute_tolerance_forward,
+    double relative_tolerance_backward,
+    const T_abs_tol_bwd& absolute_tolerance_backward,
+    double relative_tolerance_quadrature, double absolute_tolerance_quadrature,
+    long int max_num_steps,                  // NOLINT(runtime/int)
+    long int num_steps_between_checkpoints,  // NOLINT(runtime/int)
+    int interpolation_polynomial, int solver_forward, int solver_backward,
+    std::ostream* msgs, const T_Args&... args) {
+  std::vector<Eigen::VectorXd> ode_solution;
+  {
+    nested_rev_autodiff nested;
+
+    using integrator_vari
+        = cvodes_integrator_adjoint_vari<F, plain_type_t<T_y0>, T_t0, T_ts,
+                                         plain_type_t<T_Args>...>;
+
+    auto integrator = new integrator_vari(
+        function_name, std::forward<F>(f), y0, t0, ts,
+        relative_tolerance_forward, absolute_tolerance_forward,
+        relative_tolerance_backward, absolute_tolerance_backward,
+        relative_tolerance_quadrature, absolute_tolerance_quadrature,
+        max_num_steps, num_steps_between_checkpoints, interpolation_polynomial,
+        solver_forward, solver_backward, msgs, args...);
+
+    ode_solution = integrator->solution();
+  }
+  return ode_solution;
+}
+
 /**
  * Solve the ODE initial value problem y' = f(t, y), y(t0) = y0 at a set of
  * times, { t1, t2, t3, ... } using the stiff backward differentiation formula