Coverage for pySDC/implementations/convergence_controller_classes/estimate_extrapolation_error.py: 98%

167 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-29 09:02 +0000

1import numpy as np 

2from scipy.special import factorial 

3 

4from pySDC.core.ConvergenceController import ConvergenceController, Status 

5from pySDC.core.Errors import DataError 

6from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh 

7from pySDC.implementations.hooks.log_extrapolated_error_estimate import LogExtrapolationErrorEstimate 

8 

9 

10class EstimateExtrapolationErrorBase(ConvergenceController): 

11 """ 

12 Abstract base class for extrapolated error estimates 

13 ---------------------------------------------------- 

14 This error estimate extrapolates a solution based on Taylor expansions using solutions of previous time steps. 

15 In particular, child classes need to implement how to make these solutions available, which works differently for 

16 MPI and non-MPI versions. 

17 """ 

18 

19 def __init__(self, controller, params, description, **kwargs): 

20 """ 

21 Initialization routine 

22 

23 Args: 

24 controller (pySDC.Controller): The controller 

25 params (dict): The params passed for this specific convergence controller 

26 description (dict): The description object used to instantiate the controller 

27 """ 

28 self.prev = Status(["t", "u", "f", "dt"]) # store solutions etc. of previous steps here 

29 self.coeff = Status(["u", "f", "prefactor"]) # store coefficients for extrapolation here 

30 super().__init__(controller, params, description) 

31 controller.add_hook(LogExtrapolationErrorEstimate) 

32 

33 def setup(self, controller, params, description, **kwargs): 

34 """ 

35 The extrapolation based method requires storage of previous values of u, f, t and dt and also requires solving 

36 a linear system of equations to compute the Taylor expansion finite difference style. Here, all variables are 

37 initialized which are needed for this process. 

38 

39 Args: 

40 controller (pySDC.Controller): The controller 

41 params (dict): The params passed for this specific convergence controller 

42 description (dict): The description object used to instantiate the controller 

43 

44 Returns: 

45 dict: Updated parameters with default values 

46 """ 

47 

48 from pySDC.implementations.convergence_controller_classes.hotrod import HotRod 

49 from pySDC.implementations.convergence_controller_classes.adaptivity import ( 

50 Adaptivity, 

51 ) 

52 

53 default_params = { 

54 "control_order": -75, 

55 "use_adaptivity": True in [me == Adaptivity for me in description.get("convergence_controllers", {})], 

56 "use_HotRod": True in [me == HotRod for me in description.get("convergence_controllers", {})], 

57 "order_time_marching": description["step_params"]["maxiter"], 

58 } 

59 

60 new_params = {**default_params, **super().setup(controller, params, description, **kwargs)} 

61 

62 # Do a sufficiently high order Taylor expansion 

63 new_params["Taylor_order"] = new_params["order_time_marching"] + 2 

64 

65 # Estimate and store values from this iteration 

66 new_params["estimate_iter"] = new_params["order_time_marching"] - (1 if new_params["use_HotRod"] else 0) 

67 

68 # Store n values. Since we store u and f, we need only half of each (the +1 is for rounding) 

69 new_params["n"] = (new_params["Taylor_order"] + 1) // 2 

70 new_params["n_per_proc"] = new_params["n"] * 1 

71 

72 return new_params 

73 

74 def setup_status_variables(self, controller, **kwargs): 

75 """ 

76 Initialize coefficient variables and add variable to the levels for extrapolated error 

77 

78 Args: 

79 controller (pySDC.controller): The controller 

80 

81 Returns: 

82 None 

83 """ 

84 self.coeff.u = [None] * self.params.n 

85 self.coeff.f = [0.0] * self.params.n 

86 

87 self.reset_status_variables(controller, **kwargs) 

88 return None 

89 

90 def reset_status_variables(self, controller, **kwargs): 

91 """ 

92 Add variable for extrapolated error 

93 

94 Args: 

95 controller (pySDC.Controller): The controller 

96 

97 Returns: 

98 None 

99 """ 

100 if 'comm' in kwargs.keys(): 

101 steps = [controller.S] 

102 else: 

103 if 'active_slots' in kwargs.keys(): 

104 steps = [controller.MS[i] for i in kwargs['active_slots']] 

105 else: 

106 steps = controller.MS 

107 where = ["levels", "status"] 

108 for S in steps: 

