Coverage for pySDC/projects/Resilience/work_precision.py: 0%

348 statements  

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

1from mpi4py import MPI 

2import numpy as np 

3import matplotlib.pyplot as plt 

4import pickle 

5import logging 

6from time import perf_counter 

7import copy 

8 

9from pySDC.projects.Resilience.strategies import merge_descriptions 

10from pySDC.projects.Resilience.Lorenz import run_Lorenz 

11from pySDC.projects.Resilience.vdp import run_vdp 

12from pySDC.projects.Resilience.Schroedinger import run_Schroedinger 

13from pySDC.projects.Resilience.quench import run_quench 

14from pySDC.projects.Resilience.AC import run_AC 

15 

16from pySDC.helpers.stats_helper import get_sorted, filter_stats 

17from pySDC.helpers.plot_helper import setup_mpl, figsize_by_journal 

18 

19setup_mpl(reset=True) 

20LOGGER_LEVEL = 25 

21LOG_TO_FILE = False 

22 

23logging.getLogger('matplotlib.texmanager').setLevel(90) 

24 

25 

26def std_log(x): 

27 return np.std(np.log(x)) 

28 

29 

30MAPPINGS = { 

31 'e_global': ('e_global_post_run', max, False), 

32 'e_global_rel': ('e_global_rel_post_run', max, False), 

33 't': ('timing_run', max, False), 

34 # 'e_local_max': ('e_local_post_step', max, False), 

35 'k_SDC': ('k', sum, None), 

36 'k_SDC_no_restart': ('k', sum, False), 

37 'k_Newton': ('work_newton', sum, None), 

38 'k_linear': ('work_linear', sum, None), 

39 'k_Newton_no_restart': ('work_newton', sum, False), 

40 'k_rhs': ('work_rhs', sum, None), 

41 'num_steps': ('dt', len, None), 

42 'restart': ('restart', sum, None), 

43 'dt_mean': ('dt', np.mean, False), 

44 'dt_max': ('dt', max, False), 

45 'dt_min': ('dt', min, False), 

46 'dt_sigma': ('dt', std_log, False), 

47 'e_embedded_max': ('error_embedded_estimate', max, False), 

48 'u0_increment_max': ('u0_increment', max, None), 

49 'u0_increment_mean': ('u0_increment', np.mean, None), 

50 'u0_increment_max_no_restart': ('u0_increment', max, False), 

51 'u0_increment_mean_no_restart': ('u0_increment', np.mean, False), 

52} 

53 

54logger = logging.getLogger('WorkPrecision') 

55logger.setLevel(LOGGER_LEVEL) 

56 

57 

58def get_forbidden_combinations(problem, strategy, **kwargs): 

59 """ 

60 Check if the combination of strategy and problem is forbidden 

61 

62 Args: 

63 problem (function): A problem to run 

64 strategy (Strategy): SDC strategy 

65 """ 

66 if strategy.name == 'ERK': 

67 if problem.__name__ in ['run_quench', 'run_Schroedinger', 'run_AC']: 

68 return True 

69 

70 return False 

71 

72 

73def single_run( 

74 problem, 

75 strategy, 

76 data, 

77 custom_description, 

78 num_procs=1, 

79 comm_world=None, 

80 problem_args=None, 

81 hooks=None, 

82 Tend=None, 

83 num_procs_sweeper=1, 

84): 

85 """ 

86 Make a single run of a particular problem with a certain strategy. 

87 

88 Args: 

89 problem (function): A problem to run 

90 strategy (Strategy): SDC strategy 

91 data (dict): Put the results in here 

92 custom_description (dict): Overwrite presets 

93 num_procs (int): Number of processes for the time communicator 

94 comm_world (mpi4py.MPI.Intracomm): Communicator that is available for the entire script 

95 hooks (list): List of additional hooks 

96 num_procs_sweeper (int): Number of processes for the sweeper 

97 

98 Returns: 

99 dict: Stats generated by the run 

100 """ 

101 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun 

102 from pySDC.implementations.hooks.log_work import LogWork 

103 from pySDC.projects.Resilience.hook import LogData 

104 

105 hooks = hooks if hooks else [] 

106 

107 t_last = perf_counter() 

108 

109 num_procs_tot = num_procs * num_procs_sweeper 

110 comm = comm_world.Split(comm_world.rank < num_procs_tot) 

111 if comm_world.rank >= num_procs_tot: 

112 comm.Free() 

113 return None 

114 

115 # make communicators for time and sweepers 

