Add some safety checks

Author: aseyboldt

sunode/problem.py

@@ -49,7 +49,7 @@ def make_rhs_jac_prod(self):  # type: ignore
         return NotImplemented
 
     def make_user_data(self) -> np.ndarray:
-        return np.recarray((), dtype=self.user_data_dtype)
+        return np.zeros((), dtype=self.user_data_dtype).view(np.recarray)
 
     def update_params(self, user_data: np.ndarray, params: np.ndarray) -> None:
         if not self.user_data_dtype == self.params_dtype:
@@ -166,10 +166,14 @@ def make_sundials_rhs(self) -> Any:
         func_type = numba.core.typing.cffi_utils.map_type(ffi.typeof('CVRhsFn'))
         func_type = func_type.return_type(*(func_type.args[:-1] + (user_ndtype_p,)))
 
+        n_states = self.n_states
+
         @numba.cfunc(func_type)
         def rhs_wrapper(t, y_, out_, user_data_):  # type: ignore
             y_ptr = N_VGetArrayPointer_Serial(y_)
             n_vars = N_VGetLength_Serial(y_)
+            if n_vars != n_states:
+                return -1
             out_ptr = N_VGetArrayPointer_Serial(out_)
             y = numba.carray(y_ptr, (n_vars,)).view(state_dtype)[0]
             out = numba.carray(out_ptr, (n_vars,))

sunode/solver.py

@@ -221,8 +221,12 @@ def _init_sundials(self):
 
         self._ode = check(lib.CVodeCreate(self._solver_kind))
         rhs = self._problem.make_sundials_rhs()
+        self._rhs = rhs
         check(lib.CVodeInit(self._ode, rhs.cffi, 0., self._state_buffer.c_ptr))
 
+        user_data_p = ffi.cast('void *', ffi.addressof(ffi.from_buffer(self._user_data.data)))
+        check(lib.CVodeSetUserData(self._ode, user_data_p))
+
         self._set_tolerances(self._abstol, self._reltol)
         if self._constraints is not None:
             assert self._constraints.shape == (n_states,)
@@ -231,9 +235,6 @@ def _init_sundials(self):
 
         self._make_linsol(self._linear_solver_kind)
 
-        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 = self._sens_mode is not None
         if self._compute_sens:
             sens_rhs = self._problem.make_sundials_sensitivity_rhs()
@@ -340,15 +341,20 @@ def _init_sens(self, sens_rhs, sens_mode, scaling_factors=None) -> None:
     def _set_tolerances(self, atol=None, rtol=None):
         atol = np.array(atol)
         rtol = np.array(rtol)
+        n_states = self._problem.n_states
         if atol.ndim == 1 and rtol.ndim == 1:
             atol = sunode.from_numpy(atol)
             rtol = sunode.from_numpy(rtol)
+            assert atol.shape == (n_states,)
+            assert rtol.shape == (n_states,)
             check(lib.CVodeVVtolerances(self._ode, rtol.c_ptr, atol.c_ptr))
         elif atol.ndim == 1 and rtol.ndim == 0:
             atol = sunode.from_numpy(atol)
+            assert atol.shape == (n_states,)
             check(lib.CVodeSVtolerances(self._ode, rtol, atol.c_ptr))
         elif atol.ndim == 0 and rtol.ndim == 1:
             rtol = sunode.from_numpy(rtol)
+            assert rtol.shape == (n_states,)
             check(lib.CVodeVStolerances(self._ode, rtol.c_ptr, atol))
         elif atol.ndim == 0 and rtol.ndim == 0:
             check(lib.CVodeSStolerances(self._ode, rtol, atol))