109 self.add_variable(S, name='error_extrapolation_estimate', where=where, init=None) 

110 

111 def check_parameters(self, controller, params, description, **kwargs): 

112 """ 

113 Check whether parameters are compatible with whatever assumptions went into the step size functions etc. 

114 

115 Args: 

116 controller (pySDC.Controller): The controller 

117 params (dict): The params passed for this specific convergence controller 

118 description (dict): The description object used to instantiate the controller 

119 

120 Returns: 

121 bool: Whether the parameters are compatible 

122 str: Error message 

123 """ 

124 if description["step_params"].get("restol", -1.0) >= 0: 

125 return ( 

126 False, 

127 "Extrapolation error needs constant order in time and hence restol in the step parameters \ 

128has to be smaller than 0!", 

129 ) 

130 

131 if controller.params.mssdc_jac: 

132 return ( 

133 False, 

134 "Extrapolation error estimator needs the same order on all steps, please activate Gauss-Seid\ 

135el multistep mode!", 

136 ) 

137 

138 return True, "" 

139 

140 def store_values(self, S, **kwargs): 

141 """ 

142 Store the required attributes of the step to do the extrapolation. We only care about the last collocation 

143 node on the finest level at the moment. 

144 

145 Args: 

146 S (pySDC.Step): The current step 

147 

148 Returns: 

149 None 

150 """ 

151 # figure out which values are to be replaced by the new ones 

152 if None in self.prev.t: 

153 oldest_val = len(self.prev.t) - len(self.prev.t[self.prev.t == [None]]) 

154 else: 

155 oldest_val = np.argmin(self.prev.t) 

156 

157 # figure out how to store the right hand side 

158 f = S.levels[0].f[-1] 

159 if type(f) == imex_mesh: 

160 self.prev.f[oldest_val] = f.impl + f.expl 

161 elif type(f) == mesh: 

162 self.prev.f[oldest_val] = f 

163 else: 

164 raise DataError( 

165 f"Unable to store f from datatype {type(f)} \ 

166 , extrapolation based error estimate only\ 

167 works with types imex_mesh and mesh" 

168 ) 

169 

170 # store the rest of the values 

171 self.prev.u[oldest_val] = S.levels[0].u[-1] 

172 self.prev.t[oldest_val] = S.time + S.dt 

173 self.prev.dt[oldest_val] = S.dt 

174 

175 return None 

176 

177 def get_extrapolation_coefficients(self, t, dt, t_eval): 

178 """ 

179 This function solves a linear system where in the matrix A, the row index reflects the order of the derivative 

180 in the Taylor expansion and the column index reflects the particular step and whether its u or f from that 

181 step. The vector b on the other hand, contains a 1 in the first entry and zeros elsewhere, since we want to 

182 compute the value itself and all the derivatives should vanish after combining the Taylor expansions. This 

183 works to the order of the number of rows and since we want a square matrix for solving, we need the same amount 

184 of columns, which determines the memory overhead, since it is equal to the solutions / rhs that we need in 

185 memory at the time of evaluation. 

186 

187 This is enough to get the extrapolated solution, but if we want to compute the local error, we have to compute 

188 a prefactor. This is based on error accumulation between steps (first step's solution is exact plus 1 LTE, 

189 second solution is exact plus 2 LTE and so on), which can be computed for adaptive step sizes as well. However, 

190 this is only true for linear problems, which means we expect the error estimate to work less well for non-linear 

191 problems. 

192 

193 Since only time differences are important for computing the coefficients, we need to compute this only once when 

194 using constant step sizes. When we allow the step size to change, however, we need to recompute this in every 

195 step, which is activated by the `use_adaptivity` parameter. 

196 

197 Solving for the coefficients requires solving a dense linear system of equations. The number of unknowns is 

198 equal to the order of the Taylor expansion, so this step should be cheap compared to the solves in each SDC 

199 iteration. 

200 

201 The function stores the computed coefficients in the `self.coeff` variables. 

202 

203 Args: 

204 t (list): The list of times at which we have solutions available 

205 dt (list): The step sizes used for computing these solutions (needed for the prefactor) 

206 t_eval (float): The time we want to extrapolate to 

207 

208 Returns: 

209 None 

210 """ 

211 

212 # prepare A matrix 

213 A = np.zeros((self.params.Taylor_order, self.params.Taylor_order)) 

214 A[0, 0 : self.params.n] = 1.0 

215 j = np.arange(self.params.Taylor_order) 

216 inv_facs = 1.0 / factorial(j) 

217 

218 # get the steps backwards from the point of evaluation 

219 idx = np.argsort(t) 

220 steps_from_now = t[idx] - t_eval 

221 

222 # fill A matrix 

223 for i in range(1, self.params.Taylor_order): 

224 # Taylor expansions of the solutions 

225 A[i, : self.params.n] = steps_from_now ** j[i] * inv_facs[i] 

226 

227 # Taylor expansions of the first derivatives a.k.a. right hand side evaluations 

228 A[i, self.params.n : self.params.Taylor_order] = ( 

229 steps_from_now[2 * self.params.n - self.params.Taylor_order :] ** (j[i] - 1) * inv_facs[i - 1] 

230 ) 

231 

232 # prepare rhs 

233 b = np.zeros(self.params.Taylor_order) 

234 b[0] = 1.0 

235 

236 # solve linear system for the coefficients 

237 coeff = np.linalg.solve(A, b) 

238 self.coeff.u = coeff[: self.params.n] 

239 self.coeff.f[self.params.n * 2 - self.params.Taylor_order :] = coeff[self.params.n : self.params.Taylor_order] 

240 

241 # determine prefactor 

242 step_size_ratios = abs(dt[len(dt) - len(self.coeff.u) :] / dt[-1]) ** (self.params.Taylor_order - 1) 

243 inv_prefactor = -sum(step_size_ratios[1:]) - 1.0 

244 for i in range(len(self.coeff.u)): 

245 inv_prefactor += sum(step_size_ratios[1 : i + 1]) * self.coeff.u[i] 

246 self.coeff.prefactor = 1.0 / abs(inv_prefactor) 

247 

248 return None 

249 

250 

251class EstimateExtrapolationErrorNonMPI(EstimateExtrapolationErrorBase): 

252 """ 

253 Implementation of the extrapolation error estimate for the non-MPI controller. 

254 """ 

255 

256 def setup(self, controller, params, description, **kwargs): 

257 """ 

258 Add a no parameter 'no_storage' which decides whether the standard or the no-memory-overhead version is run, 

259 where only values are used for extrapolation which are in memory of other processes 

260 

261 Args: 

262 controller (pySDC.Controller): The controller 

263 params (dict): The params passed for this specific convergence controller 

264 description (dict): The description object used to instantiate the controller 

265 

266 Returns: 

267 dict: Updated parameters with default values 

268 """ 

269 default_params = super().setup(controller, params, description) 

270 

271 non_mpi_defaults = { 

272 "no_storage": False, 

273 } 

274 

275 return {**non_mpi_defaults, **default_params} 

276 

277 def setup_status_variables(self, controller, **kwargs): 

278 """ 

279 Initialize storage variables. 

280 

281 Args: 

282 controller (pySDC.controller): The controller 

283 

284 Returns: 

285 None 

286 """ 

287 super().setup_status_variables(controller, **kwargs) 

288 

289 self.prev.t = np.array([None] * self.params.n) 

290 self.prev.dt = np.array([None] * self.params.n) 

291 self.prev.u = [None] * self.params.n 

292 self.prev.f = [None] * self.params.n 

293 

294 return None 

295 

296 def post_iteration_processing(self, controller, S, **kwargs): 

297 """ 

298 We perform three key operations here in the last iteration: 

299 - Compute the error estimate 

300 - Compute the coefficients if needed 

301 - Store the values of the step if we pretend not to for the no-memory overhead version 

302 

303 Args: 

304 controller (pySDC.Controller): The controller 

305 S (pySDC.Step): The current step 

306 

307 Returns: 

308 None 

309 """ 

310 if S.status.iter == self.params.estimate_iter: 

311 t_eval = S.time + S.dt 

312 

313 # compute the extrapolation coefficients if needed 

314 if ( 

315 ( 

316 None in self.coeff.u 

317 or self.params.use_adaptivity 

318 or (not self.params.no_storage and S.status.time_size > 1) 

319 ) 

320 and None not in self.prev.t 

321 and t_eval > max(self.prev.t) 

322 ): 

323 self.get_extrapolation_coefficients(self.prev.t, self.prev.dt, t_eval) 

324 

325 # compute the error if we can 

326 if None not in self.coeff.u and None not in self.prev.t: 

327 self.get_extrapolated_error(S) 

328 