116 comm_time = comm.Split(comm.rank // num_procs) 

117 comm_sweep = comm.Split(comm_time.rank) 

118 

119 if comm_time.size < num_procs: 

120 raise Exception(f'Need at least {num_procs*num_procs_sweeper} processes, got only {comm.size}') 

121 

122 strategy_description = strategy.get_custom_description(problem, num_procs) 

123 description = merge_descriptions(strategy_description, custom_description) 

124 if comm_sweep.size > 1: 

125 description['sweeper_params']['comm'] = comm_sweep 

126 

127 controller_params = { 

128 'logger_level': LOGGER_LEVEL, 

129 'log_to_file': LOG_TO_FILE, 

130 'fname': 'out.txt', 

131 **strategy.get_controller_params(), 

132 } 

133 problem_args = {} if problem_args is None else problem_args 

134 

135 stats, controller, crash = problem( 

136 custom_description=description, 

137 Tend=strategy.get_Tend(problem, num_procs) if Tend is None else Tend, 

138 hook_class=[LogData, LogWork, LogGlobalErrorPostRun] + hooks, 

139 custom_controller_params=controller_params, 

140 use_MPI=True, 

141 comm=comm_time, 

142 **problem_args, 

143 ) 

144 

145 t_now = perf_counter() 

146 logger.debug(f'Finished run in {t_now - t_last:.2e} s') 

147 t_last = perf_counter() 

148 

149 # record all the metrics 

150 stats_all = filter_stats(stats, comm=comm_sweep) 

151 comm_sweep.Free() 

152 

153 for key, mapping in MAPPINGS.items(): 

154 if crash: 

155 data[key] += [np.nan] 

156 continue 

157 me = get_sorted(stats_all, comm=comm_time, type=mapping[0], recomputed=mapping[2]) 

158 if len(me) == 0: 

159 data[key] += [np.nan] 

160 else: 

161 data[key] += [mapping[1]([you[1] for you in me])] 

162 

163 t_now = perf_counter() 

164 logger.debug(f'Recorded all data after {t_now - t_last:.2e} s') 

165 t_last = perf_counter() 

166 

167 comm_time.Free() 

168 comm.Free() 

169 return stats 

170 

171 

172def get_parameter(dictionary, where): 

173 """ 

174 Get a parameter at a certain position in a dictionary of dictionaries. 

175 

176 Args: 

177 dictionary (dict): The dictionary 

178 where (list): The list of keys leading to the value you want 

179 

180 Returns: 

181 The value of the dictionary 

182 """ 

183 if len(where) == 1: 

184 return dictionary[where[0]] 

185 else: 

186 return get_parameter(dictionary[where[0]], where[1:]) 

187 

188 

189def set_parameter(dictionary, where, parameter): 

190 """ 

191 Set a parameter at a certain position in a dictionary of dictionaries 

192 

193 Args: 

194 dictionary (dict): The dictionary 

195 where (list): The list of keys leading to the value you want to set 

196 parameter: Whatever you want to set the parameter to 

197 

198 Returns: 

199 None 

200 """ 

201 if len(where) == 1: 

202 dictionary[where[0]] = parameter 

203 else: 

204 set_parameter(dictionary[where[0]], where[1:], parameter) 

205 

206 

207def get_path(problem, strategy, num_procs, handle='', base_path='data/work_precision', num_procs_sweeper=1, mode=''): 

208 """ 

209 Get the path to a certain data. 

210 

211 Args: 

212 problem (function): A problem to run 

213 strategy (Strategy): SDC strategy 

214 num_procs (int): Number of processes for the time communicator 

215 handle (str): The name of the configuration 

216 base_path (str): Some path where all the files are stored 

217 num_procs_sweeper (int): Number of processes for the sweeper 

218 mode (str): The mode this was generated for 

219 

220 Returns: 

221 str: The path to the data you are looking for 

222 """ 

223 return f'{base_path}/{problem.__name__}-{strategy.__class__.__name__}-{handle}{"-wp" if handle else "wp"}-{num_procs}-{num_procs_sweeper}procs-{mode}.pickle' 

224 

225 

226def record_work_precision( 

227 problem, 

228 strategy, 

229 num_procs=1, 

230 custom_description=None, 

231 handle='', 

232 runs=1, 

233 comm_world=None, 

234 problem_args=None, 

235 param_range=None, 

236 Tend=None, 

237 hooks=None, 

238 num_procs_sweeper=1, 

239 mode='', 

240): 

241 """ 

242 Run problem with strategy and record the cost parameters. 

243 

244 Args: 

245 problem (function): A problem to run 

246 strategy (Strategy): SDC strategy 

247 num_procs (int): Number of processes for the time communicator 

248 custom_description (dict): Overwrite presets 

249 handle (str): The name of the configuration 

250 runs (int): Number of runs you want to do 

251 comm_world (mpi4py.MPI.Intracomm): Communicator that is available for the entire script 

252 num_procs_sweeper (int): Number of processes for the sweeper 

253 

254 Returns: 

255 None 

256 """ 

257 if get_forbidden_combinations(problem, strategy): 

258 return None 

259 

260 data = {} 

261 

262 # prepare precision parameters 

263 param = strategy.precision_parameter 

264 description = merge_descriptions( 

265 strategy.get_custom_description(problem, num_procs), 

266 {} if custom_description is None else custom_description, 

267 ) 

268 if param == 'e_tol': 

269 power = 10.0 

270 set_parameter(description, strategy.precision_parameter_loc[:-1] + ['dt_min'], 0) 

271 exponents = [-3, -2, -1, 0, 1, 2, 3][::-1] 

272 if problem.__name__ == 'run_vdp': 

273 if type(strategy).__name__ in ["AdaptivityPolynomialError"]: 

274 exponents = [0, 1, 2, 3, 5][::-1] 

275 else: 

276 exponents = [-3, -2, -1, 0, 0.2, 0.8, 1][::-1] 

277 elif param == 'dt': 

278 power = 2.0 

279 exponents = [-1, 0, 1, 2, 3][::-1] 

280 elif param == 'restol': 

281 power = 10.0 

282 exponents = [-2, -1, 0, 1, 2, 3] 

283 if problem.__name__ == 'run_vdp': 

284 exponents = [-4, -3, -2, -1, 0, 1] 

285 else: 

286 raise NotImplementedError(f"I don't know how to get default value for parameter \"{param}\"") 

287 

288 where = strategy.precision_parameter_loc 

289 default = get_parameter(description, where) 

290 param_range = [default * power**i for i in exponents] if param_range is None else param_range 

291 

292 if problem.__name__ == 'run_quench': 

293 if param == 'restol': 

294 param_range = [1e-5, 1e-6, 1e-7, 1e-8, 1e-9] 

295 elif param == 'dt': 

296 param_range = [1.25, 2.5, 5.0, 10.0, 20.0][::-1] 

297 

298 # run multiple times with different parameters 

299 for i in range(len(param_range)): 

300 set_parameter(description, where, param_range[i]) 

301 

302 data[param_range[i]] = {key: [] for key in MAPPINGS.keys()} 

303 data[param_range[i]]['param'] = [param_range[i]] 

304 data[param_range[i]][param] = [param_range[i]] 

305 

306 description = merge_descriptions( 

307 descA=description, descB=strategy.get_description_for_tolerance(problem=problem, param=param_range[i]) 

308 ) 

309 for _j in range(runs): 

310 if comm_world.rank == 0: 

311 logger.log( 

312 24, 

313 f'Starting: {problem.__name__}: {strategy} {handle} {num_procs}-{num_procs_sweeper} procs, {param}={param_range[i]:.2e}', 

314 ) 

315 single_run( 

316 problem, 

317 strategy, 

318 data[param_range[i]], 

319 custom_description=description, 

320 comm_world=comm_world, 

321 problem_args=problem_args, 

322 num_procs=num_procs, 

323 hooks=hooks, 

324 Tend=Tend, 

325 num_procs_sweeper=num_procs_sweeper, 

326 ) 

327 

328 comm_world.Barrier() 

329 

330 if comm_world.rank == 0: 

331 if np.isfinite(data[param_range[i]]["k_linear"][-1]): 

332 k_type = "k_linear" 

333 elif np.isfinite(data[param_range[i]]["k_Newton"][-1]): 

334 k_type = 'k_Newton' 

335 else: 

336 k_type = "k_SDC" 

337 logger.log( 

338 25, 

339 f'{problem.__name__}: {strategy} {handle} {num_procs}-{num_procs_sweeper} procs, {param}={param_range[i]:.2e}: e={data[param_range[i]]["e_global"][-1]}, t={data[param_range[i]]["t"][-1]}, {k_type}={data[param_range[i]][k_type][-1]}', 

340 ) 

341 

342 if comm_world.rank == 0: 

343 import socket 

344 import time 

345 

346 data['meta'] = { 

347 'hostname': socket.gethostname(), 

348 'time': time.time, 

349 'runs': runs, 

350 } 

351 path = get_path(problem, strategy, num_procs, handle, num_procs_sweeper=num_procs_sweeper, mode=mode) 

352 with open(path, 'wb') as f: 

353 logger.debug(f'Dumping file \"{path}\"') 

354 pickle.dump(data, f) 

355 return data 

356 

357 

358def load(**kwargs): 

359 """ 

360 Load stored data. Arguments are passed on to `get_path` 

361 

362 Returns: 

363 dict: The data 

364 """ 

365 path = get_path(**kwargs) 

366 with open(path, 'rb') as f: 

367 logger.debug(f'Loading file \"{path}\"') 

368 data = pickle.load(f) 

369 return data 

370 

371 

372def extract_data(data, work_key, precision_key): 

373 """ 

374 Get the work and precision from a data object. 

375 

376 Args: 

377 data (dict): Data from work-precision measurements 

378 work_key (str): Name of variable on x-axis 

379 precision_key (str): Name of variable on y-axis 

380 

381 Returns: 

382 numpy array: Work 

383 numpy array: Precision 

384 """ 

385 keys = [key for key in data.keys() if key not in ['meta']] 

386 work = [np.nanmean(data[key][work_key]) for key in keys] 

387 precision = [np.nanmean(data[key][precision_key]) for key in keys] 

388 return np.array(work), np.array(precision) 

389 

390 

391def get_order(work_key='e_global', precision_key='param', strategy=None, handle=None, **kwargs): 

392 data = load(**kwargs, strategy=strategy, handle=handle) 

393 work, precision = extract_data(data, work_key, precision_key) 

394 

395 order = [np.log(precision[i + 1] / precision[i]) / np.log(work[i + 1] / work[i]) for i in range(len(work) - 1)] 

396 

397 print(f'Order for {strategy} {handle}: {np.mean(order):.2f}') 

398 

399 

400def plot_work_precision( 

401 problem, 

402 strategy, 

403 num_procs, 

404 ax, 

405 work_key='k_SDC', 

406 precision_key='e_global', 

407 handle='', 

408 plotting_params=None, 

409 comm_world=None, 

410 num_procs_sweeper=1, 

411 mode='', 

412): # pragma: no cover 

413 """ 

414 Plot data from running a problem with a strategy. 

415 

416 Args: 

417 problem (function): A problem to run 

418 strategy (Strategy): SDC strategy 

419 num_procs (int): Number of processes for the time communicator 

420 ax (matplotlib.pyplot.axes): Somewhere to plot 

421 work_key (str): The key in the recorded data you want on the x-axis 

422 precision_key (str): The key in the recorded data you want on the y-axis 

423 handle (str): The name of the configuration 

424 plotting_params (dict): Will be passed when plotting 

425 comm_world (mpi4py.MPI.Intracomm): Communicator that is available for the entire script 

426 num_procs_sweeper (int): Number of processes for the sweeper 

427 mode (str): The of the configurations you want to retrieve 

428 

429 Returns: 

430 None 

431 """ 

432 if comm_world.rank > 0 or get_forbidden_combinations(problem, strategy): 

433 return None 

434 

435 data = load( 

436 problem=problem, 

437 strategy=strategy, 

438 num_procs=num_procs, 

439 handle=handle, 

440 num_procs_sweeper=num_procs_sweeper, 

441 mode=mode, 

442 ) 

443 

444 work, precision = extract_data(data, work_key, precision_key) 

445 keys = [key for key in data.keys() if key not in ['meta']] 

446 

447 for key in [work_key, precision_key]: 

448 rel_variance = [np.std(data[me][key]) / max([np.nanmean(data[me][key]), 1.0]) for me in keys] 

449 if not all(me < 1e-1 or not np.isfinite(me) for me in rel_variance): 

450 logger.warning( 

451 f"Variance in \"{key}\" for {get_path(problem, strategy, num_procs, handle, num_procs_sweeper=num_procs_sweeper, mode=mode)} too large! Got {rel_variance}" 

452 ) 

453 

454 style = merge_descriptions( 

455 {**strategy.style, 'label': f'{strategy.style["label"]}{f" {handle}" if handle else ""}'}, 

456 plotting_params if plotting_params else {}, 

457 ) 

458 

459 mask = np.logical_and(np.isfinite(work), np.isfinite(precision)) 

460 ax.loglog(work[mask], precision[mask], **style) 

461 

462 # get_order( 

463 # problem=problem, 

464 # strategy=strategy, 

465 # num_procs=num_procs, 

466 # handle=handle, 

467 # work_key=work_key, 

468 # precision_key=precision_key, 

469 # ) 

470 

471 if 't' in [work_key, precision_key]: 

472 meta = data.get('meta', {}) 

473 

474 if meta.get('hostname', None) in ['thomas-work']: 

475 ax.text(0.1, 0.1, "Laptop timings!", transform=ax.transAxes) 

476 if meta.get('runs', None) == 1: 

477 ax.text(0.1, 0.2, "No sampling!", transform=ax.transAxes) 

478 

479 

480def plot_parallel_efficiency_diagonalSDC( 

481 ax, work_key, precision_key, num_procs_sweeper, num_procs=1, **kwargs 

482): # pragma: no cover 

483 serial_data = load( 

484 num_procs=num_procs, 

485 num_procs_sweeper=1, 

486 **kwargs, 

487 ) 

488 parallel_data = load( 

489 num_procs=num_procs, 

490 num_procs_sweeper=num_procs_sweeper, 

491 **kwargs, 

492 ) 

493 serial_work, serial_precision = extract_data(serial_data, work_key, precision_key) 

494 parallel_work, parallel_precision = extract_data(parallel_data, work_key, precision_key) 

495 # assert np.allclose(serial_precision, parallel_precision) 

496 

497 serial_work = np.asarray(serial_work) 

498 parallel_work = np.asarray(parallel_work) 

499 

500 # ax.loglog(serial_work, serial_precision) 

501 # ax.loglog(parallel_work, parallel_precision) 

502 

503 speedup = serial_work / parallel_work 

504 parallel_efficiency = np.median(speedup) / num_procs_sweeper 

505 ax.plot(serial_precision, speedup) 

506 ax.set_xscale('log') 

507 ax.set_ylabel('speedup') 

508 

509 if 't' in [work_key, precision_key]: 

510 meta = parallel_data.get('meta', {}) 

511 

512 if meta.get('hostname', None) in ['thomas-work']: 

513 ax.text(0.1, 0.1, "Laptop timings!", transform=ax.transAxes) 

514 if meta.get('runs', None) == 1: 

515 ax.text(0.1, 0.2, "No sampling!", transform=ax.transAxes) 

516 

517 return np.median(speedup), parallel_efficiency 

518 

519 

520def decorate_panel(ax, problem, work_key, precision_key, num_procs=1, title_only=False): # pragma: no cover 

521 """ 

522 Decorate a plot 

523 

524 Args: 

525 ax (matplotlib.pyplot.axes): Somewhere to plot 

526 problem (function): A problem to run 

527 work_key (str): The key in the recorded data you want on the x-axis 

528 precision_key (str): The key in the recorded data you want on the y-axis 

529 num_procs (int): Number of processes for the time communicator 

530 title_only (bool): Put only the title on top, or do the whole shebang 

531 

532 Returns: 

533 None 

534 """ 

535 labels = { 

536 'k_SDC': 'SDC iterations', 

537 'k_SDC_no_restart': 'SDC iterations (restarts excluded)', 

538 'k_Newton': 'Newton iterations', 

539 'k_Newton_no_restart': 'Newton iterations (restarts excluded)', 

540 'k_rhs': 'right hand side evaluations', 

541 't': 'wall clock time / s', 

542 'e_global': 'global error', 

543 'e_global_rel': 'relative global error', 

544 'e_local_max': 'max. local error', 

545 'restart': 'restarts', 

546 'dt_max': r'$\Delta t_\mathrm{max}$', 

547 'dt_min': r'$\Delta t_\mathrm{min}$', 

548 'dt_sigma': r'$\sigma(\Delta t)$', 

549 'dt_mean': r'$\bar{\Delta t}$', 

550 'param': 'accuracy parameter', 

551 'u0_increment_max': r'$\| \Delta u_0 \|_{\infty} $', 

552 'u0_increment_mean': r'$\bar{\Delta u_0}$', 

553 'u0_increment_max_no_restart': r'$\| \Delta u_0 \|_{\infty} $ (restarts excluded)', 

554 'u0_increment_mean_no_restart': r'$\bar{\Delta u_0}$ (restarts excluded)', 

555 'k_linear': 'Linear solver iterations', 

556 'speedup': 'Speedup', 

557 'nprocs': r'$N_\mathrm{procs}$', 

558 '': '', 

559 } 

560 

561 if not title_only: 

562 ax.set_xlabel(labels.get(work_key, 'work')) 

563 ax.set_ylabel(labels.get(precision_key, 'precision')) 

564 # ax.legend(frameon=False) 

565 

566 titles = { 

567 'run_vdp': 'Van der Pol', 

568 'run_Lorenz': 'Lorenz attractor', 

569 'run_Schroedinger': r'Schr\"odinger', 

570 'run_quench': 'Quench', 

571 'run_AC': 'Allen-Cahn', 

572 } 

573 ax.set_title(titles.get(problem.__name__, '')) 

574 

575 

576def execute_configurations( 

577 problem, 

578 configurations, 

579 work_key, 

580 precision_key, 

581 num_procs, 

582 ax, 

583 decorate, 

584 record, 

585 runs, 

586 comm_world, 

587 plotting, 

588 Tend=None, 

589 num_procs_sweeper=1, 

590 mode='', 

591): 

592 """ 

593 Run for multiple configurations. 

594 

595 Args: 

596 problem (function): A problem to run 

597 configurations (dict): The configurations you want to run with 

598 work_key (str): The key in the recorded data you want on the x-axis 

599 precision_key (str): The key in the recorded data you want on the y-axis 

600 num_procs (int): Number of processes for the time communicator 

601 ax (matplotlib.pyplot.axes): Somewhere to plot 

602 decorate (bool): Whether to decorate fully or only put the title 

603 record (bool): Whether to only plot or also record the data first 

604 runs (int): Number of runs you want to do 

605 comm_world (mpi4py.MPI.Intracomm): Communicator that is available for the entire script 

606 plotting (bool): Whether to plot something 

607 num_procs_sweeper (int): Number of processes for the sweeper 

608 mode (str): What you want to look at 

609 

610 Returns: 

611 None 

612 """ 

613 for _, config in configurations.items(): 

614 for strategy in config['strategies']: 

615 shared_args = { 

616 'problem': problem, 

617 'strategy': strategy, 

618 'handle': config.get('handle', ''), 

619 'num_procs': config.get('num_procs', num_procs), 

620 'num_procs_sweeper': config.get('num_procs_sweeper', num_procs_sweeper), 

621 } 

622 if record: 

623 logger.debug('Recording work precision') 

624 record_work_precision( 

625 **shared_args, 

626 custom_description=config.get('custom_description', {}), 

627 runs=runs, 

628 comm_world=comm_world, 

629 problem_args=config.get('problem_args', {}), 

630 param_range=config.get('param_range', None), 

631 hooks=config.get('hooks', None), 

632 Tend=config.get('Tend') if Tend is None else Tend, 

633 mode=mode, 

634 ) 

635 if plotting and comm_world.rank == 0: 

636 logger.debug('Plotting') 

637 plot_work_precision( 

638 **shared_args, 

639 work_key=work_key, 

640 precision_key=precision_key, 

641 ax=ax, 

642 plotting_params=config.get('plotting_params', {}), 

643 comm_world=comm_world, 

644 mode=mode, 

645 ) 

646 

647 if comm_world.rank == 0: 

648 decorate_panel( 

649 ax=ax, 

650 problem=problem, 

651 work_key=work_key, 

652 precision_key=precision_key, 

653 num_procs=num_procs, 

654 title_only=not decorate, 

655 ) 

656 

657 

658def get_configs(mode, problem): 

659 """ 

660 Get configurations for work-precision plots. These are dictionaries containing strategies and handles and so on. 

661 

662 Args: 

663 mode (str): The of the configurations you want to retrieve 

664 problem (function): A problem to run 

665 

666 Returns: 

667 dict: Configurations 

668 """ 

669 configurations = {} 

670 if mode == 'regular': 

671 from pySDC.projects.Resilience.strategies import AdaptivityStrategy, BaseStrategy, IterateStrategy 

672 

673 handle = 'regular' 

674 configurations[0] = { 

675 'handle': handle, 

676 'strategies': [AdaptivityStrategy(useMPI=True), BaseStrategy(useMPI=True), IterateStrategy(useMPI=True)], 

677 } 

678 elif mode == 'step_size_limiting': 

679 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter 

680 from pySDC.projects.Resilience.strategies import AdaptivityStrategy, ESDIRKStrategy 

681 

682 limits = [ 

683 25.0, 

684 50.0, 

685 ] 

686 colors = ['teal', 'magenta'] 

687 markers = ['v', 'x'] 

688 markersize = 9 

689 for i in range(len(limits)): 

690 configurations[i] = { 

691 'custom_description': {'convergence_controllers': {StepSizeLimiter: {'dt_max': limits[i]}}}, 

692 'handle': f'steplimiter{limits[i]:.0f}', 

693 'strategies': [AdaptivityStrategy(useMPI=True)], 

694 'plotting_params': { 

695 'color': colors[i], 

696 'marker': markers[i], 

697 'label': rf'$\Delta t \leq { {limits[i]:.0f}} $', 

698 # 'ls': '', 

699 'markersize': markersize, 

700 }, 

701 'num_procs': 1, 

702 } 

703 configurations[99] = { 

704 'custom_description': {}, 

705 'handle': 'no limits', 

706 'plotting_params': { 

707 'label': 'no limiter', 

708 # 'ls': '', 

709 'markersize': markersize, 

710 }, 

711 'strategies': [AdaptivityStrategy(useMPI=True)], 

712 'num_procs': 1, 

713 } 

714 elif mode == 'dynamic_restarts': 

715 """ 

716 Compare Block Gauss-Seidel SDC with restarting the first step in the block or the first step that exceeded the error threshold. 

717 """ 

718 from pySDC.projects.Resilience.strategies import AdaptivityStrategy, AdaptivityRestartFirstStep 

719 

720 desc = {} 

721 desc['sweeper_params'] = {'num_nodes': 3, 'QI': 'IE'} 

722 desc['step_params'] = {'maxiter': 5} 

723 

724 ls = { 

725 1: '-', 

726 2: '--', 

727 3: '-.', 

728 4: ':', 

729 5: ':', 

730 } 

731 

732 configurations[-1] = { 

733 'strategies': [AdaptivityStrategy(useMPI=True)], 

734 'num_procs': 1, 

735 } 

736 

737 for num_procs in [4, 2]: 

738 plotting_params = {'ls': ls[num_procs], 'label': f'adaptivity {num_procs} procs'} 

739 configurations[num_procs] = { 

740 'strategies': [AdaptivityStrategy(useMPI=True), AdaptivityRestartFirstStep(useMPI=True)], 

741 'custom_description': desc, 

742 'num_procs': num_procs, 

743 'plotting_params': plotting_params, 

744 } 

745 

746 elif mode == 'compare_strategies': 

747 """ 

748 Compare the different SDC strategies. 

749 """ 

750 from pySDC.projects.Resilience.strategies import ( 

751 AdaptivityStrategy, 

752 kAdaptivityStrategy, 

753 AdaptivityPolynomialError, 

754 BaseStrategy, 

755 ) 

756 

757 newton_inexactness = False if problem.__name__ in ['run_vdp'] else True 

758 

759 configurations[1] = { 

760 'strategies': [AdaptivityPolynomialError(useMPI=True, newton_inexactness=newton_inexactness)], 

761 } 

762 configurations[2] = { 

763 'strategies': [kAdaptivityStrategy(useMPI=True)], 

764 } 

765 configurations[0] = { 

766 'custom_description': { 

767 'step_params': {'maxiter': 5}, 

768 'sweeper_params': {'num_nodes': 3, 'quad_type': 'RADAU-RIGHT'}, 

769 }, 

770 'strategies': [ 

771 AdaptivityStrategy(useMPI=True), 

772 BaseStrategy(useMPI=True), 

773 ], 

774 } 

775 

776 elif mode == 'RK_comp': 

777 """ 

778 Compare parallel adaptive SDC to Runge-Kutta 

779 """ 

780 from pySDC.projects.Resilience.strategies import ( 

781 AdaptivityStrategy, 

782 ERKStrategy, 

783 ESDIRKStrategy, 

784 ARKStrategy, 

785 AdaptivityPolynomialError, 

786 ) 

787 

788 if problem.__name__ in ['run_Schroedinger', 'run_AC']: 

789 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

790 else: 

791 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

792 generic_implicit_MPI as parallel_sweeper, 

793 ) 

