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 __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)
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:
assert False
self._init_backward(checkpoint_n, interpolation)