Coverage for pySDC/projects/Resilience/vdp.py: 82%

1# script to run a van der Pol problem

2import numpy as np

3import matplotlib.pyplot as plt

5from pySDC.helpers.stats_helper import get_sorted, get_list_of_types

6from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol

7from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI

8from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity

9from pySDC.core.Errors import ProblemError, ConvergenceError

10from pySDC.projects.Resilience.hook import LogData, hook_collection

11from pySDC.projects.Resilience.strategies import merge_descriptions

12from pySDC.projects.Resilience.sweepers import generic_implicit_efficient

15def plot_step_sizes(stats, ax, e_em_key='error_embedded_estimate'):

16 """

17 Plot solution and step sizes to visualize the dynamics in the van der Pol equation.

19 Args:

20 stats (pySDC.stats): The stats object of the run

21 ax: Somewhere to plot

23 Returns:

24 None

25 """

27 # convert filtered statistics to list of iterations count, sorted by process

28 u = np.array([me[1][0] for me in get_sorted(stats, type='u', recomputed=False, sortby='time')])

29 p = np.array([me[1][1] for me in get_sorted(stats, type='u', recomputed=False, sortby='time')])

30 t = np.array([me[0] for me in get_sorted(stats, type='u', recomputed=False, sortby='time')])

32 e_em = np.array(get_sorted(stats, type=e_em_key, recomputed=False, sortby='time'))[:, 1]

33 dt = np.array(get_sorted(stats, type='dt', recomputed=False, sortby='time'))

34 restart = np.array(get_sorted(stats, type='restart', recomputed=None, sortby='time'))

36 ax.plot(t, u, label=r'$u$')

37 ax.plot(t, p, label=r'$p$')

39 dt_ax = ax.twinx()

40 dt_ax.plot(dt[:, 0], dt[:, 1], color='black')

41 dt_ax.plot(t, e_em, color='magenta')

42 dt_ax.set_yscale('log')

43 dt_ax.set_ylim((5e-10, 3e-1))

45 ax.plot([None], [None], label=r'$\Delta t$', color='black')

46 ax.plot([None], [None], label=r'$\epsilon_\mathrm{embedded}$', color='magenta')

47 ax.plot([None], [None], label='restart', color='grey', ls='-.')

49 for i in range(len(restart)):

50 if restart[i, 1] > 0:

51 ax.axvline(restart[i, 0], color='grey', ls='-.')

52 ax.legend(frameon=False)

54 ax.set_xlabel('time')

57def plot_avoid_restarts(stats, ax, avoid_restarts):

58 """

59 Make a plot that shows how many iterations where required to solve to a point in time in the simulation.

60 Also restarts are shown as vertical lines.

62 Args:

63 stats (pySDC.stats): The stats object of the run

64 ax: Somewhere to plot

65 avoid_restarts (bool): Whether the `avoid_restarts` option was set in order to choose a color

67 Returns:

68 None

69 """

70 sweeps = get_sorted(stats, type='sweeps', recomputed=None)

71 restarts = get_sorted(stats, type='restart', recomputed=None)

73 color = 'blue' if avoid_restarts else 'red'

74 ls = ':' if not avoid_restarts else '-.'

75 label = 'with' if avoid_restarts else 'without'

77 ax.plot([me[0] for me in sweeps], np.cumsum([me[1] for me in sweeps]), color=color, label=f'{label} avoid_restarts')

78 [ax.axvline(me[0], color=color, ls=ls) for me in restarts if me[1]]

80 ax.set_xlabel(r'$t$')

81 ax.set_ylabel(r'$k$')

82 ax.legend(frameon=False)

85def run_vdp(

86 custom_description=None,

87 num_procs=1,

88 Tend=10.0,

89 hook_class=LogData,

90 fault_stuff=None,

91 custom_controller_params=None,

92 use_MPI=False,

93 **kwargs,

94):

95 """

96 Run a van der Pol problem with default parameters.

98 Args:

99 custom_description (dict): Overwrite presets

100 num_procs (int): Number of steps for MSSDC

101 Tend (float): Time to integrate to

102 hook_class (pySDC.Hook): A hook to store data

103 fault_stuff (dict): A dictionary with information on how to add faults

104 custom_controller_params (dict): Overwrite presets

105 use_MPI (bool): Whether or not to use MPI

106

107 Returns:

108 dict: The stats object

109 controller: The controller

110 bool: If the code crashed

111 """

112

113 # initialize level parameters

114 level_params = {}

115 level_params['dt'] = 1e-2

116

117 # initialize sweeper parameters

