Skip to content

Commit 30633be

Browse files
wds15wds15
committed
align Sundials code with 6.0 conventions
1 parent c9c4bc9 commit 30633be

8 files changed

Lines changed: 64 additions & 44 deletions

‎stan/math/rev/functor/algebra_solver_fp.hpp‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <stan/math/prim/fun/to_array_1d.hpp>
1212
#include <stan/math/prim/fun/to_vector.hpp>
1313
#include <stan/math/prim/fun/value_of.hpp>
14+
#include <sundials/sundials_context.h>
1415
#include <kinsol/kinsol.h>
1516
#include <sunmatrix/sunmatrix_dense.h>
1617
#include <sunlinsol/sunlinsol_dense.h>
@@ -32,6 +33,8 @@ namespace math {
3233
*/
3334
template <typename F>
3435
struct KinsolFixedPointEnv {
36+
/** Sundials context **/
37+
sundials::Context sundials_context_;
3538
/** RHS functor. */
3639
const F& f_;
3740
/** val of params for @c y_ to refer to when
@@ -65,18 +68,19 @@ struct KinsolFixedPointEnv {
6568
const std::vector<int>& x_i, std::ostream* msgs,
6669
const std::vector<T_u>& u_scale,
6770
const std::vector<T_f>& f_scale)
68-
: f_(f),
71+
: sundials_context_(),
72+
f_(f),
6973
y_dummy(),
7074
y_(y),
7175
N_(x.size()),
7276
M_(y.size()),
7377
x_r_(x_r),
7478
x_i_(x_i),
7579
msgs_(msgs),
76-
mem_(KINCreate()),
77-
nv_x_(N_VNew_Serial(N_)),
78-
nv_u_scal_(N_VNew_Serial(N_)),
79-
nv_f_scal_(N_VNew_Serial(N_)) {
80+
mem_(KINCreate(sundials_context_)),
81+
nv_x_(N_VNew_Serial(N_, sundials_context_)),
82+
nv_u_scal_(N_VNew_Serial(N_, sundials_context_)),
83+
nv_f_scal_(N_VNew_Serial(N_, sundials_context_)) {
8084
for (int i = 0; i < N_; ++i) {
8185
NV_Ith_S(nv_x_, i) = stan::math::value_of(x(i));
8286
NV_Ith_S(nv_u_scal_, i) = stan::math::value_of(u_scale[i]);
@@ -92,18 +96,19 @@ struct KinsolFixedPointEnv {
9296
const std::vector<int>& x_i, std::ostream* msgs,
9397
const std::vector<T_u>& u_scale,
9498
const std::vector<T_f>& f_scale)
95-
: f_(f),
99+
: sundials_context_(),
100+
f_(f),
96101
y_dummy(stan::math::value_of(y)),
97102
y_(y_dummy),
98103
N_(x.size()),
99104
M_(y.size()),
100105
x_r_(x_r),
101106
x_i_(x_i),
102107
msgs_(msgs),
103-
mem_(KINCreate()),
104-
nv_x_(N_VNew_Serial(N_)),
105-
nv_u_scal_(N_VNew_Serial(N_)),
106-
nv_f_scal_(N_VNew_Serial(N_)) {
108+
mem_(KINCreate(sundials_context_)),
109+
nv_x_(N_VNew_Serial(N_, sundials_context_)),
110+
nv_u_scal_(N_VNew_Serial(N_, sundials_context_)),
111+
nv_f_scal_(N_VNew_Serial(N_, sundials_context_)) {
107112
for (int i = 0; i < N_; ++i) {
108113
NV_Ith_S(nv_x_, i) = stan::math::value_of(x(i));
109114
NV_Ith_S(nv_u_scal_, i) = stan::math::value_of(u_scale[i]);

‎stan/math/rev/functor/cvodes_integrator.hpp‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include <stan/math/rev/functor/ode_store_sensitivities.hpp>
88
#include <stan/math/prim/err.hpp>
99
#include <stan/math/prim/fun/value_of.hpp>
10+
#include <sundials/sundials_context.h>
1011
#include <cvodes/cvodes.h>
1112
#include <nvector/nvector_serial.h>
1213
#include <sunlinsol/sunlinsol_dense.h>
@@ -35,6 +36,7 @@ class cvodes_integrator {
3536
using T_y0_t0 = return_type_t<T_y0, T_t0>;
3637

3738
const char* function_name_;
39+
sundials::Context sundials_context_;
3840
const F& f_;
3941
const Eigen::Matrix<T_y0_t0, Eigen::Dynamic, 1> y0_;
4042
const T_t0 t0_;
@@ -190,6 +192,7 @@ class cvodes_integrator {
190192
long int max_num_steps, // NOLINT(runtime/int)
191193
std::ostream* msgs, const T_Args&... args)
192194
: function_name_(function_name),
195+
sundials_context_(),
193196
f_(f),
194197
y0_(y0.template cast<T_y0_t0>()),
195198
t0_(t0),
@@ -229,13 +232,13 @@ class cvodes_integrator {
229232
absolute_tolerance_);
230233
check_positive(function_name, "max_num_steps", max_num_steps_);
231234

232-
nv_state_ = N_VMake_Serial(N_, &coupled_state_[0]);
235+
nv_state_ = N_VMake_Serial(N_, &coupled_state_[0], sundials_context_);
233236
nv_state_sens_ = nullptr;
234-
A_ = SUNDenseMatrix(N_, N_);
235-
LS_ = SUNDenseLinearSolver(nv_state_, A_);
237+
A_ = SUNDenseMatrix(N_, N_, sundials_context_);
238+
LS_ = SUNLinSol_Dense(nv_state_, A_, sundials_context_);
236239

237240
if (num_y0_vars_ + num_args_vars_ > 0) {
238-
nv_state_sens_ = N_VCloneVectorArrayEmpty_Serial(
241+
nv_state_sens_ = N_VCloneEmptyVectorArray(
239242
num_y0_vars_ + num_args_vars_, nv_state_);
240243
for (std::size_t i = 0; i < num_y0_vars_ + num_args_vars_; i++) {
241244
NV_DATA_S(nv_state_sens_[i]) = &coupled_state_[N_] + i * N_;
@@ -248,8 +251,8 @@ class cvodes_integrator {
248251
SUNMatDestroy(A_);
249252
N_VDestroy_Serial(nv_state_);
250253
if (num_y0_vars_ + num_args_vars_ > 0) {
251-
N_VDestroyVectorArray_Serial(nv_state_sens_,
252-
num_y0_vars_ + num_args_vars_);
254+
N_VDestroyVectorArray(nv_state_sens_,
255+
num_y0_vars_ + num_args_vars_);
253256
}
254257
}
255258

@@ -264,7 +267,7 @@ class cvodes_integrator {
264267
std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> operator()() {
265268
std::vector<Eigen::Matrix<T_Return, Eigen::Dynamic, 1>> y;
266269

267-
void* cvodes_mem = CVodeCreate(Lmm);
270+
void* cvodes_mem = CVodeCreate(Lmm, sundials_context_);
268271
if (cvodes_mem == nullptr) {
269272
throw std::runtime_error("CVodeCreate failed to allocate memory");
270273
}

‎stan/math/rev/functor/cvodes_integrator_adjoint.hpp‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <stan/math/prim/err.hpp>
99
#include <stan/math/prim/fun/value_of.hpp>
1010
#include <stan/math/prim/functor/for_each.hpp>
11+
#include <sundials/sundials_context.h>
1112
#include <cvodes/cvodes.h>
1213
#include <nvector/nvector_serial.h>
1314
#include <sunmatrix/sunmatrix_dense.h>
@@ -70,6 +71,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
7071
* vari class).
7172
*/
7273
struct cvodes_solver : public chainable_alloc {
74+
sundials::Context sundials_context_;
7375
const std::string function_name_str_;
7476
const std::decay_t<F> f_;
7577
const size_t N_;
@@ -107,6 +109,7 @@ class cvodes_integrator_adjoint_vari : public vari_base {
107109
size_t num_args_vars, int solver_forward,
108110
const T_Args&... args)
109111
: chainable_alloc(),
112+
sundials_context_(),
110113
function_name_str_(function_name),
111114
f_(std::forward<FF>(f)),
112115
N_(N),
@@ -119,22 +122,22 @@ class cvodes_integrator_adjoint_vari : public vari_base {
119122
state_backward_(Eigen::VectorXd::Zero(N)),
120123
quad_(Eigen::VectorXd::Zero(num_args_vars)),
121124
t0_(t0),
122-
nv_state_forward_(N_VMake_Serial(N, state_forward_.data())),
123-
nv_state_backward_(N_VMake_Serial(N, state_backward_.data())),
124-
nv_quad_(N_VMake_Serial(num_args_vars, quad_.data())),
125+
nv_state_forward_(N_VMake_Serial(N, state_forward_.data(), sundials_context_)),
126+
nv_state_backward_(N_VMake_Serial(N, state_backward_.data(), sundials_context_)),
127+
nv_quad_(N_VMake_Serial(num_args_vars, quad_.data(), sundials_context_)),
125128
nv_absolute_tolerance_forward_(
126-
N_VMake_Serial(N, absolute_tolerance_forward_.data())),
129+
N_VMake_Serial(N, absolute_tolerance_forward_.data(), sundials_context_)),
127130
nv_absolute_tolerance_backward_(
128-
N_VMake_Serial(N, absolute_tolerance_backward_.data())),
129-
A_forward_(SUNDenseMatrix(N, N)),
130-
A_backward_(SUNDenseMatrix(N, N)),
131+
N_VMake_Serial(N, absolute_tolerance_backward_.data(), sundials_context_)),
132+
A_forward_(SUNDenseMatrix(N, N, sundials_context_)),
133+
A_backward_(SUNDenseMatrix(N, N, sundials_context_)),
131134
LS_forward_(
132135
N == 0 ? nullptr
133-
: SUNDenseLinearSolver(nv_state_forward_, A_forward_)),
136+
: SUNLinSol_Dense(nv_state_forward_, A_forward_, sundials_context_)),
134137
LS_backward_(
135138
N == 0 ? nullptr
136-
: SUNDenseLinearSolver(nv_state_backward_, A_backward_)),
137-
cvodes_mem_(CVodeCreate(solver_forward)),
139+
: SUNLinSol_Dense(nv_state_backward_, A_backward_, sundials_context_)),
140+
cvodes_mem_(CVodeCreate(solver_forward, sundials_context_)),
138141
local_args_tuple_(deep_copy_vars(args)...),
139142
value_of_args_tuple_(value_of(args)...) {
140143
if (cvodes_mem_ == nullptr) {

‎stan/math/rev/functor/idas_forward_system.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,8 @@ class idas_forward_system : public idas_system<F, Tyy, Typ, Tpar> {
6666
*/
6767
~idas_forward_system() {
6868
if (this->need_sens) {
69-
N_VDestroyVectorArray_Serial(this->nv_yys_, this->ns_);
70-
N_VDestroyVectorArray_Serial(this->nv_yps_, this->ns_);
69+
N_VDestroyVectorArray(this->nv_yys_, this->ns_);
70+
N_VDestroyVectorArray(this->nv_yps_, this->ns_);
7171
}
7272
}
7373

‎stan/math/rev/functor/idas_integrator.hpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ namespace math {
2121
* IDAS DAE integrator.
2222
*/
2323
class idas_integrator {
24+
sundials::Context sundials_context_;
25+
2426
const double rtol_;
2527
const double atol_;
2628
const int64_t max_num_steps_;
@@ -80,7 +82,7 @@ class idas_integrator {
8082
*/
8183
idas_integrator(const double rtol, const double atol,
8284
const int64_t max_num_steps = IDAS_MAX_STEPS)
83-
: rtol_(rtol), atol_(atol), max_num_steps_(max_num_steps) {
85+
: sundials_context_(), rtol_(rtol), atol_(atol), max_num_steps_(max_num_steps) {
8486
if (rtol_ <= 0) {
8587
invalid_argument("idas_integrator", "relative tolerance,", rtol_, "",
8688
", must be greater than 0");
@@ -137,8 +139,8 @@ class idas_integrator {
137139
typename Dae::return_type res_yy(
138140
ts.size(), std::vector<typename Dae::scalar_type>(n, 0));
139141

140-
auto A = SUNDenseMatrix(n, n);
141-
auto LS = SUNDenseLinearSolver(yy, A);
142+
auto A = SUNDenseMatrix(n, n, sundials_context_);
143+
auto LS = SUNDenseLinearSolver(yy, A, sundials_context_);
142144

143145
try {
144146
CHECK_IDAS_CALL(IDASetUserData(mem, dae.to_user_data()));

‎stan/math/rev/functor/idas_system.hpp‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include <stan/math/prim/fun/dot_self.hpp>
88
#include <stan/math/prim/fun/typedefs.hpp>
99
#include <stan/math/prim/fun/value_of.hpp>
10+
#include <sundials/sundials_context.h>
1011
#include <idas/idas.h>
1112
#include <nvector/nvector_serial.h>
1213
#include <ostream>
@@ -87,6 +88,7 @@ namespace math {
8788
template <typename F, typename Tyy, typename Typ, typename Tpar>
8889
class idas_system {
8990
protected:
91+
sundials::Context sundials_context_;
9092
const F& f_;
9193
const std::vector<Tyy>& yy_;
9294
const std::vector<Typ>& yp_;
@@ -132,7 +134,8 @@ class idas_system {
132134
const std::vector<Tyy>& yy0, const std::vector<Typ>& yp0,
133135
const std::vector<Tpar>& theta, const std::vector<double>& x_r,
134136
const std::vector<int>& x_i, std::ostream* msgs)
135-
: f_(f),
137+
: sundials_context_(),
138+
f_(f),
136139
yy_(yy0),
137140
yp_(yp0),
138141
yy_val_(value_of(yy0)),
@@ -144,12 +147,12 @@ class idas_system {
144147
M_(theta.size()),
145148
ns_((is_var_yy0 ? N_ : 0) + (is_var_yp0 ? N_ : 0)
146149
+ (is_var_par ? M_ : 0)),
147-
nv_yy_(N_VMake_Serial(N_, yy_val_.data())),
148-
nv_yp_(N_VMake_Serial(N_, yp_val_.data())),
150+
nv_yy_(N_VMake_Serial(N_, yy_val_.data(), sundials_context_)),
151+
nv_yp_(N_VMake_Serial(N_, yp_val_.data(), sundials_context_)),
149152
rr_val_(N_, 0.0),
150-
nv_rr_(N_VMake_Serial(N_, rr_val_.data())),
151-
id_(N_VNew_Serial(N_)),
152-
mem_(IDACreate()),
153+
nv_rr_(N_VMake_Serial(N_, rr_val_.data(), sundials_context_)),
154+
id_(N_VNew_Serial(N_, sundials_context_)),
155+
mem_(IDACreate(sundials_context_)),
153156
msgs_(msgs) {
154157
try {
155158
if (nv_yy_ == NULL || nv_yp_ == NULL) {

‎stan/math/rev/functor/kinsol_data.hpp‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <stan/math/prim/fun/to_array_1d.hpp>
77
#include <stan/math/prim/fun/to_vector.hpp>
88
#include <stan/math/prim/functor/apply.hpp>
9+
#include <sundials/sundials_context.h>
910
#include <kinsol/kinsol.h>
1011
#include <sunmatrix/sunmatrix_dense.h>
1112
#include <sunlinsol/sunlinsol_dense.h>
@@ -35,6 +36,7 @@ class kinsol_system_data {
3536
typedef kinsol_system_data<F1, Args...> system_data;
3637

3738
public:
39+
sundials::Context sundials_context_;
3840
N_Vector nv_x_;
3941
SUNMatrix J_;
4042
SUNLinearSolver LS_;
@@ -48,10 +50,11 @@ class kinsol_system_data {
4850
N_(x.size()),
4951
msgs_(msgs),
5052
args_tuple_(args...),
51-
nv_x_(N_VMake_Serial(N_, &to_array_1d(x_)[0])),
52-
J_(SUNDenseMatrix(N_, N_)),
53-
LS_(SUNLinSol_Dense(nv_x_, J_)),
54-
kinsol_memory_(KINCreate()) {}
53+
sundials_context_(),
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_)),
57+
kinsol_memory_(KINCreate(sundials_context_)) {}
5558

5659
~kinsol_system_data() {
5760
N_VDestroy_Serial(nv_x_);

‎stan/math/rev/functor/kinsol_solve.hpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <stan/math/prim/err.hpp>
77
#include <stan/math/prim/fun/to_array_1d.hpp>
88
#include <stan/math/prim/fun/to_vector.hpp>
9+
#include <sundials/sundials_context.h>
910
#include <kinsol/kinsol.h>
1011
#include <sunmatrix/sunmatrix_dense.h>
1112
#include <sunlinsol/sunlinsol_dense.h>
@@ -71,8 +72,8 @@ Eigen::VectorXd kinsol_solve(const F1& f, const Eigen::VectorXd& x,
7172
CHECK_KINSOL_CALL(KINInit(kinsol_data.kinsol_memory_,
7273
&system_data::kinsol_f_system, kinsol_data.nv_x_));
7374

74-
N_Vector scaling = N_VNew_Serial(N);
75-
N_Vector nv_x = N_VNew_Serial(N);
75+
N_Vector scaling = N_VNew_Serial(N, kinsol_data.sundials_context_.SUNContext());
76+
N_Vector nv_x = N_VNew_Serial(N, kinsol_data.sundials_context_.SUNContext());
7677
Eigen::VectorXd x_solution(N);
7778

7879
try {

0 commit comments

Comments
 (0)