794 

795 newton_inexactness = False if problem.__name__ in ['run_vdp'] else True 

796 

797 desc = {} 

798 desc['sweeper_params'] = {'num_nodes': 3, 'QI': 'IE', 'QE': "EE"} 

799 desc['step_params'] = {'maxiter': 5} 

800 

801 desc_poly = {} 

802 desc_poly['sweeper_class'] = parallel_sweeper 

803 

804 ls = { 

805 1: '-', 

806 2: '--', 

807 3: '-.', 

808 4: ':', 

809 5: ':', 

810 } 

811 RK_strategies = [] 

812 if problem.__name__ in ['run_Lorenz']: 

813 RK_strategies.append(ERKStrategy(useMPI=True)) 

814 if problem.__name__ in ['run_Schroedinger', 'run_AC']: 

815 RK_strategies.append(ARKStrategy(useMPI=True)) 

816 else: 

817 RK_strategies.append(ESDIRKStrategy(useMPI=True)) 

818 

819 configurations[3] = { 

820 'custom_description': desc_poly, 

821 'strategies': [AdaptivityPolynomialError(useMPI=True, newton_inexactness=newton_inexactness)], 

822 'num_procs': 1, 

823 'num_procs_sweeper': 3, 

824 'plotting_params': {'label': r'$\Delta t$-$k$-adaptivity $N$=1x3'}, 

825 } 

