Can now use arrays for tols. Addressing Issue #31

CHANGES.md

@@ -9,6 +9,10 @@ Major Changes
   - Added `noexcept` to pure cython functions to avoid a potential python error check.
 
 New Features
+- Added the ability to pass arrayed versions of rtol and atol to both the numba and cython-based solvers (cyrk_ode and CySolver).
+  - For both solvers, you can pass the optional argument "rtols" and/or "atols". These must be C-contiguous numpy arrays with float64 dtypes. They must have the same size as y0.
+  - Added tests to check functionality for all solvers.
+  - This resolves [https://github.com/jrenaud90/CyRK/issues/31][Issue 31].
 - Added new optional argument to all solvers `max_steps` which allows the user to control how many steps the solver is allowed to take.
   - If exceeded the integration with fail (softly). 
   - Defaults to 95% of `sys.maxsize` (depends on system architecture).

CyRK/cy/cyrk.pyx

@@ -89,6 +89,8 @@ def cyrk_ode(
     tuple args = None,
     double rtol = 1.e-6,
     double atol = 1.e-8,
+    double[::1] rtols = None,
+    double[::1] atols = None,
     double max_step_size = MAX_STEP,
     double first_step = 0.,
     unsigned char rk_method = 1,
@@ -241,15 +243,43 @@ def cyrk_ode(
         #  on computation.
         store_extras_during_integration = False
 
-    # # Determine integration parameters
-    # Check tolerances
-    if rtol < EPS_100:
-        rtol = EPS_100
+    # # Determine integration tolerances
+    cdef double rtol_tmp, atol_tmp
+    use_arg_arrays = False
+    use_atol_array = False
+    cdef np.ndarray[np.float64_t, ndim=1, mode='c'] rtol_array, atol_array
+    rtol_array = np.empty(y_size, dtype=np.float64, order='C')
+    atol_array = np.empty(y_size, dtype=np.float64, order='C')
+    cdef double[::1] rtols_view, atols_view
+    rtols_view = rtol_array
+    atols_view = atol_array
+
+    if rtols is not None:
+        # Using arrayed rtol
+        if len(rtols) != y_size:
+            raise AttributeError('rtols must be the same size as y0.')
+        for i in range(y_size):
+            rtol_tmp = rtols[i]
+            if rtol_tmp < EPS_100:
+                rtol_tmp = EPS_100
+            rtols_view[i] = rtol_tmp
+    else:
+        # Using constant rtol
+        # Check tolerances
+        if rtol < EPS_100:
+            rtol = EPS_100
+        for i in range(y_size):
+            rtols_view[i] = rtol
 
-    #     atol_arr = np.asarray(atol, dtype=np.complex128)
-    #     if atol_arr.ndim > 0 and atol_arr.shape[0] != y_size:
-    #         # atol must be either the same for all y or must be provided as an array, one for each y.
-    #         raise Exception
+    if atols is not None:
+        # Using arrayed atol
+        if len(atols) != y_size:
+            raise AttributeError('atols must be the same size as y0.')
+        for i in range(y_size):
+            atols_view[i] = atols[i]
+    else:
+        for i in range(y_size):
+            atols_view[i] = atol
 
     # Determine maximum number of steps
     cdef Py_ssize_t max_steps_touse
@@ -522,8 +552,8 @@ def cyrk_ode(
             d0 = 0.
             d1 = 0.
             for i in range(y_size):
-                scale = atol + dabs(y_old_view[i]) * rtol
 
+                scale = atols_view[i] + dabs(y_old_view[i]) * rtols_view[i]
                 d0_abs = dabs(y_old_view[i] / scale)
                 d1_abs = dabs(dydt_old_view[i] / scale)
                 d0 += (d0_abs * d0_abs)
@@ -554,7 +584,7 @@ def cyrk_ode(
             d2 = 0.
             for i in range(y_size):
                 dydt_new_view[i] = diffeq_out_view[i]
-                scale = atol + dabs(y_old_view[i]) * rtol
+                scale = atols_view[i] + dabs(y_old_view[i]) * rtols_view[i]
                 d2_abs = dabs( (dydt_new_view[i] - dydt_old_view[i]) / scale)
                 d2 += (d2_abs * d2_abs)
 
@@ -697,9 +727,7 @@ def cyrk_ode(
                 if i < extra_start:
                     # Set diffeq results
                     dydt_new_view[i] = diffeq_out_view[i]
-
-                    # Find scale of y for error calculations
-                    scale_view[i] = atol + max(dabs(y_old_view[i]), dabs(y_new_view[i])) * rtol
+                    scale_view[i] = atols_view[i] + max(dabs(y_old_view[i]), dabs(y_new_view[i])) * rtols_view[i]
 
                     # Set last array of K equal to dydt
                     K_view[rk_n_stages, i] = dydt_new_view[i]

CyRK/cy/cysolver.pxd

@@ -38,7 +38,7 @@ cdef class CySolver:
     cdef public bool_cpp_t success
     cdef double t_start, t_end, t_delta, t_delta_abs, direction_inf
     cdef bool_cpp_t direction_flag
-    cdef double rtol, atol
+    cdef double[::1] rtols_view, atols_view
     cdef double step_size, max_step_size
     cdef double first_step
     cdef Py_ssize_t expected_size, num_concats, max_steps
@@ -65,20 +65,48 @@ cdef class CySolver:
 
     # Class functions
     cpdef void reset_state(self)
+
     cdef double calc_first_step(self) noexcept nogil
+
     cdef void rk_step(self) noexcept nogil
+
     cpdef void solve(self, bool_cpp_t reset = *)
+
     cdef void _solve(self, bool_cpp_t reset = *)
+
     cdef void interpolate(self)
+
     cpdef void change_t_span(self, (double, double) t_span, bool_cpp_t auto_reset_state = *)
+
     cpdef void change_y0(self, const double[::1] y0, bool_cpp_t auto_reset_state = *)
+
     cpdef void change_args(self, tuple args, bool_cpp_t auto_reset_state = *)
-    cpdef void change_tols(self, double rtol = *, double atol = *, bool_cpp_t auto_reset_state = *)
+
+    cpdef void change_tols(self, double rtol = *,
+                           double atol = *,
+                           double[::1] rtols = *,
+                           double[::1] atols = *,
+                           bool_cpp_t auto_reset_state = *)
+
     cpdef void change_max_step_size(self, double max_step_size, bool_cpp_t auto_reset_state = *)
+
     cpdef void change_first_step(self, double first_step, bool_cpp_t auto_reset_state = *)
+
     cpdef void change_t_eval(self, const double[:] t_eval, bool_cpp_t auto_reset_state = *)
-    cpdef void change_parameters(self, (double, double) t_span = *, const double[::1] y0 = *, tuple args = *,
-                                double rtol = *, double atol = *, double max_step_size = *, double first_step = *,
-                                const double[::1] t_eval = *, bool_cpp_t auto_reset_state = *, bool_cpp_t auto_solve = *)
+
+    cpdef void change_parameters(self, (double, double) t_span = *,
+                                const double[::1] y0 = *,
+                                tuple args = *,
+                                double rtol = *,
+                                double atol = *,
+                                double[::1] rtols = *,
+                                double[::1] atols = *,
+                                double max_step_size = *,
+                                double first_step = *,
+                                const double[::1] t_eval = *,
+                                bool_cpp_t auto_reset_state = *,
+                                bool_cpp_t auto_solve = *)
+
     cdef void update_constants(self) noexcept nogil
+
     cdef void diffeq(self) noexcept nogil

CyRK/cy/cysolver.pyx

@@ -36,21 +36,23 @@ cdef class CySolver:
 
 
     def __init__(self,
-                 (double, double) t_span,
-                 const double[::1] y0,
-                 tuple args = None,
-                 double rtol = 1.e-6,
-                 double atol = 1.e-8,
-                 double max_step_size = MAX_STEP,
-                 double first_step = 0.,
-                 unsigned char rk_method = 1,
-                 const double[::1] t_eval = None,
-                 bool_cpp_t capture_extra = False,
-                 Py_ssize_t num_extra = 0,
-                 bool_cpp_t interpolate_extra = False,
-                 Py_ssize_t expected_size = 0,
-                 Py_ssize_t max_steps = 0,
-                 bool_cpp_t auto_solve = True):
+            (double, double) t_span,
+            const double[::1] y0,
+            tuple args = None,
+            double rtol = 1.e-6,
+            double atol = 1.e-8,
+            double[::1] rtols = None,
+            double[::1] atols = None,
+            double max_step_size = MAX_STEP,
+            double first_step = 0.,
+            unsigned char rk_method = 1,
+            const double[::1] t_eval = None,
+            bool_cpp_t capture_extra = False,
+            Py_ssize_t num_extra = 0,
+            bool_cpp_t interpolate_extra = False,
+            Py_ssize_t expected_size = 0,
+            Py_ssize_t max_steps = 0,
+            bool_cpp_t auto_solve = True):
 
         # Setup loop variables
         cdef Py_ssize_t i, j
@@ -93,17 +95,40 @@ cdef class CySolver:
             self.direction_flag = False
             self.direction_inf  = -INF
 
-        # # Determine integration parameters
-        # Add tolerances
-        self.rtol = rtol
-        self.atol = atol
-        if self.rtol < EPS_100:
-            self.rtol = EPS_100
-        # TODO: array based atol
-        #     atol_arr = np.asarray(atol, dtype=)
-        #     if atol_arr.ndim > 0 and atol_arr.shape[0] != y_size:
-        #         # atol must be either the same for all y or must be provided as an array, one for each y.
-        #         raise Exception
+        # # Determine integration tolerances
+        cdef double rtol_tmp
+        cdef np.ndarray[np.float64_t, ndim=1, mode='c'] rtol_array, atol_array
+        rtol_array = np.empty(self.y_size, dtype=np.float64, order='C')
+        atol_array = np.empty(self.y_size, dtype=np.float64, order='C')
+        self.rtols_view = rtol_array
+        self.atols_view = atol_array
+
+        if rtols is not None:
+            # Using arrayed rtol
+            if len(rtols) != self.y_size:
+                raise AttributeError('rtols must be the same size as y0.')
+            for i in range(self.y_size):
+                rtol_tmp = rtols[i]
+                if rtol_tmp < EPS_100:
+                    rtol_tmp = EPS_100
+                self.rtols_view[i] = rtol_tmp
+        else:
+            # Using constant rtol
+            # Check tolerances
+            if rtol < EPS_100:
+                rtol = EPS_100
+            for i in range(self.y_size):
+                self.rtols_view[i] = rtol
+
+        if atols is not None:
+            # Using arrayed atol
+            if len(atols) != self.y_size:
+                raise AttributeError('atols must be the same size as y0.')
+            for i in range(self.y_size):
+                self.atols_view[i] = atols[i]
+        else:
+            for i in range(self.y_size):
+                self.atols_view[i] = atol
 
         # Determine maximum number of steps
         if max_steps == 0:
@@ -376,7 +401,8 @@ cdef class CySolver:
             d1 = 0.
             for i in range(self.y_size):
                 y_old_tmp = self.y_old_view[i]
-                scale = self.atol + fabs(y_old_tmp) * self.rtol
+
+                scale = self.atols_view[i] + fabs(y_old_tmp) * self.rtols_view[i]
 
                 d0_abs = fabs(y_old_tmp / scale)
                 d1_abs = fabs(self.dy_old_view[i] / scale)
@@ -406,7 +432,8 @@ cdef class CySolver:
             # Find the norm for d2
             d2 = 0.
             for i in range(self.y_size):
-                scale = self.atol + fabs(self.y_old_view[i]) * self.rtol
+
+                scale = self.atols_view[i] + fabs(self.y_old_view[i]) * self.rtols_view[i]
                 d2_abs = fabs( (self.dy_new_view[i] - self.dy_old_view[i]) / scale)
                 d2 += (d2_abs * d2_abs)
 
@@ -539,7 +566,8 @@ cdef class CySolver:
                 error_norm5 = 0.
                 for i in range(self.y_size):
                     # Find scale of y for error calculations
-                    scale = self.atol + max(fabs(self.y_old_view[i]), fabs(self.y_new_view[i])) * self.rtol
+                    scale = (self.atols_view[i] +
+                             max(fabs(self.y_old_view[i]), fabs(self.y_new_view[i])) * self.rtols_view[i])
 
                     # Set last array of K equal to dydt
                     self.K_view[self.rk_n_stages, i] = self.dy_new_view[i]
@@ -574,7 +602,8 @@ cdef class CySolver:
                 error_norm = 0.
                 for i in range(self.y_size):
                     # Find scale of y for error calculations
-                    scale = self.atol + max(fabs(self.y_old_view[i]), fabs(self.y_new_view[i])) * self.rtol
+                    scale = (self.atols_view[i] +
+                             max(fabs(self.y_old_view[i]), fabs(self.y_new_view[i])) * self.rtols_view[i])
 
                     # Set last array of K equal to dydt
                     self.K_view[self.rk_n_stages, i] = self.dy_new_view[i]
@@ -998,21 +1027,44 @@ cdef class CySolver:
             self.reset_state()
 
 
-    cpdef void change_tols(self, double rtol = NAN, double atol = NAN, bool_cpp_t auto_reset_state = False):
+    cpdef void change_tols(self, double rtol = NAN, double atol = NAN,
+                           double[::1] rtols = None, double[::1] atols = None,
+                           bool_cpp_t auto_reset_state = False):
 
         # Update tolerances
-        if not isnan(rtol):
-            self.rtol = rtol
-        if not isnan(atol):
-            self.atol = atol
-
-        if self.rtol < EPS_100:
-            self.rtol = EPS_100
-        # TODO: array based atol
-        #     atol_arr = np.asarray(atol, dtype=)
-        #     if atol_arr.ndim > 0 and atol_arr.shape[0] != y_size:
-        #         # atol must be either the same for all y or must be provided as an array, one for each y.
-        #         raise Exception
+        cdef double rtol_tmp
+        cdef np.ndarray[np.float64_t, ndim=1, mode='c'] rtol_array, atol_array
+        rtol_array = np.empty(self.y_size, dtype=np.float64, order='C')
+        atol_array = np.empty(self.y_size, dtype=np.float64, order='C')
+        self.rtols_view = rtol_array
+        self.atols_view = atol_array
+
+        if rtols is not None:
+            # Using arrayed rtol
+            if len(rtols) != self.y_size:
+                raise AttributeError('rtols must be the same size as y0.')
+            for i in range(self.y_size):
+                rtol_tmp = rtols[i]
+                if rtol_tmp < EPS_100:
+                    rtol_tmp = EPS_100
+                self.rtols_view[i] = rtol_tmp
+        elif not isnan(rtol):
+            # Using constant rtol
+            # Check tolerances
+            if rtol < EPS_100:
+                rtol = EPS_100
+            for i in range(self.y_size):
+                self.rtols_view[i] = rtol
+
+        if atols is not None:
+            # Using arrayed atol
+            if len(atols) != self.y_size:
+                raise AttributeError('atols must be the same size as y0.')
+            for i in range(self.y_size):
+                self.atols_view[i] = atols[i]
+        elif not isnan(atol):
+            for i in range(self.y_size):
+                self.atols_view[i] = atol
 
         # A change to tolerances will affect the first step's size
         self.recalc_firststep = True
@@ -1076,6 +1128,8 @@ cdef class CySolver:
             tuple args = None,
             double rtol = NAN,
             double atol = NAN,
+            double[::1] rtols = None,
+            double[::1] atols = None,
             double max_step_size = NAN,
             double first_step = NAN,
             const double[::1] t_eval = None,
@@ -1091,8 +1145,8 @@ cdef class CySolver:
         if args is not None:
             self.change_args(args, auto_reset_state=False)
 
-        if not isnan(rtol) or not isnan(atol):
-            self.change_tols(rtol=rtol, atol=atol, auto_reset_state=False)
+        if (not isnan(rtol)) or (not isnan(atol)) or (rtols is not None) or (atols is not None):
+            self.change_tols(rtol=rtol, atol=atol, rtols=rtols, atols=atols, auto_reset_state=False)
 
         if not isnan(max_step_size):
             self.change_max_step_size(max_step_size, auto_reset_state=False)