Automatic Differentiation
 
Loading...
Searching...
No Matches
kinsol_data.hpp
Go to the documentation of this file.
1#ifndef STAN_MATH_REV_FUNCTOR_KINSOL_DATA_HPP
2#define STAN_MATH_REV_FUNCTOR_KINSOL_DATA_HPP
3
9#include <sundials/sundials_context.h>
10#include <kinsol/kinsol.h>
11#include <sunmatrix/sunmatrix_dense.h>
12#include <sunlinsol/sunlinsol_dense.h>
13#include <nvector/nvector_serial.h>
14#include <vector>
15#include <tuple>
16
17namespace stan {
18namespace math {
19
28template <typename F1, typename... Args>
30 const F1& f_;
31 const Eigen::VectorXd& x_;
32 const size_t N_;
33 std::ostream* const msgs_;
34 const std::tuple<const Args&...> args_tuple_;
35
36 typedef kinsol_system_data<F1, Args...> system_data;
37
38 public:
39 sundials::Context sundials_context_;
40 N_Vector nv_x_;
41 SUNMatrix J_;
42 SUNLinearSolver LS_;
44
45 /* Constructor */
46 kinsol_system_data(const F1& f, const Eigen::VectorXd& x,
47 std::ostream* const msgs, const Args&... args)
48 : f_(f),
49 x_(x),
50 N_(x.size()),
51 msgs_(msgs),
52 args_tuple_(args...),
54 nv_x_(N_VMake_Serial(N_, &to_array_1d(x_)[0], sundials_context_)),
55 J_(SUNDenseMatrix(N_, N_, sundials_context_)),
56 LS_(SUNLinSol_Dense(nv_x_, J_, sundials_context_)),
58
60 N_VDestroy_Serial(nv_x_);
61 SUNLinSolFree(LS_);
62 SUNMatDestroy(J_);
63 KINFree(&kinsol_memory_);
64 }
65
66 /* Implements the user-defined function passed to KINSOL. */
67 static int kinsol_f_system(const N_Vector x, const N_Vector f_eval,
68 void* const user_data) {
69 const system_data* explicit_system
70 = static_cast<const system_data*>(user_data);
71
72 Eigen::VectorXd x_eigen(
73 Eigen::Map<Eigen::VectorXd>(NV_DATA_S(x), explicit_system->N_));
74
75 Eigen::Map<Eigen::VectorXd> f_eval_map(N_VGetArrayPointer(f_eval),
76 explicit_system->N_);
77 auto result = math::apply(
78 [&](const auto&... args) {
79 return explicit_system->f_(x_eigen, explicit_system->msgs_, args...);
80 },
81 explicit_system->args_tuple_);
82 check_matching_sizes("", "the algebraic system's output", result,
83 "the vector of unknowns, x,", f_eval_map);
84 f_eval_map = result;
85 return 0;
86 }
87
100 static int kinsol_jacobian(const N_Vector x, const N_Vector f_eval,
101 const SUNMatrix J, void* const user_data,
102 const N_Vector tmp1, const N_Vector tmp2) {
103 const system_data* explicit_system
104 = static_cast<const system_data*>(user_data);
105
106 Eigen::VectorXd x_eigen(
107 Eigen::Map<Eigen::VectorXd>(NV_DATA_S(x), explicit_system->N_));
108 Eigen::Map<Eigen::VectorXd> f_eval_map(N_VGetArrayPointer(f_eval),
109 explicit_system->N_);
110
111 auto f_wrt_x = [&](const auto& x) {
112 return math::apply(
113 [&](const auto&... args) {
114 return explicit_system->f_(x, explicit_system->msgs_, args...);
115 },
116 explicit_system->args_tuple_);
117 };
118
119 Eigen::MatrixXd Jf_x;
120 Eigen::VectorXd f_x;
121
122 jacobian(f_wrt_x, x_eigen, f_x, Jf_x);
123
124 f_eval_map = f_x;
125
126 for (int j = 0; j < Jf_x.cols(); j++)
127 for (int i = 0; i < Jf_x.rows(); i++)
128 SM_ELEMENT_D(J, i, j) = Jf_x(i, j);
129
130 return 0;
131 }
132};
133
134} // namespace math
135} // namespace stan
136#endif
static int kinsol_jacobian(const N_Vector x, const N_Vector f_eval, const SUNMatrix J, void *const user_data, const N_Vector tmp1, const N_Vector tmp2)
Implements the function of type CVDlsJacFn which is the user-defined callbacks for KINSOL to calculat...
kinsol_system_data(const F1 &f, const Eigen::VectorXd &x, std::ostream *const msgs, const Args &... args)
sundials::Context sundials_context_
kinsol_system_data< F1, Args... > system_data
static int kinsol_f_system(const N_Vector x, const N_Vector f_eval, void *const user_data)
const Eigen::VectorXd & x_
const std::tuple< const Args &... > args_tuple_
KINSOL algebraic system data holder.
auto to_array_1d(T_x &&x)
Returns input matrix reshaped into a vector.
int64_t size(const T &m)
Returns the size (number of the elements) of a matrix_cl or var_value<matrix_cl<T>>.
Definition size.hpp:19
void jacobian(const F &f, const Eigen::Matrix< T, Eigen::Dynamic, 1 > &x, Eigen::Matrix< T, Eigen::Dynamic, 1 > &fx, Eigen::Matrix< T, Eigen::Dynamic, Eigen::Dynamic > &J)
Definition jacobian.hpp:11
void check_matching_sizes(const char *function, const char *name1, const T_y1 &y1, const char *name2, const T_y2 &y2)
Check if two structures at the same size.
constexpr decltype(auto) apply(F &&f, Tuple &&t, PreArgs &&... pre_args)
Definition apply.hpp:52
The lgamma implementation in stan-math is based on either the reentrant safe lgamma_r implementation ...