826 configurations[-1] = { 

827 'strategies': RK_strategies, 

828 'num_procs': 1, 

829 } 

830 configurations[2] = { 

831 'strategies': [AdaptivityStrategy(useMPI=True)], 

832 'custom_description': desc, 

833 'num_procs': 4, 

834 'plotting_params': {'label': r'$\Delta t$-adaptivity $N$=4x1'}, 

835 } 

836 

837 elif mode == 'parallel_efficiency': 

838 """ 

839 Compare parallel runs of the step size adaptive SDC 

840 """ 

841 from pySDC.projects.Resilience.strategies import AdaptivityStrategy, AdaptivityPolynomialError 

842 

843 if problem.__name__ in ['run_Schroedinger', 'run_AC']: 

844 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

845 else: 

846 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

847 generic_implicit_MPI as parallel_sweeper, 

848 ) 

849 

850 desc = {} 

851 desc['sweeper_params'] = {'num_nodes': 3, 'QI': 'IE', 'QE': 'EE'} 

852 desc['step_params'] = {'maxiter': 5} 

853 

854 ls = { 

855 1: '-', 

856 2: '--', 

857 3: '-.', 

858 4: '--', 

859 5: ':', 

860 12: ':', 

861 } 

862 

863 newton_inexactness = False if problem.__name__ in ['run_vdp'] else True 

