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

195 statements

, created at 2024-09-19 09:13 +0000

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

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

107 Returns:

108 dict: The stats object

109 controller: The controller

110 bool: If the code crashed

111 """

113 # initialize level parameters

114 level_params = {}

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

117 # initialize sweeper parameters

118 sweeper_params = {}

120 sweeper_params['num_nodes'] = 3

121 sweeper_params['QI'] = 'LU'

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 'relative_tolerance': True,

129 }

131 # initialize step parameters

132 step_params = {}

133 step_params['maxiter'] = 4

135 # initialize controller parameters

136 controller_params = {}

137 controller_params['logger_level'] = 30

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

139 controller_params['mssdc_jac'] = False

141 if custom_controller_params is not None:

142 controller_params = {**controller_params, **custom_controller_params}

144 # fill description dictionary for easy step instantiation

145 description = {}

146 description['problem_class'] = vanderpol

147 description['problem_params'] = problem_params

148 description['sweeper_class'] = generic_implicit_efficient

149 description['sweeper_params'] = sweeper_params

150 description['level_params'] = level_params

151 description['step_params'] = step_params

153 if custom_description is not None:

154 description = merge_descriptions(description, custom_description)

156 # set time parameters

157 t0 = 0.0

159 # instantiate controller

160 if use_MPI:

161 from mpi4py import MPI

162 from pySDC.implementations.controller_classes.controller_MPI import controller_MPI

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

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

167 # get initial values on finest level

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

169 uinit = P.u_exact(t0)

170 else:

171 controller = controller_nonMPI(

172 num_procs=num_procs, controller_params=controller_params, description=description

173 )

175 # get initial values on finest level

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

177 uinit = P.u_exact(t0)

179 # insert faults

180 if fault_stuff is not None:

181 from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults

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

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

186 crash = False

187 try:

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

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

190 crash = True

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

192 stats = controller.return_stats()

194 return stats, controller, crash

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

198 """

199 Get data to perform tests on from stats

201 Args:

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

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

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

206 Returns:

207 dict: Key values to perform tests on

208 """

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

210 data = {}

211 for type in types:

212 if type not in get_list_of_types(stats):

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

215 data[type] = [

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

217 ]

220 data['time'] = [

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

222 ]

223 return data

226def check_if_tests_match(data_nonMPI, data_MPI):

227 """

228 Check if the data matches between MPI and nonMPI versions

230 Args:

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

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

234 Returns:

235 None

236 """

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

238 for type in data_nonMPI.keys():

239 for op in ops:

240 val_nonMPI = op(data_nonMPI[type])

241 val_MPI = op(data_MPI[type])

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

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

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

245 )

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

250 """

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

253 Args:

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

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

257 Returns:

258 None

259 """

261 size = comm.size

262 rank = comm.rank

263 use_MPI = [True, False]

264 else:

265 size = 1

266 rank = 0

267 use_MPI = [False, False]

269 if rank == 0:

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

272 custom_description = {'convergence_controllers': {}}

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

275 data = [{}, {}]

277 for i in range(2):

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

279 stats, controller, Tend = run_vdp(

280 custom_description=custom_description,

281 num_procs=size,

282 use_MPI=use_MPI[i],

283 Tend=1.0,

284 comm=comm,

285 )

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

287 data[i]['size'] = [size]

289 if rank == 0:

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

294 """

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

296 expected.

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

298 turned off or on.

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

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

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

304 Args:

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

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

308 Returns:

309 None

310 """

311 fig, ax = plt.subplots()

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

313 custom_controller_params = {'all_to_done': False}

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

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

317 for avoid_restarts in [True, False]:

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

319 stats, controller, Tend = run_vdp(

320 custom_description=custom_description,

321 num_procs=size,

322 use_MPI=comm is not None,

323 custom_controller_params=custom_controller_params,

324 Tend=10.0e0,

325 comm=comm,

326 )

327 plot_avoid_restarts(stats, ax, avoid_restarts)

329 # check error

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

331 if comm is None:

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

333 else:

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

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

337 # check iteration counts

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

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

340 )

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

343 fig.tight_layout()

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

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

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

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

349 )

350 if size == 1:

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

352 '{Expected to save 43 iterations '

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

354 )

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

356 '{Expected to save 10 restarts '

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

358 )

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

360 if size == 4:

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

362 '{Expected to save 92 iterations '

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

364 )

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

366 '{Expected to save 18 restarts '

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

368 )

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

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

373 """

374 Check the step size limiter convergence controller.

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

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

378 Args:

379 size (int): Number of steps for MSSDC

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

382 Returns:

383 None

384 """

385 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter

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

388 expect = {}

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

391 for limit_step_sizes in [False, True]:

392 if limit_step_sizes:

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

394 params['dt_min'] = np.inf

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

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

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

398 else:

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

400 params.pop(k, None)

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

404 stats, controller, Tend = run_vdp(

405 custom_description=custom_description,

406 num_procs=size,

407 use_MPI=comm is not None,

408 Tend=5.0e0,

409 comm=comm,

410 )

412 # plot the step sizes

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

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

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

417 for c in convergence_controller_classes:

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

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

421 if not limit_step_sizes:

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

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

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

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

426 else:

427 dt_max = max(dt_numpy)

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

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

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

431 assert (

432 dt_max <= expect['dt_max']

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

434 assert (

435 dt_min >= expect['dt_min']

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

437 assert (

438 dt_slope_max <= expect['dt_slope_max']

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

440 assert (

441 dt_slope_min >= expect['dt_slope_min']

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

444 assert (

445 dt_slope_max <= expect['dt_slope_max']

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

447 assert (

448 dt_slope_min >= expect['dt_slope_min']

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

451 if comm is None:

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

453 else:

454 if comm.rank == 0:

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

458def interpolation_stuff(): # pragma: no cover

459 """

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

461 """

462 from pySDC.implementations.convergence_controller_classes.interpolate_between_restarts import (

463 InterpolateBetweenRestarts,

464 )

465 from pySDC.implementations.hooks.log_errors import LogLocalErrorPostStep

466 from pySDC.implementations.hooks.log_work import LogWork

467 from pySDC.helpers.plot_helper import figsize_by_journal

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

470 restart_ax = axs[2].twinx()

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

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

475 for i in range(3):

476 convergence_controllers = {

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

478 }

479 if i == 0:

480 convergence_controllers[InterpolateBetweenRestarts] = {}

481 if i == 2:

484 problem_params = {

485 'mu': 5,

486 }

488 sweeper_params = {

489 'QI': 'LU',

490 }

492 custom_description = {

493 'convergence_controllers': convergence_controllers,

494 'problem_params': problem_params,

495 'sweeper_params': sweeper_params,

496 }

498 stats, controller, _ = run_vdp(

499 custom_description=custom_description,

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

501 )

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

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

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

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

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

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

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

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

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

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

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

516 ax.set_yscale('log')

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

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

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

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

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

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

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

524 fig.tight_layout()

525 plt.show()

528if __name__ == "__main__":

529 import sys

531 try:

532 from mpi4py import MPI

535 comm = MPI.COMM_WORLD

536 size = comm.size

537 except ModuleNotFoundError:

539 comm = None

540 size = 1

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

544 check_step_size_limiter(size, comm)

546 if size == 1:

549 elif 'mpi_vs_nonMPI' in sys.argv: