coupled_ode_system.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUNCTOR_COUPLED_ODE_SYSTEM_HPP
#define STAN_MATH_REV_FUNCTOR_COUPLED_ODE_SYSTEM_HPP
 
#include <stan/math/prim/functor/coupled_ode_system.hpp>
#include <stan/math/rev/functor/cvodes_utils.hpp>
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/prim/functor/apply.hpp>
#include <stan/math/prim/functor/for_each.hpp>
#include <stan/math/prim/err.hpp>
#include <stdexcept>
#include <ostream>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename F, typename T_y0, typename... Args>
struct coupled_ode_system_impl<false, F, T_y0, Args...> {
  const F& f_;
  const Eigen::Matrix<T_y0, Eigen::Dynamic, 1>& y0_;
  std::tuple<decltype(deep_copy_vars(std::declval<const Args&>()))...>
      local_args_tuple_;
  const size_t num_y0_vars_;
  const size_t num_args_vars;
  const size_t N_;
  Eigen::VectorXd args_adjoints_;
  Eigen::VectorXd y_adjoints_;
  std::ostream* msgs_;
 
  coupled_ode_system_impl(const F& f,
                          const Eigen::Matrix<T_y0, Eigen::Dynamic, 1>& y0,
                          std::ostream* msgs, const Args&... args)
      : f_(f),
        y0_(y0),
        local_args_tuple_(deep_copy_vars(args)...),
        num_y0_vars_(count_vars(y0_)),
        num_args_vars(count_vars(args...)),
        N_(y0.size()),
        args_adjoints_(num_args_vars),
        y_adjoints_(N_),
        msgs_(msgs) {}
 
  void operator()(const std::vector<double>& z, std::vector<double>& dz_dt,
                  double t) {
    using std::vector;
 
    dz_dt.resize(size());
 
    // Run nested autodiff in this scope
    nested_rev_autodiff nested;
 
    Eigen::Matrix<var, Eigen::Dynamic, 1> y_vars(N_);
    for (size_t n = 0; n < N_; ++n)
      y_vars.coeffRef(n) = z[n];
 
    Eigen::Matrix<var, Eigen::Dynamic, 1> f_y_t_vars = math::apply(
        [&](auto&&... args) { return f_(t, y_vars, msgs_, args...); },
        local_args_tuple_);
 
    check_size_match("coupled_ode_system", "dy_dt", f_y_t_vars.size(), "states",
                     N_);
 
    for (size_t i = 0; i < N_; ++i) {
      dz_dt[i] = f_y_t_vars.coeffRef(i).val();
      f_y_t_vars.coeffRef(i).grad();
 
      y_adjoints_ = y_vars.adj();
 
      if (args_adjoints_.size() > 0) {
        memset(args_adjoints_.data(), 0,
               sizeof(double) * args_adjoints_.size());
      }
 
      math::apply(
          [&](auto&&... args) {
            accumulate_adjoints(args_adjoints_.data(), args...);
          },
          local_args_tuple_);
 
      // 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_);
 
      // No need to zero adjoints after last sweep
      if (i + 1 < N_) {
        nested.set_zero_all_adjoints();
      }
 
      // Compute the right hand side for the sensitivities with respect to the
      // initial conditions
      for (size_t j = 0; j < num_y0_vars_; ++j) {
        double temp_deriv = 0.0;
        for (size_t k = 0; k < N_; ++k) {
          temp_deriv += z[N_ + N_ * j + k] * y_adjoints_.coeffRef(k);
        }
 
        dz_dt[N_ + N_ * j + i] = temp_deriv;
      }
 
      // Compute the right hand size for the sensitivities with respect to the
      // parameters
      for (size_t j = 0; j < num_args_vars; ++j) {
        double temp_deriv = args_adjoints_.coeffRef(j);
        for (size_t k = 0; k < N_; ++k) {
          temp_deriv += z[N_ + N_ * num_y0_vars_ + N_ * j + k]
                        * y_adjoints_.coeffRef(k);
        }
 
        dz_dt[N_ + N_ * num_y0_vars_ + N_ * j + i] = temp_deriv;
      }
    }
  }
 
  size_t size() const { return N_ + N_ * num_y0_vars_ + N_ * num_args_vars; }
 
  std::vector<double> initial_state() const {
    std::vector<double> initial(size(), 0.0);
    for (size_t i = 0; i < N_; i++) {
      initial[i] = value_of(y0_(i));
    }
    for (size_t i = 0; i < num_y0_vars_; i++) {
      initial[N_ + i * N_ + i] = 1.0;
    }
    return initial;
  }
};
 
}  // namespace math
}  // namespace stan
#endif