864 

865 for num_procs in [4, 1]: 

866 plotting_params = ( 

867 {'ls': ls[num_procs], 'label': fr'$\Delta t$-adaptivity $N$={num_procs}x1'} if num_procs > 1 else {} 

868 ) 

869 configurations[num_procs] = { 

870 'strategies': [AdaptivityStrategy(useMPI=True)], 

871 'custom_description': desc.copy(), 

872 'num_procs': num_procs, 

873 'plotting_params': plotting_params.copy(), 

874 } 

875 configurations[num_procs * 100 + 79] = { 

876 'custom_description': {'sweeper_class': parallel_sweeper}, 

877 'strategies': [ 

878 AdaptivityPolynomialError( 

879 useMPI=True, newton_inexactness=newton_inexactness, linear_inexactness=True 

880 ) 

881 ], 

882 'num_procs_sweeper': 3, 

883 'num_procs': num_procs, 

884 'plotting_params': { 

885 'ls': ls.get(num_procs * 3, '-'), 

886 'label': rf'$\Delta t$-$k$-adaptivity $N$={num_procs}x3', 

887 }, 

888 } 

889 

890 configurations[num_procs * 200 + 79] = { 

891 'strategies': [ 

892 AdaptivityPolynomialError(useMPI=True, newton_inexactness=newton_inexactness, linear_inexactness=True) 

893 ], 

894 'num_procs': 1, 

895 } 

896 

897 elif mode == 'interpolate_between_restarts': 

898 """ 

899 Compare adaptivity with interpolation between restarts and without 

900 """ 

901 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

902 

903 i = 0 

904 for interpolate_between_restarts, handle, ls in zip( 

905 [True, False], ['Interpolation between restarts', 'regular'], ['--', '-'] 

906 ): 

907 configurations[i] = { 

908 'strategies': [ 

909 AdaptivityPolynomialError(interpolate_between_restarts=interpolate_between_restarts, useMPI=True) 

910 ], 

911 'plotting_params': {'ls': ls}, 

912 'handle': handle, 

913 } 

914 i += 1 

915 elif mode == 'diagonal_SDC': 

916 """ 

917 Run diagonal SDC with different number of nodes and ranks. You can use this to compute a speedup, but it's not strong scaling! 

918 """ 

919 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

920 

921 if problem.__name__ in ['run_Schroedinger']: 

922 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

923 else: 

924 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

925 generic_implicit_MPI as parallel_sweeper, 

926 ) 

927 

928 for parallel in [False, True]: 

929 desc = {'sweeper_class': parallel_sweeper} if parallel else {} 

930 for num_nodes, ls in zip([3, 4, 2], ['-', '--', ':', '-.']): 

931 configurations[num_nodes + (99 if parallel else 0)] = { 

932 'custom_description': {**desc, 'sweeper_params': {'num_nodes': num_nodes}}, 

933 'strategies': [ 

934 AdaptivityPolynomialError(useMPI=True, newton_inexactness=True, linear_inexactness=True) 

935 ], 

936 'num_procs_sweeper': num_nodes if parallel else 1, 

937 'num_procs': 1, 

938 'handle': f'{num_nodes} nodes', 

939 'plotting_params': { 

940 'ls': ls, 

941 'label': f'{num_nodes} procs', 

942 # **{'color': 'grey' if parallel else None}, 

943 }, 

944 } 

945 

946 elif mode[:13] == 'vdp_stiffness': 

947 """ 

948 Run van der Pol with different parameter for the nonlinear term, which controls the stiffness. 

949 """ 

950 from pySDC.projects.Resilience.strategies import ( 

951 AdaptivityStrategy, 

952 ERKStrategy, 

953 ESDIRKStrategy, 

954 AdaptivityPolynomialError, 

955 ) 

956 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

957 generic_implicit_MPI as parallel_sweeper, 

958 ) 

959 

960 Tends = { 

961 1000: 2000, 

962 100: 200, 

963 10: 20, 

964 0: 2, 

965 } 

966 mu = float(mode[14:]) 

967 Tend = Tends[mu] 

968 

969 problem_desc = {'problem_params': {'mu': mu}} 

970 

971 desc = {} 

972 desc['sweeper_params'] = {'num_nodes': 3, 'QI': 'IE'} 

973 desc['step_params'] = {'maxiter': 5} 

974 desc['problem_params'] = problem_desc['problem_params'] 

975 

976 ls = { 

977 1: '-', 

978 2: '--', 

979 3: '-.', 

980 4: ':', 

981 5: ':', 

982 'MIN-SR-S': '-', 

983 'MIN-SR-NS': '--', 

984 'MIN-SR-FLEX': '-.', 

985 } 

986 

987 if mu < 100: 

988 configurations[2] = { 

989 'strategies': [ERKStrategy(useMPI=True)], 

990 'num_procs': 1, 

991 'handle': mode, 

992 'plotting_params': {'label': 'CP5(4)'}, 

993 'custom_description': problem_desc, 

994 'Tend': Tend, 

995 } 

996 configurations[1] = { 

997 'strategies': [AdaptivityStrategy(useMPI=True)], 

998 'custom_description': desc, 

999 'num_procs': 4, 

1000 'plotting_params': {'ls': ls[1], 'label': 'SDC $N$=4x1'}, 

1001 'handle': mode, 

1002 'Tend': Tend, 

1003 } 

1004 configurations[4] = { 

1005 'strategies': [ESDIRKStrategy(useMPI=True)], 

1006 'num_procs': 1, 

1007 'handle': mode, 

1008 'plotting_params': {'label': 'ESDIRK5(3)'}, 

1009 'custom_description': problem_desc, 

1010 'Tend': Tend, 

1011 } 

1012 for QI, i in zip( 

1013 [ 

1014 'MIN-SR-S', 

1015 # 'MIN-SR-FLEX', 

1016 ], 

1017 [9991, 12123127391, 1231723109247102731092], 

1018 ): 