118 sweeper_params = {}

119 sweeper_params['quad_type'] = 'RADAU-RIGHT'

120 sweeper_params['num_nodes'] = 3

121 sweeper_params['QI'] = 'LU'

122

123 problem_params = {

124 'mu': 5.0,

125 'newton_tol': 1e-9,

126 'newton_maxiter': 99,

127 'u0': np.array([2.0, 0.0]),

128 }

129

130 # initialize step parameters

131 step_params = {}

132 step_params['maxiter'] = 4

133

134 # initialize controller parameters

135 controller_params = {}

136 controller_params['logger_level'] = 30

137 controller_params['hook_class'] = hook_collection + (hook_class if type(hook_class) == list else [hook_class])

138 controller_params['mssdc_jac'] = False

139

140 if custom_controller_params is not None:

141 controller_params = {**controller_params, **custom_controller_params}

142

143 # fill description dictionary for easy step instantiation

144 description = {}

145 description['problem_class'] = vanderpol

146 description['problem_params'] = problem_params

147 description['sweeper_class'] = generic_implicit_efficient

148 description['sweeper_params'] = sweeper_params

149 description['level_params'] = level_params

150 description['step_params'] = step_params

151

152 if custom_description is not None:

153 description = merge_descriptions(description, custom_description)

154

155 # set time parameters

156 t0 = 0.0

157

158 # instantiate controller

159 if use_MPI:

160 from mpi4py import MPI

161 from pySDC.implementations.controller_classes.controller_MPI import controller_MPI

162

163 comm = kwargs.get('comm', MPI.COMM_WORLD)

164 controller = controller_MPI(controller_params=controller_params, description=description, comm=comm)

165

166 # get initial values on finest level

167 P = controller.S.levels[0].prob

168 uinit = P.u_exact(t0)

169 else:

170 controller = controller_nonMPI(

171 num_procs=num_procs, controller_params=controller_params, description=description

172 )

173

174 # get initial values on finest level

175 P = controller.MS[0].levels[0].prob

176 uinit = P.u_exact(t0)

177

178 # insert faults

179 if fault_stuff is not None:

180 from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults

181

182 prepare_controller_for_faults(controller, fault_stuff, {}, {})

183

184 # call main function to get things done...

185 crash = False

186 try:

187 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)

188 except (ProblemError, ConvergenceError, ZeroDivisionError) as e:

189 crash = True

190 print(f'Warning: Premature termination: {e}')

191 stats = controller.return_stats()

192

193 return stats, controller, crash

194

195

196def fetch_test_data(stats, comm=None, use_MPI=False):

197 """

198 Get data to perform tests on from stats

199

200 Args:

201 stats (pySDC.stats): The stats object of the run

202 comm (mpi4py.MPI.Comm): MPI communicator, or `None` for the non-MPI version

203 use_MPI (bool): Whether or not MPI was used when generating stats

204

205 Returns:

206 dict: Key values to perform tests on

207 """

208 types = ['error_embedded_estimate', 'restart', 'dt', 'sweeps', 'residual_post_step']

209 data = {}

210 for type in types:

211 if type not in get_list_of_types(stats):

212 raise ValueError(f"Can't read type \"{type}\" from stats, only got", get_list_of_types(stats))

213

214 data[type] = [

215 me[1] for me in get_sorted(stats, type=type, recomputed=None, sortby='time', comm=comm if use_MPI else None)

216 ]

217

218 # add time

219 data['time'] = [

220 me[0] for me in get_sorted(stats, type='u', recomputed=None, sortby='time', comm=comm if use_MPI else None)

221 ]

222 return data

223

224

225def check_if_tests_match(data_nonMPI, data_MPI):

226 """

227 Check if the data matches between MPI and nonMPI versions

228

229 Args:

230 data_nonMPI (dict): Key values to perform tests on obtained without MPI

231 data_MPI (dict): Key values to perform tests on obtained with MPI

232

233 Returns:

234 None

235 """

236 ops = [np.mean, np.min, np.max, len, sum]

237 for type in data_nonMPI.keys():

238 for op in ops:

239 val_nonMPI = op(data_nonMPI[type])

240 val_MPI = op(data_MPI[type])

241 assert np.isclose(val_nonMPI, val_MPI), (

242 f"Mismatch in operation {op.__name__} on type \"{type}\": with {data_MPI['size'][0]} ranks: "

243 f"nonMPI: {val_nonMPI}, MPI: {val_MPI}"

244 )

245 print(f'Passed with {data_MPI["size"][0]} ranks')

246

