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

129 }

130

131 # initialize step parameters

132 step_params = {}

133 step_params['maxiter'] = 4

134

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

140

141 if custom_controller_params is not None:

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

143

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

152

153 if custom_description is not None:

154 description = merge_descriptions(description, custom_description)

155

156 # set time parameters

157 t0 = 0.0

158

159 # instantiate controller

160 if use_MPI:

161 from mpi4py import MPI

162 from pySDC.implementations.controller_classes.controller_MPI import controller_MPI

163

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

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

166

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 )

174

175 # get initial values on finest level

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

177 uinit = P.u_exact(t0)

178

179 # insert faults

180 if fault_stuff is not None:

181 from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults

182

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

184

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

193

194 return stats, controller, crash

195

196

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

198 """

199 Get data to perform tests on from stats

200

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

205

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

214

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 ]

218

219 # add time

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

224

225

226def check_if_tests_match(data_nonMPI, data_MPI):

227 """

228 Check if the data matches between MPI and nonMPI versions

229

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

233

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')

247

248

249def mpi_vs_nonMPI(MPI_ready, comm):

250 """

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

252

253 Args:

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

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

256

257 Returns:

258 None

259 """

260 if MPI_ready:

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]

268

269 if rank == 0:

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

271

272 custom_description = {'convergence_controllers': {}}

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

274

275 data = [{}, {}]

276

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]

288

289 if rank == 0:

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

291

292

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

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.

301

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

303

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

307

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

316

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)

328

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)

336

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

342

343 fig.tight_layout()

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

345

346 # measured: 1.07e-8 with and 8.36e-9 without avoiding restarts; the step sequences are deterministic (the

347 # iteration counts below are compared exactly), so a factor 1.5 leaves room without accepting worse steps

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

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

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

351 )

352 if size == 1:

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

354 '{Expected to save 43 iterations '

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

356 )

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

358 '{Expected to save 10 restarts '

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

360 )

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

362 if size == 4:

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

364 '{Expected to save 92 iterations '

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

366 )

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

368 '{Expected to save 18 restarts '

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

370 )

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

372

373

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

375 """

376 Check the step size limiter convergence controller.

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

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

379

380 Args:

381 size (int): Number of steps for MSSDC

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

383

384 Returns:

385 None

386 """

387 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter

388

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

390 expect = {}

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

392

393 for limit_step_sizes in [False, True]:

394 if limit_step_sizes:

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

396 params['dt_min'] = np.inf

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

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

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

400 else:

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

402 params.pop(k, None)

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

404

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

406 stats, controller, Tend = run_vdp(

407 custom_description=custom_description,

408 num_procs=size,

409 use_MPI=comm is not None,

410 Tend=5.0e0,

411 comm=comm,

412 )

413

414 # plot the step sizes

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

416

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

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

419 for c in convergence_controller_classes:

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

421

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

423 if not limit_step_sizes:

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

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

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

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

428 else:

429 dt_max = max(dt_numpy)

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

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

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

433 assert (

434 dt_max <= expect['dt_max']

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

436 assert (

437 dt_min >= expect['dt_min']

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

439 assert (

440 dt_slope_max <= expect['dt_slope_max']

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

442 assert (

443 dt_slope_min >= expect['dt_slope_min']

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

445

446 assert (

447 dt_slope_max <= expect['dt_slope_max']

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

449 assert (

450 dt_slope_min >= expect['dt_slope_min']

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

452

453 if comm is None:

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

455 else:

456 if comm.rank == 0:

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

458

459

460def interpolation_stuff(): # pragma: no cover

461 """

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

463 """

464 from pySDC.implementations.convergence_controller_classes.interpolate_between_restarts import (

465 InterpolateBetweenRestarts,

466 )

467 from pySDC.implementations.hooks.log_errors import LogLocalErrorPostStep

468 from pySDC.implementations.hooks.log_work import LogWork

469 from pySDC.helpers.plot_helper import figsize_by_journal

470

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

472 restart_ax = axs[2].twinx()

473

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

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

476

477 for i in range(3):

478 convergence_controllers = {

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

480 }

481 if i == 0:

482 convergence_controllers[InterpolateBetweenRestarts] = {}

483 if i == 2:

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

485

486 problem_params = {

487 'mu': 5,

488 }

489

490 sweeper_params = {

491 'QI': 'LU',

492 }

493

494 custom_description = {

495 'convergence_controllers': convergence_controllers,

496 'problem_params': problem_params,

497 'sweeper_params': sweeper_params,

498 }

499

500 stats, controller, _ = run_vdp(

501 custom_description=custom_description,

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

503 )

504

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

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

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

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

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

510

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

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

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

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

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

516

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

518 ax.set_yscale('log')

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

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

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

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

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

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

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

526 fig.tight_layout()

527 plt.show()

528

529

530if __name__ == "__main__":

531 import sys

532

533 try:

534 from mpi4py import MPI

535

536 MPI_ready = True

537 comm = MPI.COMM_WORLD

538 size = comm.size

539 except ModuleNotFoundError:

540 MPI_ready = False

541 comm = None

542 size = 1

543

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

545 mpi_vs_nonMPI(MPI_ready, comm)

546 check_step_size_limiter(size, comm)

547

548 if size == 1:

549 check_adaptivity_with_avoid_restarts(comm=None, size=1)

550

551 elif 'mpi_vs_nonMPI' in sys.argv:

552 mpi_vs_nonMPI(MPI_ready, comm)

553 elif 'check_step_size_limiter' in sys.argv:

554 # `size` is the number of parallel steps; this passed `MPI_ready`, so it always ran on one

555 check_step_size_limiter(size, comm)

556 elif 'check_adaptivity_with_avoid_restarts' in sys.argv and size == 1:

557 # the `in sys.argv` was missing, and a non-empty string is truthy, so every unrecognised

558 # argument landed here instead of raising below

559 check_adaptivity_with_avoid_restarts(comm=None, size=1)

560 else:

561 raise NotImplementedError(f'Cannot run {sys.argv[1:]} on {size} rank(s)!')