1019 configurations[i] = { 

1020 'custom_description': { 

1021 'sweeper_params': {'num_nodes': 3, 'QI': QI}, 

1022 'problem_params': desc["problem_params"], 

1023 'sweeper_class': parallel_sweeper, 

1024 }, 

1025 'strategies': [ 

1026 AdaptivityPolynomialError( 

1027 useMPI=True, newton_inexactness=False, linear_inexactness=False, max_slope=4 

1028 ) 

1029 ], 

1030 'num_procs_sweeper': 3, 

1031 'num_procs': 1, 

1032 'plotting_params': { 

1033 'ls': ls.get(QI, '-'), 

1034 'label': rf'$\Delta t$-$k$-adaptivity $N$={1}x3', 

1035 }, 

1036 'handle': f'{mode}-{QI}', 

1037 'Tend': Tend, 

1038 } 

1039 

1040 elif mode == 'inexactness': 

1041 """ 

1042 Compare inexact SDC to exact SDC 

1043 """ 

1044 from pySDC.projects.Resilience.strategies import ( 

1045 AdaptivityPolynomialError, 

1046 ) 

1047 

1048 if problem.__name__ in ['run_Schroedinger']: 

1049 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

1050 else: 

1051 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1052 generic_implicit_MPI as parallel_sweeper, 

1053 ) 

1054 

1055 strategies = [ 

1056 AdaptivityPolynomialError, 

1057 ] 

1058 

1059 inexactness = { 

1060 'newton_inexactness': True, 

1061 'linear_inexactness': True, 

1062 } 

1063 no_inexactness = { 

1064 'newton_inexactness': False, 

1065 'linear_inexactness': False, 

1066 'SDC_maxiter': 99, 

1067 'use_restol_rel': False, 

1068 } 

1069 

1070 configurations[1] = { 

1071 'custom_description': {'sweeper_class': parallel_sweeper}, 

1072 'strategies': [me(useMPI=True, **no_inexactness) for me in strategies], 

1073 'num_procs_sweeper': 3, 

1074 'handle': 'exact', 

1075 'plotting_params': {'ls': '--'}, 

1076 } 

1077 configurations[0] = { 

1078 'custom_description': {'sweeper_class': parallel_sweeper}, 

1079 'strategies': [me(useMPI=True, **inexactness) for me in strategies], 

1080 'handle': 'inexact', 

1081 'num_procs_sweeper': 3, 

1082 } 

1083 elif mode == 'compare_adaptivity': 

1084 """ 

1085 Compare various modes of adaptivity 

1086 """ 

1087 # TODO: configurations not final! 

1088 from pySDC.projects.Resilience.strategies import ( 

1089 # AdaptivityCollocationTypeStrategy, 

1090 # AdaptivityCollocationRefinementStrategy, 

1091 AdaptivityStrategy, 

1092 # AdaptivityExtrapolationWithinQStrategy, 

1093 ESDIRKStrategy, 

1094 ARKStrategy, 

1095 AdaptivityPolynomialError, 

1096 ) 

1097 

1098 if problem.__name__ in ['run_Schroedinger']: 

1099 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

1100 else: 

1101 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1102 generic_implicit_MPI as parallel_sweeper, 

1103 ) 

1104 

1105 inexactness_params = { 

1106 # 'double_adaptivity': True, 

1107 'newton_inexactness': True, 

1108 'linear_inexactness': True, 

1109 } 

1110 

1111 strategies = [ 

1112 AdaptivityPolynomialError, 

1113 # AdaptivityCollocationTypeStrategy, 

1114 # AdaptivityExtrapolationWithinQStrategy, 

1115 ] 

1116 

1117 # restol = None 

1118 # for strategy in strategies: 

1119 # strategy.restol = restol 

1120 

1121 configurations[1] = { 

1122 'custom_description': {'sweeper_class': parallel_sweeper}, 

1123 'strategies': [me(useMPI=True, **inexactness_params) for me in strategies], 

1124 'handle': 'parallel', 

1125 'num_procs_sweeper': 3, 

1126 'plotting_params': {'ls': '-', 'label': '3 procs'}, 

1127 } 

1128 configurations[2] = { 

1129 'strategies': [me(useMPI=True, **inexactness_params) for me in strategies], 

1130 'plotting_params': {'ls': '--'}, 

1131 } 

1132 configurations[4] = { 

1133 'custom_description': {'step_params': {'maxiter': 5}}, 

1134 'strategies': [AdaptivityStrategy(useMPI=True)], 

1135 } 

1136 

1137 desc_RK = {} 

1138 configurations[-1] = { 

1139 'strategies': [ 

1140 ARKStrategy(useMPI=True) if problem.__name__ == 'run_Schroedinger' else ESDIRKStrategy(useMPI=True), 

1141 ], 

1142 'num_procs': 1, 

1143 'custom_description': desc_RK, 

1144 } 

1145 

1146 elif mode == 'preconditioners': 

1147 """ 

1148 Compare different preconditioners 

1149 """ 

1150 from pySDC.projects.Resilience.strategies import ( 

1151 AdaptivityStrategy, 

1152 IterateStrategy, 

1153 BaseStrategy, 

1154 ESDIRKStrategy, 

1155 ERKStrategy, 

1156 AdaptivityPolynomialError, 

1157 ) 

1158 

1159 inexacness = True 

1160 strategies = [ 

1161 AdaptivityPolynomialError( 

1162 useMPI=True, SDC_maxiter=29, newton_inexactness=inexacness, linear_inexactness=inexacness 

1163 ), 

1164 BaseStrategy(useMPI=True), 

1165 ] 

1166 

1167 desc = {} 

1168 desc['sweeper_params'] = { 

1169 'num_nodes': 3, 

1170 } 

1171 # desc['step_params'] = {'maxiter': 5} 

1172 

1173 precons = ['IE', 'LU'] 

1174 ls = ['-.', '--', '-', ':'] 

1175 for i in range(len(precons) + 1): 

1176 if i < len(precons): 

1177 desc['sweeper_params']['QI'] = precons[i] 

1178 handle = precons[i] 

1179 else: 

1180 handle = None 

1181 configurations[i] = { 

1182 'strategies': strategies, 

1183 'custom_description': copy.deepcopy(desc), 

1184 'handle': handle, 

1185 'plotting_params': {'ls': ls[i]}, 

1186 } 

1187 elif mode == 'RK_comp_high_order': 

1188 """ 

1189 Compare higher order SDC than we can get with RKM to RKM 

1190 """ 

1191 from pySDC.projects.Resilience.strategies import ( 

1192 AdaptivityStrategy, 

1193 ERKStrategy, 

1194 ESDIRKStrategy, 

1195 ARKStrategy, 

1196 AdaptivityPolynomialError, 

1197 ) 

1198 

1199 if problem.__name__ in ['run_Schroedinger']: 

1200 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as parallel_sweeper 

1201 else: 

1202 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1203 generic_implicit_MPI as parallel_sweeper, 

1204 ) 

1205 

1206 desc = {} 

1207 desc['sweeper_params'] = {'num_nodes': 4, 'QI': 'IE', 'QE': "EE"} 

1208 desc['step_params'] = {'maxiter': 7} 

1209 

1210 desc_poly = {} 

1211 desc_poly['sweeper_class'] = parallel_sweeper 

1212 desc_poly['sweeper_params'] = {'num_nodes': 4} 

1213 

1214 ls = { 

1215 1: '-', 

1216 2: '--', 

1217 3: '-.', 

1218 4: ':', 

1219 5: ':', 

1220 } 

1221 

1222 desc_RK = {} 

1223 if problem.__name__ in ['run_Schroedinger']: 

1224 desc_RK['problem_params'] = {'imex': True} 

