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

157 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-09-09 14:59 +0000

1import numpy as np 

2from scipy.special import factorial 

3 

4from pySDC.core.convergence_controller 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.add_status_variable_to_level('error_extrapolation_estimate') 

88 

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

90 """ 

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

92 

93 Args: 

94 controller (pySDC.Controller): The controller 

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

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

97 

98 Returns: 

99 bool: Whether the parameters are compatible 

100 str: Error message 

101 """ 

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

103 return ( 

104 False, 

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

106has to be smaller than 0!", 

107 ) 

108 

109 if controller.params.mssdc_jac: 

110 return ( 

111 False, 

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

113el multistep mode!", 

114 ) 

115 

116 return True, "" 

117 

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

119 """ 

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

121 node on the finest level at the moment. 

122 

123 Args: 

124 S (pySDC.Step): The current step 

125 

126 Returns: 

127 None 

128 """ 

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

130 if None in self.prev.t: 

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

132 else: 

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

134 

135 # figure out how to store the right hand side 

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

137 if type(f) == imex_mesh: 

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

139 elif type(f) == mesh: 

140 self.prev.f[oldest_val] = f 

141 else: 

142 raise DataError( 

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

144 , extrapolation based error estimate only\ 

145 works with types imex_mesh and mesh" 

146 ) 

147 

148 # store the rest of the values 

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

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

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

152 

153 return None 

154 

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

156 """ 

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

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

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

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

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

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

163 memory at the time of evaluation. 

164 

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

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

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

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

169 problems. 

170 

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

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

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

174 

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

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

177 iteration. 

178 

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

180 

181 Args: 

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

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

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

185 

186 Returns: 

187 None 

188 """ 

189 

190 # prepare A matrix 

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

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

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

194 inv_facs = 1.0 / factorial(j) 

195 

196 # get the steps backwards from the point of evaluation 

197 idx = np.argsort(t) 

198 steps_from_now = t[idx] - t_eval 

199 

200 # fill A matrix 

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

202 # Taylor expansions of the solutions 

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

204 

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

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

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

208 ) 

209 

210 # prepare rhs 

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

212 b[0] = 1.0 

213 

214 # solve linear system for the coefficients 

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

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

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

218 

219 # determine prefactor 

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

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

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

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

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

225 

226 return None 

227 

228 

229class EstimateExtrapolationErrorNonMPI(EstimateExtrapolationErrorBase): 

230 """ 

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

232 """ 

233 

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

235 """ 

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

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

238 

239 Args: 

240 controller (pySDC.Controller): The controller 

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

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

243 

244 Returns: 

245 dict: Updated parameters with default values 

246 """ 

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

248 

249 non_mpi_defaults = { 

250 "no_storage": False, 

251 } 

252 

253 return {**non_mpi_defaults, **default_params} 

254 

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

256 """ 

257 Initialize storage variables. 

258 

259 Args: 

260 controller (pySDC.controller): The controller 

261 

262 Returns: 

263 None 

264 """ 

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

266 

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

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

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

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

271 

272 return None 

273 

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

275 """ 

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

277 - Compute the error estimate 

278 - Compute the coefficients if needed 

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

280 

281 Args: 

282 controller (pySDC.Controller): The controller 

283 S (pySDC.Step): The current step 

284 

285 Returns: 

286 None 

287 """ 

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

289 t_eval = S.time + S.dt 

290 

291 # compute the extrapolation coefficients if needed 

292 if ( 

293 ( 

294 None in self.coeff.u 

295 or self.params.use_adaptivity 

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

297 ) 

298 and None not in self.prev.t 

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

300 ): 

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

302 

303 # compute the error if we can 

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

305 self.get_extrapolated_error(S) 

306 

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

308 if self.params.no_storage: 

309 self.store_values(S) 

310 

311 return None 

312 

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

314 """ 

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

316 need to store all the stuff that we can. 

317 

318 Args: 

319 controller (pySDC.Controller): The controller 

320 S (pySDC.step): The current step 

321 size (int): Number of ranks 

322 time (float): The current time 

323 Tend (float): The final time 

324 MS (list): Active steps 

325 

326 Returns: 

327 None 

328 """ 

329 

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

331 if self.params.no_storage: 

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

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

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

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

336 

337 else: 

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

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

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

341 

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

343 if restart_at > S.status.slot: 

344 self.store_values(S) 

345 

346 return None 

347 

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

349 """ 

350 Combine values from previous steps to extrapolate. 

351 

352 Args: 

353 S (pySDC.Step): The current step 

354 

355 Returns: 

356 dtype_u: The extrapolated solution 

357 """ 

358 if len(S.levels) > 1: 

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

360 

361 # prepare variables 

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

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

364 

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

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

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

368 else: 

369 idx_step = max(idx) + 1 

370 

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

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

373 

374 # do the extrapolation by summing everything up 

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

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

377 

378 return u_ex 

379 

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

381 """ 

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

383 solution obtained by the time marching scheme. 

384 

385 Args: 

386 S (pySDC.Step): The current step 

387 

388 Returns: 

389 None 

390 """ 

391 u_ex = self.get_extrapolated_solution(S) 

392 if u_ex is not None: 

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

394 else: 

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

396 

397 

398class EstimateExtrapolationErrorWithinQ(EstimateExtrapolationErrorBase): 

399 """ 

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

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

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

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

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

405 of the number of nodes. 

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

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

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

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

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

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

412 """ 

413 

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

415 """ 

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

417 to be converged. 

418 

419 Args: 

420 controller (pySDC.Controller): The controller 

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

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

423 

424 Returns: 

425 dict: Updated parameters with default values 

426 """ 

427 from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence 

428 

429 num_nodes = description['sweeper_params']['num_nodes'] 

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

431 if self.comm: 

432 from mpi4py import MPI 

433 

434 self.sum = MPI.SUM 

435 

436 self.check_convergence = CheckConvergence.check_convergence 

437 

438 default_params = { 

439 'Taylor_order': 2 * num_nodes, 

440 'n': num_nodes, 

441 'recompute_coefficients': False, 

442 **params, 

443 } 

444 

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

446 

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

448 """ 

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

450 

451 Args: 

452 controller (pySDC.Controller): The controller 

453 S (pySDC.Step): The current step 

454 

455 Returns: 

456 None 

457 """ 

458 

459 if not self.check_convergence(S): 

460 return None 

461 

462 lvl = S.levels[0] 

463 

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

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

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

467 

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

469 self.params.Taylor_order = len(nodes) 

470 self.params.n = len(nodes) 

471 

472 # compute the extrapolation coefficients 

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

474 self.get_extrapolation_coefficients(nodes, dts, t_eval) 

475 

476 # compute the extrapolated solution 

477 if lvl.f[0] is None: 

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

479 

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

481 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)] 

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

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

484 else: 

485 raise DataError( 

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

487 , extrapolation based error estimate only\ 

488 works with types imex_mesh and mesh" 

489 ) 

490 

491 # compute the error with the weighted sum 

492 if self.comm: 

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

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

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

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

497 else: 

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

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

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

501 

502 # store the error 

503 if self.comm: 

504 error = ( 

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

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

507 else None 

508 ) 

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

510 else: 

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

512 return None