247

248def mpi_vs_nonMPI(MPI_ready, comm):

249 """

250 Check if MPI and non-MPI versions give the same output.

251

252 Args:

253 MPI_ready (bool): Whether or not we can use MPI at all

254 comm (mpi4py.MPI.Comm): MPI communicator

255

256 Returns:

257 None

258 """

259 if MPI_ready:

260 size = comm.size

261 rank = comm.rank

262 use_MPI = [True, False]

263 else:

264 size = 1

265 rank = 0

266 use_MPI = [False, False]

267

268 if rank == 0:

269 print(f"Running with {size} ranks")

270

271 custom_description = {'convergence_controllers': {}}

272 custom_description['convergence_controllers'][Adaptivity] = {'e_tol': 1e-7, 'avoid_restarts': False}

273

274 data = [{}, {}]

275

276 for i in range(2):

277 if use_MPI[i] or rank == 0:

278 stats, controller, Tend = run_vdp(

279 custom_description=custom_description,

280 num_procs=size,

281 use_MPI=use_MPI[i],

282 Tend=1.0,

283 comm=comm,

284 )

285 data[i] = fetch_test_data(stats, comm, use_MPI=use_MPI[i])

286 data[i]['size'] = [size]

287

288 if rank == 0:

289 check_if_tests_match(data[1], data[0])

290

291

292def check_adaptivity_with_avoid_restarts(comm=None, size=1):

293 """

294 Make a test if adaptivity with the option to avoid restarts based on a contraction factor estimate works as

295 expected.

296 To this end, we run the same test of the van der Pol equation twice with the only difference being this option

297 turned off or on.

298 We recorded how many iterations we expect to avoid by avoiding restarts and check against this value.

299 Also makes a figure comparing the number of iterations over time.

300

301 In principle there is an option to test MSSDC here, but this is only preliminary and needs to be checked further.

302

303 Args:

304 comm (mpi4py.MPI.Comm): MPI communicator, or `None` for the non-MPI version

305 size (int): Number of steps for MSSDC, is overridden by communicator size if applicable

306

307 Returns:

308 None

309 """

310 fig, ax = plt.subplots()

311 custom_description = {'convergence_controllers': {}, 'level_params': {'dt': 1.0e-2}}

312 custom_controller_params = {'all_to_done': False}

313 results = {'e': {}, 'sweeps': {}, 'restarts': {}}

314 size = comm.size if comm is not None else size

315

316 for avoid_restarts in [True, False]:

317 custom_description['convergence_controllers'][Adaptivity] = {'e_tol': 1e-7, 'avoid_restarts': avoid_restarts}

318 stats, controller, Tend = run_vdp(

319 custom_description=custom_description,

320 num_procs=size,

321 use_MPI=comm is not None,

322 custom_controller_params=custom_controller_params,

323 Tend=10.0e0,

324 comm=comm,

325 )

326 plot_avoid_restarts(stats, ax, avoid_restarts)

327

328 # check error

329 u = get_sorted(stats, type='u', recomputed=False)[-1]

330 if comm is None:

331 u_exact = controller.MS[0].levels[0].prob.u_exact(t=u[0])

332 else:

333 u_exact = controller.S.levels[0].prob.u_exact(t=u[0])

334 results['e'][avoid_restarts] = abs(u[1] - u_exact)

335

336 # check iteration counts

337 results['sweeps'][avoid_restarts] = sum(

338 [me[1] for me in get_sorted(stats, type='sweeps', recomputed=None, comm=comm)]

339 )

340 results['restarts'][avoid_restarts] = sum([me[1] for me in get_sorted(stats, type='restart', comm=comm)])

341

342 fig.tight_layout()

343 fig.savefig(f'data/vdp-{size}procs{"-use_MPI" if comm is not None else ""}-avoid_restarts.png')

344

345 assert np.isclose(results['e'][True], results['e'][False], rtol=5.0), (

346 'Errors don\'t match with avoid_restarts and without, got '

347 f'{results["e"][True]:.2e} and {results["e"][False]:.2e}'

348 )

349 if size == 1:

350 assert results['sweeps'][True] - results['sweeps'][False] == 1301 - 1344, (

351 '{Expected to save 43 iterations '

352 f"with avoid_restarts, got {results['sweeps'][False] - results['sweeps'][True]}"

353 )

354 assert results['restarts'][True] - results['restarts'][False] == 0 - 10, (

355 '{Expected to save 10 restarts '

356 f"with avoid_restarts, got {results['restarts'][False] - results['restarts'][True]}"

357 )