1225 

1226 configurations[3] = { 

1227 'custom_description': desc_poly, 

1228 'strategies': [AdaptivityPolynomialError(useMPI=True)], 

1229 'num_procs': 1, 

1230 'num_procs_sweeper': 4, 

1231 } 

1232 configurations[-1] = { 

1233 'strategies': [ 

1234 ERKStrategy(useMPI=True), 

1235 ARKStrategy(useMPI=True) if problem.__name__ in ['run_Schroedinger'] else ESDIRKStrategy(useMPI=True), 

1236 ], 

1237 'num_procs': 1, 

1238 'custom_description': desc_RK, 

1239 } 

1240 

1241 configurations[2] = { 

1242 'strategies': [AdaptivityStrategy(useMPI=True)], 

1243 'custom_description': desc, 

1244 'num_procs': 4, 

1245 } 

1246 elif mode == 'avoid_restarts': 

1247 """ 

1248 Test how well avoiding restarts works. 

1249 """ 

1250 from pySDC.projects.Resilience.strategies import ( 

1251 AdaptivityStrategy, 

1252 AdaptivityAvoidRestartsStrategy, 

1253 AdaptivityPolynomialStrategy, 

1254 ) 

1255 

1256 desc = {'sweeper_params': {'QI': 'IE'}, 'step_params': {'maxiter': 3}} 

1257 param_range = [1e-3, 1e-5] 

1258 configurations[0] = { 

1259 'strategies': [AdaptivityPolynomialStrategy(useMPI=True)], 

1260 'plotting_params': {'ls': '--'}, 

1261 'custom_description': desc, 

1262 'param_range': param_range, 

1263 } 

1264 configurations[1] = { 

1265 'strategies': [AdaptivityAvoidRestartsStrategy(useMPI=True)], 

1266 'plotting_params': {'ls': '-.'}, 

1267 'custom_description': desc, 

1268 'param_range': param_range, 

1269 } 

1270 configurations[2] = { 

1271 'strategies': [AdaptivityStrategy(useMPI=True)], 

1272 'custom_description': desc, 

1273 'param_range': param_range, 

1274 } 

1275 else: 

1276 raise NotImplementedError(f'Don\'t know the mode "{mode}"!') 

1277 

1278 return configurations 

1279 

1280 

1281def get_fig(x=1, y=1, **kwargs): # pragma: no cover 

1282 """ 

1283 Get a figure to plot in. 

1284 

1285 Args: 

1286 x (int): How many panels in horizontal direction you want 

1287 y (int): How many panels in vertical direction you want 

1288 

1289 Returns: 

1290 matplotlib.pyplot.Figure 

1291 """ 

1292 width = 1.0 

1293 ratio = 1.0 if y == 2 else 0.5 

1294 keyword_arguments = { 

1295 'figsize': figsize_by_journal('Springer_Numerical_Algorithms', width, ratio), 

1296 'layout': 'constrained', 

1297 **kwargs, 

1298 } 

1299 return plt.subplots(y, x, **keyword_arguments) 

1300 

1301 

1302def save_fig( 

1303 fig, name, work_key, precision_key, legend=True, format='pdf', base_path='data', squares=True, ncols=None, **kwargs 

1304): # pragma: no cover 

1305 """ 

1306 Save a figure with a legend on the bottom. 

1307 

1308 Args: 

1309 fig (matplotlib.pyplot.Figure): Figure you want to save 

1310 name (str): Name of the plot to put in the path 

1311 work_key (str): The key in the recorded data you want on the x-axis 

1312 precision_key (str): The key in the recorded data you want on the y-axis 

1313 legend (bool): Put a legend or not 

1314 format (str): Format to store the figure with 

1315 base_path (str): Some path where all the files are stored 

1316 squares (bool): Adjust aspect ratio to squares if true 

1317 

1318 Returns: 

1319 None 

1320 """ 

1321 handles = [] 

1322 labels = [] 

1323 for ax in fig.get_axes(): 

1324 h, l = ax.get_legend_handles_labels() 

1325 handles += [h[i] for i in range(len(h)) if l[i] not in labels] 

1326 labels += [me for me in l if me not in labels] 

1327 if squares: 

1328 ax.set_box_aspect(1) 

1329 # order = np.argsort([me[0] for me in labels]) 

1330 order = np.arange(len(labels)) 

1331 fig.legend( 

1332 [handles[i] for i in order], 

1333 [labels[i] for i in order], 

1334 loc='outside lower center', 

1335 ncols=ncols if ncols else 3 if len(handles) % 3 == 0 else 4, 

1336 frameon=False, 

1337 fancybox=True, 

1338 handlelength=2.2, 

1339 ) 

1340 

1341 path = f'{base_path}/wp-{name}-{work_key}-{precision_key}.{format}' 

1342 fig.savefig(path, bbox_inches='tight', **kwargs) 

1343 print(f'Stored figure \"{path}\"') 

1344 

1345 

1346def all_problems(mode='compare_strategies', plotting=True, base_path='data', **kwargs): # pragma: no cover 

1347 """ 

1348 Make a plot comparing various strategies for all problems. 

1349 

1350 Args: 

1351 work_key (str): The key in the recorded data you want on the x-axis 

1352 precision_key (str): The key in the recorded data you want on the y-axis 

1353 

1354 Returns: 

1355 None 

1356 """ 

1357 

1358 fig, axs = get_fig(2, 2) 

1359 

1360 shared_params = { 

1361 'work_key': 'k_SDC', 

1362 'precision_key': 'e_global', 

1363 'num_procs': 1, 

1364 'runs': 1, 

1365 'comm_world': MPI.COMM_WORLD, 

1366 'record': False, 

1367 'plotting': plotting, 

1368 **kwargs, 

1369 } 

1370 

1371 problems = [run_vdp, run_quench, run_Schroedinger, run_AC] 

1372 

1373 logger.log(26, f"Doing for all problems {mode}") 

1374 for i in range(len(problems)): 

1375 execute_configurations( 

1376 **shared_params, 

1377 problem=problems[i], 

1378 ax=axs.flatten()[i], 

1379 decorate=True, 

1380 configurations=get_configs(mode, problems[i]), 

1381 mode=mode, 

1382 ) 

1383 

1384 if plotting and shared_params['comm_world'].rank == 0: 

1385 ncols = { 

1386 'parallel_efficiency': 2, 

1387 'RK_comp': 2, 

1388 } 

1389 y_right_dt_fixed = [1e18, 4e1, 5, 1e8] 

1390 y_right_dt = [1e-1, 1e4, 1, 2e-2] 

1391 y_right_dtk = [1e-4, 1e4, 1e-2, 1e-3] 

1392 

1393 if shared_params['work_key'] == 'param': 

1394 for ax, yRfixed, yRdt, yRdtk in zip(fig.get_axes(), y_right_dt_fixed, y_right_dt, y_right_dtk): 

1395 add_order_line(ax, 1, '--', yRdt, marker=None) 

1396 add_order_line(ax, 5 / 4, ':', yRdtk, marker=None) 

1397 add_order_line(ax, 5, '-.', yRfixed, marker=None) 

1398 save_fig( 

1399 fig=fig, 

1400 name=mode, 

1401 work_key=shared_params['work_key'], 

1402 precision_key=shared_params['precision_key'], 

1403 legend=True, 

1404 base_path=base_path, 

1405 ncols=ncols.get(mode, None), 

1406 ) 

1407 

1408 

1409def ODEs(mode='compare_strategies', plotting=True, base_path='data', **kwargs): # pragma: no cover 

1410 """ 

1411 Make a plot comparing various strategies for the two ODEs. 

1412 

1413 Args: 

1414 work_key (str): The key in the recorded data you want on the x-axis 

1415 precision_key (str): The key in the recorded data you want on the y-axis 

1416 

1417 Returns: 

1418 None 

1419 """ 

