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)
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*') check(