def _set_tolerances(self, atol=None, rtol=None):
atol = np.array(atol)
rtol = np.array(rtol)
if atol.ndim == 1 and rtol.ndim == 1:
atol = sunode.from_numpy(atol)
rtol = sunode.from_numpy(rtol)
check(lib.CVodeVVtolerances(self._ode, rtol.c_ptr, atol.c_ptr))
elif atol.ndim == 1 and rtol.ndim == 0:
atol = sunode.from_numpy(atol)
check(lib.CVodeSVtolerances(self._ode, rtol, atol.c_ptr))
elif atol.ndim == 0 and rtol.ndim == 1:
rtol = sunode.from_numpy(rtol)
check(lib.CVodeVStolerances(self._ode, rtol.c_ptr, atol))
elif atol.ndim == 0 and rtol.ndim == 0:
check(lib.CVodeSStolerances(self._ode, rtol, atol))
else:
raise ValueError('Invalid tolerance.')
self._atol = atol
self._rtol = rtol
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 constraints(self, constraints: Optional[np.ndarray]) -> None:
self.mark_changed()
if constraints is None:
check(lib.CVodeSetConstraints(self.c_ptr, ffi.NULL))
return
assert constraints.shape == (self.n_states, )
if not hasattr(self, '_constraints_buffer'):
self._constraints_buffer = sunode.from_numpy(constraints)
self.borrow(self._constraints_buffer)
self._constraints_buffer.data[:] = constraints
check(
lib.CVodeSetConstraints(self.c_ptr,
self._constraints_buffer.c_ptr))
def tolerance(self, rtol: float, atol: Union[np.ndarray, float]) -> None:
self.mark_changed()
self._atol = np.array(atol)
self._rtol = rtol
if self._atol.ndim == 1:
if not hasattr(self, '_atol_buffer'):
self._atol_buffer = sunode.from_numpy(atol)
self.borrow(self._atol_buffer)
self._atol_buffer.data[:] = atol
check(
lib.CVodeSVtolerances(self.c_ptr, self._rtol,
self._atol_buffer.c_ptr))
elif self._atol.ndim == 0:
check(lib.CVodeSStolerances(self.c_ptr, self._rtol, self._atol))
else:
raise ValueError('Invalid absolute tolerances.')
if adjoint_solver == 'BDF':
self._adjoint_solver_type = lib.CV_BDF
elif adjoint_solver == 'ADAMS':
self._adjoint_solver_type = lib.CV_ADAMS
else:
raise ValueError(f'Unknown solver {solver}.')
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:
