coupled_ode_system.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUNCTOR_COUPLED_ODE_SYSTEM_HPP
#define STAN_MATH_PRIM_FUNCTOR_COUPLED_ODE_SYSTEM_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/size.hpp>
#include <stan/math/prim/functor/apply.hpp>
#include <ostream>
#include <vector>
 
namespace stan {
namespace math {
 
template <bool ArithmeticArguments, typename F, typename T_y0, typename... Args>
struct coupled_ode_system_impl;
 
template <typename F, typename T_y0, typename... Args>
struct coupled_ode_system_impl<true, F, T_y0, Args...> {
  const F& f_;
  const Eigen::VectorXd& y0_;
  std::tuple<const Args&...> args_tuple_;
  const size_t N_;
  std::ostream* msgs_;
 
  coupled_ode_system_impl(const F& f, const Eigen::VectorXd& y0,
                          std::ostream* msgs, const Args&... args)
      : f_(f), y0_(y0), args_tuple_(args...), N_(y0.size()), msgs_(msgs) {}
 
  void operator()(const std::vector<double>& z, std::vector<double>& dz_dt,
                  double t) const {
    Eigen::VectorXd y(z.size());
    for (size_t n = 0; n < z.size(); ++n)
      y.coeffRef(n) = z[n];
 
    dz_dt.resize(y.size());
 
    Eigen::VectorXd f_y_t = math::apply(
        [&](const Args&... args) { return f_(t, y, msgs_, args...); },
        args_tuple_);
 
    check_size_match("coupled_ode_system", "dy_dt", f_y_t.size(), "states",
                     y.size());
 
    Eigen::Map<Eigen::VectorXd>(dz_dt.data(), dz_dt.size()) = f_y_t;
  }
 
  size_t size() const { return N_; }
 
  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));
    }
 
    return initial;
  }
};
 
template <typename F, typename T_y0, typename... Args>
struct coupled_ode_system
    : public coupled_ode_system_impl<
          std::is_arithmetic<return_type_t<T_y0, Args...>>::value, F, T_y0,
          Args...> {
  coupled_ode_system(const F& f,
                     const Eigen::Matrix<T_y0, Eigen::Dynamic, 1>& y0,
                     std::ostream* msgs, const Args&... args)
      : coupled_ode_system_impl<
          std::is_arithmetic<return_type_t<T_y0, Args...>>::value, F, T_y0,
          Args...>(f, y0, msgs, args...) {}
};
 
}  // namespace math
}  // namespace stan
#endif