329 # store the solution and pretend we didn't because in the non MPI version we take a few shortcuts 

330 if self.params.no_storage: 

331 self.store_values(S) 

332 

333 return None 

334 

335 def prepare_next_block(self, controller, S, size, time, Tend, MS, **kwargs): 

336 """ 

337 If the no-memory-overhead version is used, we need to delete stuff that shouldn't be available. Otherwise, we 

338 need to store all the stuff that we can. 

339 

340 Args: 

341 controller (pySDC.Controller): The controller 

342 S (pySDC.step): The current step 

343 size (int): Number of ranks 

344 time (float): The current time 

345 Tend (float): The final time 

346 MS (list): Active steps 

347 

348 Returns: 

349 None 

350 """ 

351 

352 # delete values that should not be available in the next step 

353 if self.params.no_storage: 

354 self.prev.t = np.array([None] * self.params.n) 

355 self.prev.dt = np.array([None] * self.params.n) 

356 self.prev.u = [None] * self.params.n 

357 self.prev.f = [None] * self.params.n 

358 

359 else: 

360 # decide where we need to restart to store everything up to that point 

361 restarts = [S.status.restart for S in MS] 

362 restart_at = np.where(restarts)[0][0] if True in restarts else len(MS) 

363 

364 # store values in the current block that don't need restarting 

365 if restart_at > S.status.slot: 

366 self.store_values(S) 

367 

368 return None 

369 

370 def get_extrapolated_solution(self, S, **kwargs): 

371 """ 

372 Combine values from previous steps to extrapolate. 

373 

374 Args: 

375 S (pySDC.Step): The current step 

376 

377 Returns: 

378 dtype_u: The extrapolated solution 

379 """ 

380 if len(S.levels) > 1: 

381 raise NotImplementedError("Extrapolated estimate only works on the finest level for now") 

382 

383 # prepare variables 

384 u_ex = S.levels[0].u[-1] * 0.0 

385 idx = np.argsort(self.prev.t) 

386 

387 # see if we have a solution for the current step already stored 

388 if (abs(S.time + S.dt - self.prev.t) < 10.0 * np.finfo(float).eps).any(): 

389 idx_step = idx[np.argmin(abs(self.prev.t - S.time - S.dt))] 

390 else: 

391 idx_step = max(idx) + 1 

392 

393 # make a mask of all the steps we want to include in the extrapolation 

394 mask = np.logical_and(idx < idx_step, idx >= idx_step - self.params.n) 

395 

396 # do the extrapolation by summing everything up 

397 for i in range(self.params.n): 

398 u_ex += self.coeff.u[i] * self.prev.u[idx[mask][i]] + self.coeff.f[i] * self.prev.f[idx[mask][i]] 

399 

400 return u_ex 

401 

402 def get_extrapolated_error(self, S, **kwargs): 

403 """ 

404 The extrapolation estimate combines values of u and f from multiple steps to extrapolate and compare to the 

405 solution obtained by the time marching scheme. 

406 

407 Args: 

408 S (pySDC.Step): The current step 

409 

410 Returns: 

411 None 

412 """ 

413 u_ex = self.get_extrapolated_solution(S) 

414 if u_ex is not None: 

415 S.levels[0].status.error_extrapolation_estimate = abs(u_ex - S.levels[0].u[-1]) * self.coeff.prefactor 

416 else: 

417 S.levels[0].status.error_extrapolation_estimate = None 

418 

419 

420class EstimateExtrapolationErrorWithinQ(EstimateExtrapolationErrorBase): 

421 """ 

422 This convergence controller estimates the local error based on comparing the SDC solution to an extrapolated 

423 solution within the quadrature matrix. Collocation methods compute a high order solution from a linear combination 

424 of solutions at intermediate time points. While the intermediate solutions (a.k.a. stages) don't share the order of 

425 accuracy with the solution at the end of the interval, for SDC we know that the order is equal to the number of 

426 nodes + 1 (locally). This is because the solution to the collocation problem is a polynomial approximation of order 

427 of the number of nodes. 

428 That means we can do a Taylor expansion around the end point of the interval to higher order and after cancelling 

429 terms just like we are used to with the extrapolation based error estimate across multiple steps, we get an error 

430 estimate that is of the order of accuracy of the stages. 

431 This can be used for adaptivity, for instance, with the nice property that it doesn't matter how we arrived at the 

432 converged collocation solution, as long as we did. We don't rely on knowing the order of accuracy after every sweep, 

433 only after convergence of the collocation problem has been achieved, which we can check from the residual. 

434 """ 

435 

436 def setup(self, controller, params, description, **kwargs): 

437 """ 

438 We need this convergence controller to become active after the check for convergence, because we need the step 

439 to be converged. 

440 

441 Args: 

442 controller (pySDC.Controller): The controller 

443 params (dict): The params passed for this specific convergence controller 

444 description (dict): The description object used to instantiate the controller 

445 

446 Returns: 

447 dict: Updated parameters with default values 

448 """ 

449 from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence 

450 

451 num_nodes = description['sweeper_params']['num_nodes'] 

452 self.comm = description['sweeper_params'].get('comm', None) 

453 if self.comm: 

454 from mpi4py import MPI 

455 

456 self.sum = MPI.SUM 

457 

458 self.check_convergence = CheckConvergence.check_convergence 

459 

460 default_params = { 

461 'Taylor_order': 2 * num_nodes, 

462 'n': num_nodes, 

463 'recompute_coefficients': False, 

464 **params, 

465 } 

466 

467 return {**super().setup(controller, params, description, **kwargs), **default_params} 

468 

469 def post_iteration_processing(self, controller, S, **kwargs): 

470 """ 

471 Compute the extrapolated error estimate here if the step is converged. 

472 

473 Args: 

474 controller (pySDC.Controller): The controller 

475 S (pySDC.Step): The current step 

476 

477 Returns: 

478 None 

479 """ 

480 

481 if not self.check_convergence(S): 

482 return None 

483 

484 lvl = S.levels[0] 

485 

486 nodes_ = lvl.sweep.coll.nodes * S.dt 

487 nodes = S.time + np.append(0, nodes_[:-1]) 

488 t_eval = S.time + nodes_[-1] 

489 

490 dts = np.append(nodes_[0], nodes_[1:] - nodes_[:-1]) 

491 self.params.Taylor_order = len(nodes) 

492 self.params.n = len(nodes) 

493 

494 # compute the extrapolation coefficients 

495 if None in self.coeff.u or self.params.recompute_coefficients: 

496 self.get_extrapolation_coefficients(nodes, dts, t_eval) 

497 

498 # compute the extrapolated solution 

499 if lvl.f[0] is None: 

500 lvl.f[0] = lvl.prob.eval_f(lvl.u[0], lvl.time) 

501 

502 if type(lvl.f[0]) == imex_mesh: 

503 f = [lvl.f[i].impl + lvl.f[i].expl if self.coeff.f[i] and lvl.f[i] else 0.0 for i in range(len(lvl.f) - 1)] 

504 elif type(lvl.f[0]) == mesh: 

505 f = [lvl.f[i] if self.coeff.f[i] else 0.0 for i in range(len(lvl.f) - 1)] 

506 else: 

507 raise DataError( 

508 f"Unable to store f from datatype {type(lvl.f[0])} \ 

509 , extrapolation based error estimate only\ 

510 works with types imex_mesh and mesh" 

511 ) 

512 

513 # compute the error with the weighted sum 

514 if self.comm: 

515 idx = (self.comm.rank + 1) % self.comm.size 

516 sendbuf = self.coeff.u[idx] * lvl.u[idx] + self.coeff.f[idx] * f[idx] 

517 u_ex = lvl.prob.dtype_u(lvl.prob.init, val=0.0) if self.comm.rank == self.comm.size - 1 else None 

518 self.comm.Reduce(sendbuf, u_ex, op=self.sum, root=self.comm.size - 1) 

519 else: 

520 u_ex = lvl.prob.dtype_u(lvl.prob.init, val=0.0) 

521 for i in range(self.params.n): 

522 u_ex += self.coeff.u[i] * lvl.u[i] + self.coeff.f[i] * f[i] 

523 

524 # store the error 

525 if self.comm: 

526 error = ( 

527 abs(u_ex - lvl.u[self.comm.rank + 1]) * self.coeff.prefactor 

528 if self.comm.rank == self.comm.size - 1 

529 else None 

530 ) 

531 lvl.status.error_extrapolation_estimate = self.comm.bcast(error, root=self.comm.size - 1) 

532 else: 

533 lvl.status.error_extrapolation_estimate = abs(u_ex - lvl.u[-1]) * self.coeff.prefactor 

534 return None