cvodes_integrator.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUNCTOR_INTEGRATE_ODE_CVODES_HPP
#define STAN_MATH_REV_FUNCTOR_INTEGRATE_ODE_CVODES_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/functor/cvodes_utils.hpp>
#include <stan/math/rev/functor/coupled_ode_system.hpp>
#include <stan/math/rev/functor/ode_store_sensitivities.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/functor/apply.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <sundials/sundials_context.h>
#include <cvodes/cvodes.h>
#include <nvector/nvector_serial.h>
#include <sunlinsol/sunlinsol_dense.h>
#include <algorithm>
#include <ostream>
#include <vector>
 
namespace stan {
namespace math {
 
template <int Lmm, typename F, typename T_y0, typename T_t0, typename T_ts,
          typename... T_Args>
class cvodes_integrator {
  using T_Return = return_type_t<T_y0, T_t0, T_ts, T_Args...>;
  using T_y0_t0 = return_type_t<T_y0, T_t0>;
 
  const char* function_name_;
  sundials::Context sundials_context_;
  const F& f_;
  const Eigen::Matrix<T_y0_t0, Eigen::Dynamic, 1> y0_;
  const T_t0 t0_;
  const std::vector<T_ts>& ts_;
  std::tuple<const T_Args&...> args_tuple_;
  std::tuple<plain_type_t<decltype(value_of(std::declval<const T_Args&>()))>...>
      value_of_args_tuple_;
  const size_t N_;
  std::ostream* msgs_;
  double relative_tolerance_;
  double absolute_tolerance_;
  long int max_num_steps_;  // NOLINT(runtime/int)
 
  const size_t num_y0_vars_;
  const size_t num_args_vars_;
 
  coupled_ode_system<F, T_y0_t0, T_Args...> coupled_ode_;
 
  std::vector<double> coupled_state_;
  N_Vector nv_state_;
  N_Vector* nv_state_sens_;
  SUNMatrix A_;
  SUNLinearSolver LS_;
 
  static int cv_rhs(realtype t, N_Vector y, N_Vector ydot, void* user_data) {
    cvodes_integrator* integrator = static_cast<cvodes_integrator*>(user_data);
    integrator->rhs(t, NV_DATA_S(y), NV_DATA_S(ydot));
    return 0;
  }
 
  static int cv_rhs_sens(int Ns, realtype t, N_Vector y, N_Vector ydot,
                         N_Vector* yS, N_Vector* ySdot, void* user_data,
                         N_Vector tmp1, N_Vector tmp2) {
    cvodes_integrator* integrator = static_cast<cvodes_integrator*>(user_data);
    integrator->rhs_sens(t, NV_DATA_S(y), yS, ySdot);
    return 0;
  }
 
  static int cv_jacobian_states(realtype t, N_Vector y, N_Vector fy,
                                SUNMatrix J, void* user_data, N_Vector tmp1,
                                N_Vector tmp2, N_Vector tmp3) {
    cvodes_integrator* integrator = static_cast<cvodes_integrator*>(user_data);
    integrator->jacobian_states(t, NV_DATA_S(y), J);
    return 0;
  }
 
  inline void rhs(double t, const double y[], double dy_dt[]) const {
    const Eigen::VectorXd y_vec = Eigen::Map<const Eigen::VectorXd>(y, N_);
 
    Eigen::VectorXd dy_dt_vec = math::apply(
        [&](auto&&... args) { return f_(t, y_vec, msgs_, args...); },
        value_of_args_tuple_);
 
    check_size_match("cvodes_integrator", "dy_dt", dy_dt_vec.size(), "states",
                     N_);
 
    std::copy(dy_dt_vec.data(), dy_dt_vec.data() + dy_dt_vec.size(), dy_dt);
  }
 
  inline void jacobian_states(double t, const double y[], SUNMatrix J) const {
    Eigen::VectorXd fy;
    Eigen::MatrixXd Jfy;
 
    auto f_wrapped = [&](const Eigen::Matrix<var, Eigen::Dynamic, 1>& y) {
      return math::apply(
          [&](auto&&... args) { return f_(t, y, msgs_, args...); },
          value_of_args_tuple_);
    };
 
    jacobian(f_wrapped, Eigen::Map<const Eigen::VectorXd>(y, N_), fy, Jfy);
 
    for (size_t j = 0; j < Jfy.cols(); ++j) {
      for (size_t i = 0; i < Jfy.rows(); ++i) {
        SM_ELEMENT_D(J, i, j) = Jfy(i, j);
      }
    }
  }
 
  inline void rhs_sens(double t, const double y[], N_Vector* yS,
                       N_Vector* ySdot) {
    std::vector<double> z(coupled_state_.size());
    std::vector<double> dz_dt;
    std::copy(y, y + N_, z.data());
    for (std::size_t s = 0; s < num_y0_vars_ + num_args_vars_; s++) {
      std::copy(NV_DATA_S(yS[s]), NV_DATA_S(yS[s]) + N_,
                z.data() + (s + 1) * N_);
    }
    coupled_ode_(z, dz_dt, t);
    for (std::size_t s = 0; s < num_y0_vars_ + num_args_vars_; s++) {
      std::move(dz_dt.data() + (s + 1) * N_, dz_dt.data() + (s + 2) * N_,
                NV_DATA_S(ySdot[s]));
    }
  }
 