358 print('Passed avoid_restarts tests with 1 process')

359 if size == 4:

360 assert results['sweeps'][True] - results['sweeps'][False] == 2916 - 3008, (

361 '{Expected to save 92 iterations '

362 f"with avoid_restarts, got {results['sweeps'][False] - results['sweeps'][True]}"

363 )

364 assert results['restarts'][True] - results['restarts'][False] == 0 - 18, (

365 '{Expected to save 18 restarts '

366 f"with avoid_restarts, got {results['restarts'][False] - results['restarts'][True]}"

367 )

368 print('Passed avoid_restarts tests with 4 processes')

369

370

371def check_step_size_limiter(size=4, comm=None):

372 """

373 Check the step size limiter convergence controller.

374 First we run without step size limits and then enforce limits that are slightly above and below what the usual

375 limits. Then we run again and see if we exceed the limits.

376

377 Args:

378 size (int): Number of steps for MSSDC

379 comm (mpi4py.MPI.Comm): MPI communicator, or `None` for the non-MPI version

380

381 Returns:

382 None

383 """

384 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter

385

386 custom_description = {'convergence_controllers': {}, 'level_params': {'dt': 1.0e-2}}

387 expect = {}

388 params = {'e_tol': 1e-6}

389

390 for limit_step_sizes in [False, True]:

391 if limit_step_sizes:

392 params['dt_max'] = expect['dt_max'] * 0.9

393 params['dt_min'] = np.inf

394 params['dt_slope_max'] = expect['dt_slope_max'] * 0.9

395 params['dt_slope_min'] = expect['dt_slope_min'] * 1.1

396 custom_description['convergence_controllers'][StepSizeLimiter] = {'dt_min': expect['dt_min'] * 1.1}

397 else:

398 for k in ['dt_max', 'dt_min', 'dt_slope_max', 'dt_slope_min']:

399 params.pop(k, None)

400 custom_description['convergence_controllers'].pop(StepSizeLimiter, None)

401

402 custom_description['convergence_controllers'][Adaptivity] = params

403 stats, controller, Tend = run_vdp(

404 custom_description=custom_description,

405 num_procs=size,

406 use_MPI=comm is not None,

407 Tend=5.0e0,

408 comm=comm,

409 )

410

411 # plot the step sizes

412 dt = get_sorted(stats, type='dt', recomputed=None, comm=comm)

413

414 # make sure that the convergence controllers are only added once

415 convergence_controller_classes = [type(me) for me in controller.convergence_controllers]

416 for c in convergence_controller_classes:

417 assert convergence_controller_classes.count(c) == 1, f'Convergence controller {c} added multiple times'

418

419 dt_numpy = np.array([me[1] for me in dt])

420 if not limit_step_sizes:

421 expect['dt_max'] = max(dt_numpy)

422 expect['dt_min'] = min(dt_numpy)

423 expect['dt_slope_max'] = max(dt_numpy[:-2] / dt_numpy[1:-1])

424 expect['dt_slope_min'] = min(dt_numpy[:-2] / dt_numpy[1:-1])

425 else:

426 dt_max = max(dt_numpy)

427 dt_min = min(dt_numpy[size:-size]) # The first and last step might fall below the limits

428 dt_slope_max = max(dt_numpy[:-2] / dt_numpy[1:-1])

429 dt_slope_min = min(dt_numpy[:-2] / dt_numpy[1:-1])

430 assert (

431 dt_max <= expect['dt_max']

432 ), f"Exceeded maximum allowed step size! Got {dt_max:.4e}, allowed {params['dt_max']:.4e}."

433 assert (

434 dt_min >= expect['dt_min']

435 ), f"Exceeded minimum allowed step size! Got {dt_min:.4e}, allowed {params['dt_min']:.4e}."

436 assert (

437 dt_slope_max <= expect['dt_slope_max']

438 ), f"Exceeded maximum allowed step size slope! Got {dt_slope_max:.4e}, allowed {params['dt_slope_max']:.4e}."

439 assert (

440 dt_slope_min >= expect['dt_slope_min']

441 ), f"Exceeded minimum allowed step size slope! Got {dt_slope_min:.4e}, allowed {params['dt_slope_min']:.4e}."

442

443 assert (

444 dt_slope_max <= expect['dt_slope_max']

445 ), f"Exceeded maximum allowed step size slope! Got {dt_slope_max:.4e}, allowed {params['dt_slope_max']:.4e}."

446 assert (

447 dt_slope_min >= expect['dt_slope_min']

448 ), f"Exceeded minimum allowed step size slope! Got {dt_slope_min:.4e}, allowed {params['dt_slope_min']:.4e}."

