def __init__(self,
problem: Problem,
*,
compute_sens: bool = False,
abstol: float = 1e-10,
reltol: float = 1e-10,
sens_mode: Optional[str] = None,
scaling_factors: Optional[np.ndarray] = None,
constraints: Optional[np.ndarray] = None,
solver='BDF'):
self._problem = problem
self._user_data = problem.make_user_data()
n_states = self._problem.n_states
n_params = self._problem.n_params
self._state_buffer = sunode.empty_vector(n_states)
self._state_buffer.data[:] = 0
self._jac = check(lib.SUNDenseMatrix(n_states, n_states))
self._constraints = constraints
if solver == 'BDF':
solver_kind = lib.CV_BDF
elif solver == 'ADAMS':
solver_kind = lib.CV_ADAMS
else:
assert False
self._ode = check(lib.CVodeCreate(solver_kind))
rhs = problem.make_sundials_rhs()
check(lib.CVodeInit(self._ode, rhs.cffi, 0., self._state_buffer.c_ptr))
self._set_tolerances(abstol, reltol)
if self._constraints is not None:
assert constraints.shape == (n_states, )
self._constraints_vec = sunode.from_numpy(constraints)
check(
lib.CVodeSetConstraints(self._ode,
self._constraints_vec.c_ptr))
self._make_linsol()
user_data_p = ffi.cast(
'void *', ffi.addressof(ffi.from_buffer(self._user_data.data)))
check(lib.CVodeSetUserData(self._ode, user_data_p))
self._compute_sens = compute_sens
if compute_sens:
sens_rhs = self._problem.make_sundials_sensitivity_rhs()
self._init_sens(sens_rhs, sens_mode)
def _init_backward(self, checkpoint_n, interpolation):
check(lib.CVodeAdjInit(self._ode, checkpoint_n, interpolation))
# Initialized by CVodeCreateB
backward_ode = ffi.new('int*')
check(
lib.CVodeCreateB(self._ode, self._adjoint_solver_type,
backward_ode))
self._odeB = backward_ode[0]
self._state_bufferB = sunode.empty_vector(self._problem.n_states)
check(
lib.CVodeInitB(self._ode, self._odeB, self._adj_rhs.cffi, 0.,
self._state_bufferB.c_ptr))
# TODO
check(lib.CVodeSStolerancesB(self._ode, self._odeB, 1e-10, 1e-10))
linsolver = check(
lib.SUNLinSol_Dense(self._state_bufferB.c_ptr, self._jacB))
check(
lib.CVodeSetLinearSolverB(self._ode, self._odeB, linsolver,
self._jacB))
self._jac_funcB = self._problem.make_sundials_adjoint_jac_dense()
check(lib.CVodeSetJacFnB(self._ode, self._odeB, self._jac_funcB.cffi))
user_data_p = ffi.cast(
'void *', ffi.addressof(ffi.from_buffer(self._user_data.data)))
check(lib.CVodeSetUserDataB(self._ode, self._odeB, user_data_p))
self._quad_buffer = sunode.empty_vector(self._problem.n_params)
self._quad_buffer_out = sunode.empty_vector(self._problem.n_params)
check(
lib.CVodeQuadInitB(self._ode, self._odeB, self._quad_rhs.cffi,
self._quad_buffer.c_ptr))
check(lib.CVodeQuadSStolerancesB(self._ode, self._odeB, 1e-10, 1e-10))
check(lib.CVodeSetQuadErrConB(self._ode, self._odeB, 1))
def _init_sens(self, sens_rhs, sens_mode, scaling_factors=None) -> None:
if sens_mode == 'simultaneous':
sens_mode = lib.CV_SIMULTANEOUS
elif sens_mode == 'staggered':
sens_mode = lib.CV_STAGGERED
elif sens_mode == 'staggered1':
raise ValueError('staggered1 not implemented.')
else:
raise ValueError(
'sens_mode must be one of "simultaneous" and "staggered".')
self._sens_mode = sens_mode
n_params = self._problem.n_params
yS = check(lib.N_VCloneVectorArray(n_params, self._state_buffer.c_ptr))
vecs = [sunode.vector.Vector(yS[i]) for i in range(n_params)]
for vec in vecs:
vec.data[:] = 0
self._sens_buffer_array = yS
self._sens_buffers = vecs
check(
lib.CVodeSensInit(self._ode, n_params, sens_mode, sens_rhs.cffi,
yS))
if scaling_factors is not None:
if scaling_factors.shape != (n_params, ):
raise ValueError('Invalid shape of scaling_factors.')
self._scaling_factors = scaling_factors
NULL_D = ffi.cast('double *', 0)
NULL_I = ffi.cast('int *', 0)
pbar_p = ffi.cast(
'double *',
ffi.addressof(ffi.from_buffer(scaling_factors.data)))
check(lib.CVodeSetSensParams(self._ode, NULL_D, pbar_p, NULL_I))
check(lib.CVodeSensEEtolerances(self._ode)) # TODO
check(lib.CVodeSetSensErrCon(self._ode, 1)) # TODO
self._ode = check(lib.CVodeCreate(self._solver_type))
check(lib.CVodeInit(self._ode, rhs.cffi, 0., self._state_buffer.c_ptr))
self._set_tolerances(abstol, reltol)
if self._constraints is not None:
self._constraints = np.broadcast_to(constraints,
(n_states, )).copy()
self._constraints_vec = sunode.from_numpy(self._constraints)
check(
lib.CVodeSetConstraints(self._ode,
self._constraints_vec.c_ptr))
self._make_linsol()
user_data_p = ffi.cast(
'void *', ffi.addressof(ffi.from_buffer(self._user_data.data)))
check(lib.CVodeSetUserData(self._ode, user_data_p))
if interpolation == 'polynomial':
interpolation = lib.CV_POLYNOMIAL
elif interpolation == 'hermite':
interpolation = lib.CV_HERMITE
else:
assert False
self._init_backward(checkpoint_n, interpolation)
def _init_backward(self, checkpoint_n, interpolation):
check(lib.CVodeAdjInit(self._ode, checkpoint_n, interpolation))
# Initialized by CVodeCreateB
backward_ode = ffi.new('int*')