1from itertools import groupby 2from warnings import warn 3import numpy as np 4from scipy.sparse import find, coo_matrix 5 6 7EPS = np.finfo(float).eps 8 9 10def validate_first_step(first_step, t0, t_bound): 11 """Assert that first_step is valid and return it.""" 12 if first_step <= 0: 13 raise ValueError("`first_step` must be positive.") 14 if first_step > np.abs(t_bound - t0): 15 raise ValueError("`first_step` exceeds bounds.") 16 return first_step 17 18 19def validate_max_step(max_step): 20 """Assert that max_Step is valid and return it.""" 21 if max_step <= 0: 22 raise ValueError("`max_step` must be positive.") 23 return max_step 24 25 26def warn_extraneous(extraneous): 27 """Display a warning for extraneous keyword arguments. 28 29 The initializer of each solver class is expected to collect keyword 30 arguments that it doesn't understand and warn about them. This function 31 prints a warning for each key in the supplied dictionary. 32 33 Parameters 34 ---------- 35 extraneous : dict 36 Extraneous keyword arguments 37 """ 38 if extraneous: 39 warn("The following arguments have no effect for a chosen solver: {}." 40 .format(", ".join("`{}`".format(x) for x in extraneous))) 41 42 43def validate_tol(rtol, atol, n): 44 """Validate tolerance values.""" 45 if rtol < 100 * EPS: 46 warn("`rtol` is too low, setting to {}".format(100 * EPS)) 47 rtol = 100 * EPS 48 49 atol = np.asarray(atol) 50 if atol.ndim > 0 and atol.shape != (n,): 51 raise ValueError("`atol` has wrong shape.") 52 53 if np.any(atol < 0): 54 raise ValueError("`atol` must be positive.") 55 56 return rtol, atol 57 58 59def norm(x): 60 """Compute RMS norm.""" 61 return np.linalg.norm(x) / x.size ** 0.5 62 63 64def select_initial_step(fun, t0, y0, f0, direction, order, rtol, atol): 65 """Empirically select a good initial step. 66 67 The algorithm is described in [1]_. 68 69 Parameters 70 ---------- 71 fun : callable 72 Right-hand side of the system. 73 t0 : float 74 Initial value of the independent variable. 75 y0 : ndarray, shape (n,) 76 Initial value of the dependent variable. 77 f0 : ndarray, shape (n,) 78 Initial value of the derivative, i.e., ``fun(t0, y0)``. 79 direction : float 80 Integration direction. 81 order : float 82 Error estimator order. It means that the error controlled by the 83 algorithm is proportional to ``step_size ** (order + 1)`. 84 rtol : float 85 Desired relative tolerance. 86 atol : float 87 Desired absolute tolerance. 88 89 Returns 90 ------- 91 h_abs : float 92 Absolute value of the suggested initial step. 93 94 References 95 ---------- 96 .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential 97 Equations I: Nonstiff Problems", Sec. II.4. 98 """ 99 if y0.size == 0: 100 return np.inf 101 102 scale = atol + np.abs(y0) * rtol 103 d0 = norm(y0 / scale) 104 d1 = norm(f0 / scale) 105 if d0 < 1e-5 or d1 < 1e-5: 106 h0 = 1e-6 107 else: 108 h0 = 0.01 * d0 / d1 109 110 y1 = y0 + h0 * direction * f0 111 f1 = fun(t0 + h0 * direction, y1) 112 d2 = norm((f1 - f0) / scale) / h0 113 114 if d1 <= 1e-15 and d2 <= 1e-15: 115 h1 = max(1e-6, h0 * 1e-3) 116 else: 117 h1 = (0.01 / max(d1, d2)) ** (1 / (order + 1)) 118 119 return min(100 * h0, h1) 120 121 122class OdeSolution: 123 """Continuous ODE solution. 124 125 It is organized as a collection of `DenseOutput` objects which represent 126 local interpolants. It provides an algorithm to select a right interpolant 127 for each given point. 128 129 The interpolants cover the range between `t_min` and `t_max` (see 130 Attributes below). Evaluation outside this interval is not forbidden, but 131 the accuracy is not guaranteed. 132 133 When evaluating at a breakpoint (one of the values in `ts`) a segment with 134 the lower index is selected. 135 136 Parameters 137 ---------- 138 ts : array_like, shape (n_segments + 1,) 139 Time instants between which local interpolants are defined. Must 140 be strictly increasing or decreasing (zero segment with two points is 141 also allowed). 142 interpolants : list of DenseOutput with n_segments elements 143 Local interpolants. An i-th interpolant is assumed to be defined 144 between ``ts[i]`` and ``ts[i + 1]``. 145 146 Attributes 147 ---------- 148 t_min, t_max : float 149 Time range of the interpolation. 150 """ 151 def __init__(self, ts, interpolants): 152 ts = np.asarray(ts) 153 d = np.diff(ts) 154 # The first case covers integration on zero segment. 155 if not ((ts.size == 2 and ts[0] == ts[-1]) 156 or np.all(d > 0) or np.all(d < 0)): 157 raise ValueError("`ts` must be strictly increasing or decreasing.") 158 159 self.n_segments = len(interpolants) 160 if ts.shape != (self.n_segments + 1,): 161 raise ValueError("Numbers of time stamps and interpolants " 162 "don't match.") 163 164 self.ts = ts 165 self.interpolants = interpolants 166 if ts[-1] >= ts[0]: 167 self.t_min = ts[0] 168 self.t_max = ts[-1] 169 self.ascending = True 170 self.ts_sorted = ts 171 else: 172 self.t_min = ts[-1] 173 self.t_max = ts[0] 174 self.ascending = False 175 self.ts_sorted = ts[::-1] 176 177 def _call_single(self, t): 178 # Here we preserve a certain symmetry that when t is in self.ts, 179 # then we prioritize a segment with a lower index. 180 if self.ascending: 181 ind = np.searchsorted(self.ts_sorted, t, side='left') 182 else: 183 ind = np.searchsorted(self.ts_sorted, t, side='right') 184 185 segment = min(max(ind - 1, 0), self.n_segments - 1) 186 if not self.ascending: 187 segment = self.n_segments - 1 - segment 188 189 return self.interpolants[segment](t) 190 191 def __call__(self, t): 192 """Evaluate the solution. 193 194 Parameters 195 ---------- 196 t : float or array_like with shape (n_points,) 197 Points to evaluate at. 198 199 Returns 200 ------- 201 y : ndarray, shape (n_states,) or (n_states, n_points) 202 Computed values. Shape depends on whether `t` is a scalar or a 203 1-D array. 204 """ 205 t = np.asarray(t) 206 207 if t.ndim == 0: 208 return self._call_single(t) 209 210 order = np.argsort(t) 211 reverse = np.empty_like(order) 212 reverse[order] = np.arange(order.shape[0]) 213 t_sorted = t[order] 214 215 # See comment in self._call_single. 216 if self.ascending: 217 segments = np.searchsorted(self.ts_sorted, t_sorted, side='left') 218 else: 219 segments = np.searchsorted(self.ts_sorted, t_sorted, side='right') 220 segments -= 1 221 segments[segments < 0] = 0 222 segments[segments > self.n_segments - 1] = self.n_segments - 1 223 if not self.ascending: 224 segments = self.n_segments - 1 - segments 225 226 ys = [] 227 group_start = 0 228 for segment, group in groupby(segments): 229 group_end = group_start + len(list(group)) 230 y = self.interpolants[segment](t_sorted[group_start:group_end]) 231 ys.append(y) 232 group_start = group_end 233 234 ys = np.hstack(ys) 235 ys = ys[:, reverse] 236 237 return ys 238 239 240NUM_JAC_DIFF_REJECT = EPS ** 0.875 241NUM_JAC_DIFF_SMALL = EPS ** 0.75 242NUM_JAC_DIFF_BIG = EPS ** 0.25 243NUM_JAC_MIN_FACTOR = 1e3 * EPS 244NUM_JAC_FACTOR_INCREASE = 10 245NUM_JAC_FACTOR_DECREASE = 0.1 246 247 248def num_jac(fun, t, y, f, threshold, factor, sparsity=None): 249 """Finite differences Jacobian approximation tailored for ODE solvers. 250 251 This function computes finite difference approximation to the Jacobian 252 matrix of `fun` with respect to `y` using forward differences. 253 The Jacobian matrix has shape (n, n) and its element (i, j) is equal to 254 ``d f_i / d y_j``. 255 256 A special feature of this function is the ability to correct the step 257 size from iteration to iteration. The main idea is to keep the finite 258 difference significantly separated from its round-off error which 259 approximately equals ``EPS * np.abs(f)``. It reduces a possibility of a 260 huge error and assures that the estimated derivative are reasonably close 261 to the true values (i.e., the finite difference approximation is at least 262 qualitatively reflects the structure of the true Jacobian). 263 264 Parameters 265 ---------- 266 fun : callable 267 Right-hand side of the system implemented in a vectorized fashion. 268 t : float 269 Current time. 270 y : ndarray, shape (n,) 271 Current state. 272 f : ndarray, shape (n,) 273 Value of the right hand side at (t, y). 274 threshold : float 275 Threshold for `y` value used for computing the step size as 276 ``factor * np.maximum(np.abs(y), threshold)``. Typically, the value of 277 absolute tolerance (atol) for a solver should be passed as `threshold`. 278 factor : ndarray with shape (n,) or None 279 Factor to use for computing the step size. Pass None for the very 280 evaluation, then use the value returned from this function. 281 sparsity : tuple (structure, groups) or None 282 Sparsity structure of the Jacobian, `structure` must be csc_matrix. 283 284 Returns 285 ------- 286 J : ndarray or csc_matrix, shape (n, n) 287 Jacobian matrix. 288 factor : ndarray, shape (n,) 289 Suggested `factor` for the next evaluation. 290 """ 291 y = np.asarray(y) 292 n = y.shape[0] 293 if n == 0: 294 return np.empty((0, 0)), factor 295 296 if factor is None: 297 factor = np.full(n, EPS ** 0.5) 298 else: 299 factor = factor.copy() 300 301 # Direct the step as ODE dictates, hoping that such a step won't lead to 302 # a problematic region. For complex ODEs it makes sense to use the real 303 # part of f as we use steps along real axis. 304 f_sign = 2 * (np.real(f) >= 0).astype(float) - 1 305 y_scale = f_sign * np.maximum(threshold, np.abs(y)) 306 h = (y + factor * y_scale) - y 307 308 # Make sure that the step is not 0 to start with. Not likely it will be 309 # executed often. 310 for i in np.nonzero(h == 0)[0]: 311 while h[i] == 0: 312 factor[i] *= 10 313 h[i] = (y[i] + factor[i] * y_scale[i]) - y[i] 314 315 if sparsity is None: 316 return _dense_num_jac(fun, t, y, f, h, factor, y_scale) 317 else: 318 structure, groups = sparsity 319 return _sparse_num_jac(fun, t, y, f, h, factor, y_scale, 320 structure, groups) 321 322 323def _dense_num_jac(fun, t, y, f, h, factor, y_scale): 324 n = y.shape[0] 325 h_vecs = np.diag(h) 326 f_new = fun(t, y[:, None] + h_vecs) 327 diff = f_new - f[:, None] 328 max_ind = np.argmax(np.abs(diff), axis=0) 329 r = np.arange(n) 330 max_diff = np.abs(diff[max_ind, r]) 331 scale = np.maximum(np.abs(f[max_ind]), np.abs(f_new[max_ind, r])) 332 333 diff_too_small = max_diff < NUM_JAC_DIFF_REJECT * scale 334 if np.any(diff_too_small): 335 ind, = np.nonzero(diff_too_small) 336 new_factor = NUM_JAC_FACTOR_INCREASE * factor[ind] 337 h_new = (y[ind] + new_factor * y_scale[ind]) - y[ind] 338 h_vecs[ind, ind] = h_new 339 f_new = fun(t, y[:, None] + h_vecs[:, ind]) 340 diff_new = f_new - f[:, None] 341 max_ind = np.argmax(np.abs(diff_new), axis=0) 342 r = np.arange(ind.shape[0]) 343 max_diff_new = np.abs(diff_new[max_ind, r]) 344 scale_new = np.maximum(np.abs(f[max_ind]), np.abs(f_new[max_ind, r])) 345 346 update = max_diff[ind] * scale_new < max_diff_new * scale[ind] 347 if np.any(update): 348 update, = np.nonzero(update) 349 update_ind = ind[update] 350 factor[update_ind] = new_factor[update] 351 h[update_ind] = h_new[update] 352 diff[:, update_ind] = diff_new[:, update] 353 scale[update_ind] = scale_new[update] 354 max_diff[update_ind] = max_diff_new[update] 355 356 diff /= h 357 358 factor[max_diff < NUM_JAC_DIFF_SMALL * scale] *= NUM_JAC_FACTOR_INCREASE 359 factor[max_diff > NUM_JAC_DIFF_BIG * scale] *= NUM_JAC_FACTOR_DECREASE 360 factor = np.maximum(factor, NUM_JAC_MIN_FACTOR) 361 362 return diff, factor 363 364 365def _sparse_num_jac(fun, t, y, f, h, factor, y_scale, structure, groups): 366 n = y.shape[0] 367 n_groups = np.max(groups) + 1 368 h_vecs = np.empty((n_groups, n)) 369 for group in range(n_groups): 370 e = np.equal(group, groups) 371 h_vecs[group] = h * e 372 h_vecs = h_vecs.T 373 374 f_new = fun(t, y[:, None] + h_vecs) 375 df = f_new - f[:, None] 376 377 i, j, _ = find(structure) 378 diff = coo_matrix((df[i, groups[j]], (i, j)), shape=(n, n)).tocsc() 379 max_ind = np.array(abs(diff).argmax(axis=0)).ravel() 380 r = np.arange(n) 381 max_diff = np.asarray(np.abs(diff[max_ind, r])).ravel() 382 scale = np.maximum(np.abs(f[max_ind]), 383 np.abs(f_new[max_ind, groups[r]])) 384 385 diff_too_small = max_diff < NUM_JAC_DIFF_REJECT * scale 386 if np.any(diff_too_small): 387 ind, = np.nonzero(diff_too_small) 388 new_factor = NUM_JAC_FACTOR_INCREASE * factor[ind] 389 h_new = (y[ind] + new_factor * y_scale[ind]) - y[ind] 390 h_new_all = np.zeros(n) 391 h_new_all[ind] = h_new 392 393 groups_unique = np.unique(groups[ind]) 394 groups_map = np.empty(n_groups, dtype=int) 395 h_vecs = np.empty((groups_unique.shape[0], n)) 396 for k, group in enumerate(groups_unique): 397 e = np.equal(group, groups) 398 h_vecs[k] = h_new_all * e 399 groups_map[group] = k 400 h_vecs = h_vecs.T 401 402 f_new = fun(t, y[:, None] + h_vecs) 403 df = f_new - f[:, None] 404 i, j, _ = find(structure[:, ind]) 405 diff_new = coo_matrix((df[i, groups_map[groups[ind[j]]]], 406 (i, j)), shape=(n, ind.shape[0])).tocsc() 407 408 max_ind_new = np.array(abs(diff_new).argmax(axis=0)).ravel() 409 r = np.arange(ind.shape[0]) 410 max_diff_new = np.asarray(np.abs(diff_new[max_ind_new, r])).ravel() 411 scale_new = np.maximum( 412 np.abs(f[max_ind_new]), 413 np.abs(f_new[max_ind_new, groups_map[groups[ind]]])) 414 415 update = max_diff[ind] * scale_new < max_diff_new * scale[ind] 416 if np.any(update): 417 update, = np.nonzero(update) 418 update_ind = ind[update] 419 factor[update_ind] = new_factor[update] 420 h[update_ind] = h_new[update] 421 diff[:, update_ind] = diff_new[:, update] 422 scale[update_ind] = scale_new[update] 423 max_diff[update_ind] = max_diff_new[update] 424 425 diff.data /= np.repeat(h, np.diff(diff.indptr)) 426 427 factor[max_diff < NUM_JAC_DIFF_SMALL * scale] *= NUM_JAC_FACTOR_INCREASE 428 factor[max_diff > NUM_JAC_DIFF_BIG * scale] *= NUM_JAC_FACTOR_DECREASE 429 factor = np.maximum(factor, NUM_JAC_MIN_FACTOR) 430 431 return diff, factor 432