@@ -416,6 +422,7 @@ def solve(self, t0, tvals, y0, y_out, *, sens0=None, sens_out=None, max_retries=
         TOO_MUCH_WORK = lib.CV_TOO_MUCH_WORK
 
         n_params = self._problem.n_params
+        n_states = self._problem.n_states
 
         state_data = self._state_buffer.data
         state_c_ptr = self._state_buffer.c_ptr
@@ -428,6 +435,9 @@ def solve(self, t0, tvals, y0, y_out, *, sens0=None, sens_out=None, max_retries=
 
         if y0.dtype == self._problem.state_dtype:
             y0 = y0[None].view(np.float64)
+
+        if y0.shape != (n_states,):
+            raise ValueError(f"y0 should have shape {(n_states,)} but has shape {y0.shape}.")
         state_data[:] = y0
 
         time_p = ffi.new('double*')

sunode/wrappers/as_pytensor.py

@@ -53,14 +53,14 @@ def read_dict(vals, name=None):
                 tensor, dim_names = vals
             else:
                 try:
-                    tensor, dim_names = vals, pt.as_tensor_variable(vals).shape.eval()
+                    tensor, dim_names = vals, pt.as_tensor_variable(vals, dtype="float64").shape.eval()
                 except MissingInputError as e:
                     raise ValueError(
                         'Shapes of tensors need to be statically '
                         'known or given explicitly.') from e
             if isinstance(dim_names, (str, int)):
                 dim_names = (dim_names,)
-            tensor = pt.as_tensor_variable(tensor)
+            tensor = pt.as_tensor_variable(tensor, dtype="float64")
             if tensor.ndim != len(dim_names):
                 raise ValueError(
                     f"Dimension mismatch for {name}: Value has rank {tensor.ndim}, "
@@ -97,30 +97,30 @@ def read_dict(vals, name=None):
         tensor = flat_tensors[path]
         if isinstance(tensor, tuple):
             tensor, _ = tensor
-        vars.append(pt.as_tensor_variable(tensor).reshape((-1,)))
+        vars.append(pt.as_tensor_variable(tensor, dtype="float64").reshape((-1,)))
     if vars:
         params_subs_flat = pt.concatenate(vars)
     else:
-        params_subs_flat = pt.as_tensor_variable(np.zeros(0))
+        params_subs_flat = pt.as_tensor_variable(np.zeros(0), dtype="float64")
 
     vars = []
     for path in problem.params_subset.remainder.subset_paths:
         tensor = flat_tensors[path]
         if isinstance(tensor, tuple):
             tensor, _ = tensor
-        vars.append(pt.as_tensor_variable(tensor).reshape((-1,)))
+        vars.append(pt.as_tensor_variable(tensor, dtype="float64").reshape((-1,)))
     if vars:
         params_remaining_flat = pt.concatenate(vars)
     else:
-        params_remaining_flat = pt.as_tensor_variable(np.zeros(0))
+        params_remaining_flat = pt.as_tensor_variable(np.zeros(0), dtype="float64")
 
     flat_tensors = as_flattened(y0)
     vars = []
     for path in problem.state_subset.paths:
         tensor = flat_tensors[path]
         if isinstance(tensor, tuple):
             tensor, _ = tensor
-        vars.append(pt.as_tensor_variable(tensor).reshape((-1,)))
+        vars.append(pt.as_tensor_variable(tensor, dtype="float64").reshape((-1,)))
     y0_flat = pt.concatenate(vars)
 
     if derivatives == 'adjoint':
@@ -261,6 +261,25 @@ class SolveODEAdjointBackward(Op):
 
     __props__ = ('_solver_id', '_t0', '_tvals_id')
 
+    def make_nodes(self, *inputs):
+        if len(inputs) != len(self.itypes):
+            raise ValueError(
+                f"We expected {len(self.itypes)} inputs but got {len(inputs)}."
+            )
+        if not all(it.in_same_class(inp.type) for inp, it in zip(inputs, self.itypes)):
+            raise TypeError(
+                f"Invalid input types for Op {self}:\n"
+                + "\n".join(
+                    f"Input {i}/{len(inputs)}: Expected {inp}, got {out}"
+                    for i, (inp, out) in enumerate(
+                        zip(self.itypes, (inp.type for inp in inputs)),
+                        start=1,
+                    )
+                    if inp != out
+                )
+            )
+        return Apply(self, inputs, [o() for o in self.otypes])
+
     def __init__(self, solver, t0, tvals):
         self._solver = solver
         self._t0 = t0