Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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(object): 

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