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

370

371

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.

377

378 Args:

379 size (int): Number of steps for MSSDC

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

381

382 Returns:

383 None

384 """

385 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter

386

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

388 expect = {}

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

390

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)

402

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

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 )

411

412 # plot the step sizes

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

414

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'

419

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}."

443

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}."

450

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

456

457

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

468

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

471

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

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

474

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:

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

483

484 problem_params = {

485 'mu': 5,

486 }

487

488 sweeper_params = {

489 'QI': 'LU',

490 }

491

492 custom_description = {

493 'convergence_controllers': convergence_controllers,

494 'problem_params': problem_params,

495 'sweeper_params': sweeper_params,

496 }

497

498 stats, controller, _ = run_vdp(

499 custom_description=custom_description,

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

501 )

502

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)

508

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

514

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

526

527

528if __name__ == "__main__":

529 import sys

530

531 try:

532 from mpi4py import MPI

533

534 MPI_ready = True

535 comm = MPI.COMM_WORLD

536 size = comm.size

537 except ModuleNotFoundError:

538 MPI_ready = False

539 comm = None

540 size = 1

541

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

543 mpi_vs_nonMPI(MPI_ready, comm)

544 check_step_size_limiter(size, comm)

545

546 if size == 1:

547 check_adaptivity_with_avoid_restarts(comm=None, size=1)

548

549 elif 'mpi_vs_nonMPI' in sys.argv:

550 mpi_vs_nonMPI(MPI_ready, comm)

551 elif 'check_step_size_limiter' in sys.argv:

552 check_step_size_limiter(MPI_ready, comm)

553 elif 'check_adaptivity_with_avoid_restarts' and size == 1:

554 check_adaptivity_with_avoid_restarts(comm=None, size=1)

555 else:

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