449

450 if comm is None:

451 print(f'Passed step size limiter test with {size} ranks in nonMPI implementation')

452 else:

453 if comm.rank == 0:

454 print(f'Passed step size limiter test with {size} ranks in MPI implementation')

455

456

457def interpolation_stuff(): # pragma: no cover

458 """

459 Plot interpolation vdp with interpolation after a restart and compare it to other modes of adaptivity.

460 """

461 from pySDC.implementations.convergence_controller_classes.interpolate_between_restarts import (

462 InterpolateBetweenRestarts,

463 )

464 from pySDC.implementations.hooks.log_errors import LogLocalErrorPostStep

465 from pySDC.implementations.hooks.log_work import LogWork

466 from pySDC.helpers.plot_helper import figsize_by_journal

467

468 fig, axs = plt.subplots(4, 1, figsize=figsize_by_journal('Springer_Numerical_Algorithms', 1.0, 1.0), sharex=True)

469 restart_ax = axs[2].twinx()

470

471 colors = ['black', 'red', 'blue']

472 labels = ['interpolate', 'regular', 'keep iterating']

473

474 for i in range(3):

475 convergence_controllers = {

476 Adaptivity: {'e_tol': 1e-7, 'dt_max': 9.0e-1},

477 }

478 if i == 0:

479 convergence_controllers[InterpolateBetweenRestarts] = {}

480 if i == 2:

481 convergence_controllers[Adaptivity]['avoid_restarts'] = True

482

483 problem_params = {

484 'mu': 5,

485 }

486

487 sweeper_params = {

488 'QI': 'LU',

489 }

490

491 custom_description = {

492 'convergence_controllers': convergence_controllers,

493 'problem_params': problem_params,

494 'sweeper_params': sweeper_params,

495 }

496

497 stats, controller, _ = run_vdp(

498 custom_description=custom_description,

499 hook_class=[LogLocalErrorPostStep, LogData, LogWork] + hook_collection,

500 )

501

502 k = get_sorted(stats, type='work_newton')

503 restarts = get_sorted(stats, type='restart')

504 u = get_sorted(stats, type='u', recomputed=False)

505 e_loc = get_sorted(stats, type='e_local_post_step', recomputed=False)

506 dt = get_sorted(stats, type='dt', recomputed=False)

507

508 axs[0].plot([me[0] for me in u], [me[1][1] for me in u], color=colors[i], label=labels[i])

509 axs[1].plot([me[0] for me in e_loc], [me[1] for me in e_loc], color=colors[i])

510 axs[2].plot([me[0] for me in k], np.cumsum([me[1] for me in k]), color=colors[i])

511 restart_ax.plot([me[0] for me in restarts], np.cumsum([me[1] for me in restarts]), color=colors[i], ls='--')

512 axs[3].plot([me[0] for me in dt], [me[1] for me in dt], color=colors[i])

513

514 for ax in [axs[1], axs[3]]:

515 ax.set_yscale('log')

516 axs[0].set_ylabel(r'$u$')

517 axs[1].set_ylabel(r'$e_\mathrm{local}$')

518 axs[2].set_ylabel(r'Newton iterations')

519 restart_ax.set_ylabel(r'restarts (dashed)')

520 axs[3].set_ylabel(r'$\Delta t$')

521 axs[3].set_xlabel(r'$t$')

522 axs[0].legend(frameon=False)

523 fig.tight_layout()

524 plt.show()

525

526

527if __name__ == "__main__":

528 import sys

529

530 try:

531 from mpi4py import MPI

532

533 MPI_ready = True

534 comm = MPI.COMM_WORLD

535 size = comm.size

536 except ModuleNotFoundError:

537 MPI_ready = False

538 comm = None

539 size = 1

540

541 if len(sys.argv) == 1:

542 mpi_vs_nonMPI(MPI_ready, comm)

543 check_step_size_limiter(size, comm)

544

545 if size == 1:

546 check_adaptivity_with_avoid_restarts(comm=None, size=1)

547

548 elif 'mpi_vs_nonMPI' in sys.argv:

549 mpi_vs_nonMPI(MPI_ready, comm)

550 elif 'check_step_size_limiter' in sys.argv:

551 check_step_size_limiter(MPI_ready, comm)

552 elif 'check_adaptivity_with_avoid_restarts' and size == 1:

553 check_adaptivity_with_avoid_restarts(comm=None, size=1)

554 else:

555 raise NotImplementedError('Your test is not implemented!')