 public:
  template <require_eigen_col_vector_t<T_y0>* = nullptr>
  cvodes_integrator(const char* function_name, const F& f, const T_y0& y0,
                    const T_t0& t0, const std::vector<T_ts>& ts,
                    double relative_tolerance, double absolute_tolerance,
                    long int max_num_steps,  // NOLINT(runtime/int)
                    std::ostream* msgs, const T_Args&... args)
      : function_name_(function_name),
        sundials_context_(),
        f_(f),
        y0_(y0.template cast<T_y0_t0>()),
        t0_(t0),
        ts_(ts),
        args_tuple_(args...),
        value_of_args_tuple_(value_of(args)...),
        N_(y0.size()),
        msgs_(msgs),
        relative_tolerance_(relative_tolerance),
        absolute_tolerance_(absolute_tolerance),
        max_num_steps_(max_num_steps),
        num_y0_vars_(count_vars(y0_)),
        num_args_vars_(count_vars(args...)),
        coupled_ode_(f, y0_, msgs, args...),
        coupled_state_(coupled_ode_.initial_state()) {
    check_finite(function_name, "initial state", y0_);
    check_finite(function_name, "initial time", t0_);
    check_finite(function_name, "times", ts_);
 
    // Code from: https://stackoverflow.com/a/17340003 . Should probably do
    // something better
    math::apply(
        [&](auto&&... args) {
          std::vector<int> unused_temp{
              0, (check_finite(function_name, "ode parameters and data", args),
                  0)...};
        },
        args_tuple_);
 
    check_nonzero_size(function_name, "times", ts_);
    check_nonzero_size(function_name, "initial state", y0_);
    check_sorted(function_name, "times", ts_);
    check_less(function_name, "initial time", t0_, ts_[0]);
    check_positive_finite(function_name, "relative_tolerance",
                          relative_tolerance_);
    check_positive_finite(function_name, "absolute_tolerance",
                          absolute_tolerance_);
    check_positive(function_name, "max_num_steps", max_num_steps_);
 
    nv_state_ = N_VMake_Serial(N_, &coupled_state_[0], sundials_context_);
    nv_state_sens_ = nullptr;
    A_ = SUNDenseMatrix(N_, N_, sundials_context_);
    LS_ = SUNLinSol_Dense(nv_state_, A_, sundials_context_);
 
    if (num_y0_vars_ + num_args_vars_ > 0) {
      nv_state_sens_
          = N_VCloneEmptyVectorArray(num_y0_vars_ + num_args_vars_, nv_state_);
      for (std::size_t i = 0; i < num_y0_vars_ + num_args_vars_; i++) {
        NV_DATA_S(nv_state_sens_[i]) = &coupled_state_[N_] + i * N_;
      }
    }
  }
 
  ~cvodes_integrator() {
    SUNLinSolFree(LS_);
    SUNMatDestroy(A_);
    N_VDestroy_Serial(nv_state_);
    if (num_y0_vars_ + num_args_vars_ > 0) {
      N_VDestroyVectorArray(nv_state_sens_, num_y0_vars_ + num_args_vars_);
    }
  }
 
  std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> operator()() {
    std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> y;
 
    void* cvodes_mem = CVodeCreate(Lmm, sundials_context_);
    if (cvodes_mem == nullptr) {
      throw std::runtime_error("CVodeCreate failed to allocate memory");
    }
 
    try {
      CHECK_CVODES_CALL(CVodeInit(cvodes_mem, &cvodes_integrator::cv_rhs,
                                  value_of(t0_), nv_state_));
 
      // Assign pointer to this as user data
      CHECK_CVODES_CALL(
          CVodeSetUserData(cvodes_mem, reinterpret_cast<void*>(this)));
 
      cvodes_set_options(cvodes_mem, max_num_steps_);
 
      CHECK_CVODES_CALL(CVodeSStolerances(cvodes_mem, relative_tolerance_,
                                          absolute_tolerance_));
 
      CHECK_CVODES_CALL(CVodeSetLinearSolver(cvodes_mem, LS_, A_));
      CHECK_CVODES_CALL(
          CVodeSetJacFn(cvodes_mem, &cvodes_integrator::cv_jacobian_states));
 
      // initialize forward sensitivity system of CVODES as needed
      if (num_y0_vars_ + num_args_vars_ > 0) {
        CHECK_CVODES_CALL(CVodeSensInit(
            cvodes_mem, static_cast<int>(num_y0_vars_ + num_args_vars_),
            CV_STAGGERED, &cvodes_integrator::cv_rhs_sens, nv_state_sens_));
 
        CHECK_CVODES_CALL(CVodeSetSensErrCon(cvodes_mem, SUNTRUE));
 
        CHECK_CVODES_CALL(CVodeSensEEtolerances(cvodes_mem));
      }
 
      double t_init = value_of(t0_);
      for (size_t n = 0; n < ts_.size(); ++n) {
        double t_final = value_of(ts_[n]);
 
        if (t_final != t_init) {
          CHECK_CVODES_CALL(
              CVode(cvodes_mem, t_final, nv_state_, &t_init, CV_NORMAL));
 
          if (num_y0_vars_ + num_args_vars_ > 0) {
            CHECK_CVODES_CALL(
                CVodeGetSens(cvodes_mem, &t_init, nv_state_sens_));
          }
        }
 
        y.emplace_back(math::apply(
            [&](auto&&... args) {
              return ode_store_sensitivities(f_, coupled_state_, y0_, t0_,
                                             ts_[n], msgs_, args...);
            },
            args_tuple_));
 
        t_init = t_final;
      }
    } catch (const std::exception& e) {
      CVodeFree(&cvodes_mem);
      throw;
    }
 
    CVodeFree(&cvodes_mem);
 
    return y;
  }
};  // cvodes integrator
 
}  // namespace math
}  // namespace stan
#endif