1420 

1421 fig, axs = get_fig(x=2, y=1) 

1422 

1423 shared_params = { 

1424 'work_key': 'k_SDC', 

1425 'precision_key': 'e_global', 

1426 'num_procs': 1, 

1427 'runs': 1, 

1428 'comm_world': MPI.COMM_WORLD, 

1429 'record': False, 

1430 'plotting': plotting, 

1431 **kwargs, 

1432 } 

1433 

1434 problems = [run_vdp, run_Lorenz] 

1435 

1436 for i in range(len(problems)): 

1437 execute_configurations( 

1438 **shared_params, 

1439 problem=problems[i], 

1440 ax=axs.flatten()[i], 

1441 decorate=i == 0, 

1442 configurations=get_configs(mode, problems[i]), 

1443 ) 

1444 

1445 if plotting and shared_params['comm_world'].rank == 0: 

1446 save_fig( 

1447 fig=fig, 

1448 name=f'ODEs-{mode}', 

1449 work_key=shared_params['work_key'], 

1450 precision_key=shared_params['precision_key'], 

1451 legend=True, 

1452 base_path=base_path, 

1453 ) 

1454 

1455 

1456def single_problem(mode, problem, plotting=True, base_path='data', **kwargs): # pragma: no cover 

1457 """ 

1458 Make a plot for a single problem 

1459 

1460 Args: 

1461 mode (str): What you want to look at 

1462 problem (function): A problem to run 

1463 """ 

1464 fig, ax = get_fig(1, 1, figsize=figsize_by_journal('Springer_Numerical_Algorithms', 1, 0.8)) 

1465 

1466 params = { 

1467 'work_key': 'k_SDC', 

1468 'precision_key': 'e_global', 

1469 'num_procs': 1, 

1470 'runs': 1, 

1471 'comm_world': MPI.COMM_WORLD, 

1472 'record': False, 

1473 'plotting': plotting, 

1474 **kwargs, 

1475 } 

1476 

1477 logger.log(26, f"Doing single problem {mode}") 

1478 execute_configurations( 

1479 **params, problem=problem, ax=ax, decorate=True, configurations=get_configs(mode, problem), mode=mode 

1480 ) 

1481 

1482 if plotting: 

1483 save_fig( 

1484 fig=fig, 

1485 name=f'{problem.__name__}-{mode}', 

1486 work_key=params['work_key'], 

1487 precision_key=params['precision_key'], 

1488 legend=False, 

1489 base_path=base_path, 

1490 ) 

1491 

1492 

1493def vdp_stiffness_plot(base_path='data', format='pdf', **kwargs): # pragma: no cover 

1494 fig, axs = get_fig(3, 1, sharex=False, sharey=False) 

1495 

1496 mus = [10, 100, 1000] 

1497 

1498 for i in range(len(mus)): 

1499 params = { 

1500 'runs': 1, 

1501 'problem': run_vdp, 

1502 'record': False, 

1503 'work_key': 't', 

1504 'precision_key': 'e_global', 

1505 'comm_world': MPI.COMM_WORLD, 

1506 **kwargs, 

1507 } 

1508 params['num_procs'] = min(params['comm_world'].size, 5) 

1509 params['plotting'] = params['comm_world'].rank == 0 

1510 

1511 mode = f'vdp_stiffness-{mus[i]}' 

1512 configurations = get_configs(mode=mode, problem=run_vdp) 

1513 execute_configurations(**params, ax=axs.flatten()[i], decorate=True, configurations=configurations, mode=mode) 

1514 axs.flatten()[i].set_title(rf'$\mu={ {mus[i]}} $') 

1515 

1516 fig.suptitle('Van der Pol') 

1517 if params['comm_world'].rank == 0: 

1518 save_fig( 

1519 fig=fig, 

1520 name='vdp-stiffness', 

1521 work_key=params['work_key'], 

1522 precision_key=params['precision_key'], 

1523 legend=False, 

1524 base_path=base_path, 

1525 format=format, 

1526 ) 

1527 

1528 

1529def add_order_line(ax, order, ls, y_right=1.0, marker='.'): 

1530 x_min = min([min(line.get_xdata()) for line in ax.get_lines()]) 

1531 x_max = max([max(line.get_xdata()) for line in ax.get_lines()]) 

1532 y_min = min([min(line.get_ydata()) for line in ax.get_lines()]) 

1533 y_max = max([max(line.get_ydata()) for line in ax.get_lines()]) 

1534 x = np.logspace(np.log10(x_min), np.log10(x_max), 100) 

1535 y = y_right * (x / x_max) ** order 

1536 mask = np.logical_and(y > y_min, y < y_max) 

1537 ax.loglog(x[mask], y[mask], ls=ls, color='black', label=f'Order {order}', marker=marker, markevery=5) 

1538 

1539 

1540def aggregate_parallel_efficiency_plot(): # pragma: no cover 

1541 """ 

1542 Make a "weak scaling" plot for diagonal SDC 

1543 """ 

1544 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

1545 

1546 fig, axs = plt.subplots(2, 2) 

1547 

1548 _fig, _ax = plt.subplots(1, 1) 

1549 num_procs = 1 

1550 num_procs_sweeper = 2 

1551 problem = run_quench 

1552 

1553 num_procs_sweeper_list = [2, 3, 4] 

1554 

1555 for problem, ax in zip([run_vdp, run_Lorenz, run_quench], axs.flatten()): 

1556 speedup = [] 

1557 for num_procs_sweeper in num_procs_sweeper_list: 

1558 s, e = plot_parallel_efficiency_diagonalSDC( 

1559 ax=_ax, 

1560 work_key='t', 

1561 precision_key='e_global_rel', 

1562 num_procs=num_procs, 

1563 num_procs_sweeper=num_procs_sweeper, 

1564 problem=problem, 

1565 strategy=AdaptivityPolynomialError(), 

1566 mode='diagonal_SDC', 

1567 handle=f'{num_procs_sweeper} nodes', 

1568 ) 

1569 speedup += [s] 

1570 decorate_panel(ax, problem, work_key='nprocs', precision_key='') 

1571 

1572 ax.plot(num_procs_sweeper_list, speedup, label='speedup') 

1573 ax.plot( 

1574 num_procs_sweeper_list, 

1575 [speedup[i] / num_procs_sweeper_list[i] for i in range(len(speedup))], 

1576 label='parallel efficiency', 

1577 ) 

1578 

1579 fig.tight_layout() 

1580 save_fig(fig, 'parallel_efficiency', 'nprocs', 'speedup') 

1581 

1582 

1583if __name__ == "__main__": 

1584 comm_world = MPI.COMM_WORLD 

1585 

1586 record = comm_world.size > 1 

1587 for mode in [ 

1588 # 'compare_strategies', 

1589 # 'RK_comp', 

1590 # 'parallel_efficiency', 

1591 ]: 

1592 params = { 

1593 'mode': mode, 

1594 'runs': 5, 

1595 'plotting': comm_world.rank == 0, 

1596 } 

1597 params_single = { 

1598 **params, 

1599 'problem': run_AC, 

1600 } 

1601 single_problem(**params_single, work_key='t', precision_key='e_global_rel', record=record) 

1602 

1603 all_params = { 

1604 'record': True, 

1605 'runs': 5, 

1606 'work_key': 't', 

1607 'precision_key': 'e_global_rel', 

1608 'plotting': comm_world.rank == 0, 

1609 } 

1610 

1611 for mode in [ 

1612 'RK_comp', 

1613 'parallel_efficiency', 

1614 'compare_strategies', 

1615 ]: 

1616 all_problems(**all_params, mode=mode) 

1617 comm_world.Barrier() 

1618 

1619 if comm_world.rank == 0: 

1620 plt.show()