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

358 statements  

« prev     ^ index     » next       coverage.py v7.6.7, created at 2024-11-16 14:51 +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 

15from pySDC.projects.Resilience.RBC import run_RBC 

16 

17from pySDC.helpers.stats_helper import get_sorted, filter_stats 

18from pySDC.helpers.plot_helper import setup_mpl, figsize_by_journal 

19 

20setup_mpl(reset=True) 

21LOGGER_LEVEL = 25 

22LOG_TO_FILE = False 

23 

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

25 

26 

27def std_log(x): 

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

29 

30 

31MAPPINGS = { 

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

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

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

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

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

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

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

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

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

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

42 'k_factorizations': ('work_factorizations', sum, None), 

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

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

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

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

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

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

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

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

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

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

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

54} 

55 

56logger = logging.getLogger('WorkPrecision') 

57logger.setLevel(LOGGER_LEVEL) 

58 

59 

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

61 """ 

62 Check if the combination of strategy and problem is forbidden 

63 

64 Args: 

65 problem (function): A problem to run 

66 strategy (Strategy): SDC strategy 

67 """ 

68 if strategy.name == 'ERK': 

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

70 return True 

71 

72 return False 

73 

74 

75Tends = { 

76 'run_RBC': 16, 

77} 

78t0s = { 

79 'run_RBC': 0.2, 

80} 

81 

82 

83def single_run( 

84 problem, 

85 strategy, 

86 data, 

87 custom_description, 

88 num_procs=1, 

89 comm_world=None, 

90 problem_args=None, 

91 hooks=None, 

92 Tend=None, 

93 num_procs_sweeper=1, 

94): 

95 """ 

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

97 

98 Args: 

99 problem (function): A problem to run 

100 strategy (Strategy): SDC strategy 

101 data (dict): Put the results in here 

102 custom_description (dict): Overwrite presets 

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

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

105 hooks (list): List of additional hooks 

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

107 

108 Returns: 

109 dict: Stats generated by the run 

110 """ 

111 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun 

112 from pySDC.implementations.hooks.log_work import LogWork 

113 from pySDC.projects.Resilience.hook import LogData 

114 

115 hooks = hooks if hooks else [] 

116 

117 t_last = perf_counter() 

118 

119 num_procs_tot = num_procs * num_procs_sweeper 

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

121 if comm_world.rank >= num_procs_tot: 

122 comm.Free() 

123 return None 

124 

125 # make communicators for time and sweepers 

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

127 comm_sweep = comm.Split(comm_time.rank) 

128 

129 if comm_time.size < num_procs: 

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

131 

132 strategy_description = strategy.get_custom_description(problem, num_procs) 

133 description = merge_descriptions(strategy_description, custom_description) 

134 if comm_sweep.size > 1: 

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

136 

137 controller_params = { 

138 'logger_level': LOGGER_LEVEL, 

139 'log_to_file': LOG_TO_FILE, 

140 'fname': 'out.txt', 

141 **strategy.get_controller_params(), 

142 } 

143 problem_args = {} if problem_args is None else problem_args 

144 

145 Tend = Tends.get(problem.__name__, None) if Tend is None else Tend 

146 t0 = t0s.get(problem.__name__, None) 

147 

148 stats, controller, crash = problem( 

149 custom_description=description, 

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

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

152 custom_controller_params=controller_params, 

153 use_MPI=True, 

154 t0=t0, 

155 comm=comm_time, 

156 **problem_args, 

157 ) 

158 

159 t_now = perf_counter() 

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

161 t_last = perf_counter() 

162 

163 # record all the metrics 

164 stats_all = filter_stats(stats, comm=comm_sweep) 

165 comm_sweep.Free() 

166 

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

168 if crash: 

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

170 continue 

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

172 if len(me) == 0: 

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

174 else: 

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

176 

177 t_now = perf_counter() 

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

179 t_last = perf_counter() 

180 

181 comm_time.Free() 

182 comm.Free() 

183 return stats 

184 

185 

186def get_parameter(dictionary, where): 

187 """ 

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

189 

190 Args: 

191 dictionary (dict): The dictionary 

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

193 

194 Returns: 

195 The value of the dictionary 

196 """ 

197 if len(where) == 1: 

198 return dictionary[where[0]] 

199 else: 

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

201 

202 

203def set_parameter(dictionary, where, parameter): 

204 """ 

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

206 

207 Args: 

208 dictionary (dict): The dictionary 

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

210 parameter: Whatever you want to set the parameter to 

211 

212 Returns: 

213 None 

214 """ 

215 if len(where) == 1: 

216 dictionary[where[0]] = parameter 

217 else: 

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

219 

220 

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

222 """ 

223 Get the path to a certain data. 

224 

225 Args: 

226 problem (function): A problem to run 

227 strategy (Strategy): SDC strategy 

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

229 handle (str): The name of the configuration 

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

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

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

233 

234 Returns: 

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

236 """ 

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

238 

239 

240def record_work_precision( 

241 problem, 

242 strategy, 

243 num_procs=1, 

244 custom_description=None, 

245 handle='', 

246 runs=1, 

247 comm_world=None, 

248 problem_args=None, 

249 param_range=None, 

250 Tend=None, 

251 hooks=None, 

252 num_procs_sweeper=1, 

253 mode='', 

254): 

255 """ 

256 Run problem with strategy and record the cost parameters. 

257 

258 Args: 

259 problem (function): A problem to run 

260 strategy (Strategy): SDC strategy 

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

262 custom_description (dict): Overwrite presets 

263 handle (str): The name of the configuration 

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

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

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

267 

268 Returns: 

269 None 

270 """ 

271 if get_forbidden_combinations(problem, strategy): 

272 return None 

273 

274 data = {} 

275 

276 # prepare precision parameters 

277 param = strategy.precision_parameter 

278 description = merge_descriptions( 

279 strategy.get_custom_description(problem, num_procs), 

280 {} if custom_description is None else custom_description, 

281 ) 

282 if param == 'e_tol': 

283 power = 10.0 

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

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

286 if problem.__name__ == 'run_vdp': 

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

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

289 else: 

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

291 if problem.__name__ == 'run_RBC': 

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

293 elif param == 'dt': 

294 power = 2.0 

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

296 elif param == 'restol': 

297 power = 10.0 

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

299 if problem.__name__ == 'run_vdp': 

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

301 else: 

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

303 

304 where = strategy.precision_parameter_loc 

305 default = get_parameter(description, where) 

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

307 

308 if problem.__name__ == 'run_quench': 

309 if param == 'restol': 

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

311 elif param == 'dt': 

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

313 if problem.__name__ == 'run_RBC': 

314 if param == 'dt': 

315 param_range = [2e-1, 1e-1, 8e-2, 6e-2] 

316 

317 # run multiple times with different parameters 

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

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

320 

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

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

323 data[param_range[i]][param] = [param_range[i]] 

324 

325 description = merge_descriptions( 

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

327 ) 

328 for _j in range(runs): 

329 if comm_world.rank == 0: 

330 logger.log( 

331 24, 

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

333 ) 

334 single_run( 

335 problem, 

336 strategy, 

337 data[param_range[i]], 

338 custom_description=description, 

339 comm_world=comm_world, 

340 problem_args=problem_args, 

341 num_procs=num_procs, 

342 hooks=hooks, 

343 Tend=Tend, 

344 num_procs_sweeper=num_procs_sweeper, 

345 ) 

346 

347 comm_world.Barrier() 

348 

349 if comm_world.rank == 0: 

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

351 k_type = "k_linear" 

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

353 k_type = 'k_Newton' 

354 else: 

355 k_type = "k_SDC" 

356 logger.log( 

357 25, 

358 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]}', 

359 ) 

360 

361 if comm_world.rank == 0: 

362 import socket 

363 import time 

364 

365 data['meta'] = { 

366 'hostname': socket.gethostname(), 

367 'time': time.time, 

368 'runs': runs, 

369 } 

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

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

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

373 pickle.dump(data, f) 

374 return data 

375 

376 

377def load(**kwargs): 

378 """ 

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

380 

381 Returns: 

382 dict: The data 

383 """ 

384 path = get_path(**kwargs) 

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

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

387 data = pickle.load(f) 

388 return data 

389 

390 

391def extract_data(data, work_key, precision_key): 

392 """ 

393 Get the work and precision from a data object. 

394 

395 Args: 

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

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

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

399 

400 Returns: 

401 numpy array: Work 

402 numpy array: Precision 

403 """ 

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

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

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

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

408 

409 

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

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

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

413 

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

415 

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

417 

418 

419def plot_work_precision( 

420 problem, 

421 strategy, 

422 num_procs, 

423 ax, 

424 work_key='k_SDC', 

425 precision_key='e_global', 

426 handle='', 

427 plotting_params=None, 

428 comm_world=None, 

429 num_procs_sweeper=1, 

430 mode='', 

431): # pragma: no cover 

432 """ 

433 Plot data from running a problem with a strategy. 

434 

435 Args: 

436 problem (function): A problem to run 

437 strategy (Strategy): SDC strategy 

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

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

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

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

442 handle (str): The name of the configuration 

443 plotting_params (dict): Will be passed when plotting 

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

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

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

447 

448 Returns: 

449 None 

450 """ 

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

452 return None 

453 

454 data = load( 

455 problem=problem, 

456 strategy=strategy, 

457 num_procs=num_procs, 

458 handle=handle, 

459 num_procs_sweeper=num_procs_sweeper, 

460 mode=mode, 

461 ) 

462 

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

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

465 

466 for key in [work_key, precision_key]: 

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

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

469 logger.warning( 

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

471 ) 

472 

473 style = merge_descriptions( 

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

475 plotting_params if plotting_params else {}, 

476 ) 

477 

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

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

480 

481 # get_order( 

482 # problem=problem, 

483 # strategy=strategy, 

484 # num_procs=num_procs, 

485 # handle=handle, 

486 # work_key=work_key, 

487 # precision_key=precision_key, 

488 # ) 

489 

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

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

492 

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

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

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

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

497 

498 

499def plot_parallel_efficiency_diagonalSDC( 

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

501): # pragma: no cover 

502 serial_data = load( 

503 num_procs=num_procs, 

504 num_procs_sweeper=1, 

505 **kwargs, 

506 ) 

507 parallel_data = load( 

508 num_procs=num_procs, 

509 num_procs_sweeper=num_procs_sweeper, 

510 **kwargs, 

511 ) 

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

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

514 # assert np.allclose(serial_precision, parallel_precision) 

515 

516 serial_work = np.asarray(serial_work) 

517 parallel_work = np.asarray(parallel_work) 

518 

519 # ax.loglog(serial_work, serial_precision) 

520 # ax.loglog(parallel_work, parallel_precision) 

521 

522 speedup = serial_work / parallel_work 

523 parallel_efficiency = np.median(speedup) / num_procs_sweeper 

524 ax.plot(serial_precision, speedup) 

525 ax.set_xscale('log') 

526 ax.set_ylabel('speedup') 

527 

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

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

530 

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

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

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

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

535 

536 return np.median(speedup), parallel_efficiency 

537 

538 

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

540 """ 

541 Decorate a plot 

542 

543 Args: 

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

545 problem (function): A problem to run 

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

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

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

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

550 

551 Returns: 

552 None 

553 """ 

554 labels = { 

555 'k_SDC': 'SDC iterations', 

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

557 'k_Newton': 'Newton iterations', 

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

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

560 'k_factorizations': 'matrix factorizations', 

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

562 'e_global': 'global error', 

563 'e_global_rel': 'relative global error', 

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

565 'restart': 'restarts', 

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

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

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

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

570 'param': 'accuracy parameter', 

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

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

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

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

575 'k_linear': 'Linear solver iterations', 

576 'speedup': 'Speedup', 

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

578 '': '', 

579 } 

580 

581 if not title_only: 

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

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

584 # ax.legend(frameon=False) 

585 

586 titles = { 

587 'run_vdp': 'Van der Pol', 

588 'run_Lorenz': 'Lorenz attractor', 

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

590 'run_quench': 'Quench', 

591 'run_AC': 'Allen-Cahn', 

592 } 

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

594 

595 

596def execute_configurations( 

597 problem, 

598 configurations, 

599 work_key, 

600 precision_key, 

601 num_procs, 

602 ax, 

603 decorate, 

604 record, 

605 runs, 

606 comm_world, 

607 plotting, 

608 Tend=None, 

609 num_procs_sweeper=1, 

610 mode='', 

611): 

612 """ 

613 Run for multiple configurations. 

614 

615 Args: 

616 problem (function): A problem to run 

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

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

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

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

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

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

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

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

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

626 plotting (bool): Whether to plot something 

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

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

629 

630 Returns: 

631 None 

632 """ 

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

634 for strategy in config['strategies']: 

635 shared_args = { 

636 'problem': problem, 

637 'strategy': strategy, 

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

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

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

641 } 

642 if record: 

643 logger.debug('Recording work precision') 

644 record_work_precision( 

645 **shared_args, 

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

647 runs=runs, 

648 comm_world=comm_world, 

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

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

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

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

653 mode=mode, 

654 ) 

655 if plotting and comm_world.rank == 0: 

656 logger.debug('Plotting') 

657 plot_work_precision( 

658 **shared_args, 

659 work_key=work_key, 

660 precision_key=precision_key, 

661 ax=ax, 

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

663 comm_world=comm_world, 

664 mode=mode, 

665 ) 

666 

667 if comm_world.rank == 0: 

668 decorate_panel( 

669 ax=ax, 

670 problem=problem, 

671 work_key=work_key, 

672 precision_key=precision_key, 

673 num_procs=num_procs, 

674 title_only=not decorate, 

675 ) 

676 

677 

678def get_configs(mode, problem): 

679 """ 

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

681 

682 Args: 

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

684 problem (function): A problem to run 

685 

686 Returns: 

687 dict: Configurations 

688 """ 

689 configurations = {} 

690 if mode == 'regular': 

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

692 

693 handle = 'regular' 

694 configurations[0] = { 

695 'handle': handle, 

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

697 } 

698 elif mode == 'step_size_limiting': 

699 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeLimiter 

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

701 

702 limits = [ 

703 25.0, 

704 50.0, 

705 ] 

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

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

708 markersize = 9 

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

710 configurations[i] = { 

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

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

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

714 'plotting_params': { 

715 'color': colors[i], 

716 'marker': markers[i], 

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

718 # 'ls': '', 

719 'markersize': markersize, 

720 }, 

721 'num_procs': 1, 

722 } 

723 configurations[99] = { 

724 'custom_description': {}, 

725 'handle': 'no limits', 

726 'plotting_params': { 

727 'label': 'no limiter', 

728 # 'ls': '', 

729 'markersize': markersize, 

730 }, 

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

732 'num_procs': 1, 

733 } 

734 elif mode == 'dynamic_restarts': 

735 """ 

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

737 """ 

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

739 

740 desc = {} 

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

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

743 

744 ls = { 

745 1: '-', 

746 2: '--', 

747 3: '-.', 

748 4: ':', 

749 5: ':', 

750 } 

751 

752 configurations[-1] = { 

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

754 'num_procs': 1, 

755 } 

756 

757 for num_procs in [4, 2]: 

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

759 configurations[num_procs] = { 

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

761 'custom_description': desc, 

762 'num_procs': num_procs, 

763 'plotting_params': plotting_params, 

764 } 

765 

766 elif mode == 'compare_strategies': 

767 """ 

768 Compare the different SDC strategies. 

769 """ 

770 from pySDC.projects.Resilience.strategies import ( 

771 AdaptivityStrategy, 

772 kAdaptivityStrategy, 

773 AdaptivityPolynomialError, 

774 BaseStrategy, 

775 ) 

776 

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

778 

779 configurations[1] = { 

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

781 } 

782 configurations[2] = { 

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

784 } 

785 configurations[0] = { 

786 'custom_description': { 

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

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

789 }, 

790 'strategies': [ 

791 AdaptivityStrategy(useMPI=True), 

792 BaseStrategy(useMPI=True), 

793 ], 

794 } 

795 

796 elif mode == 'RK_comp': 

797 """ 

798 Compare parallel adaptive SDC to Runge-Kutta 

799 """ 

800 from pySDC.projects.Resilience.strategies import ( 

801 AdaptivityStrategy, 

802 ERKStrategy, 

803 ESDIRKStrategy, 

804 ARKStrategy, 

805 AdaptivityPolynomialError, 

806 ) 

807 

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

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

810 else: 

811 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

812 generic_implicit_MPI as parallel_sweeper, 

813 ) 

814 

815 newton_inexactness = False if problem.__name__ in ['run_vdp', 'run_RBC'] else True 

816 

817 desc = {} 

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

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

820 

821 desc_poly = {} 

822 desc_poly['sweeper_class'] = parallel_sweeper 

823 

824 ls = { 

825 1: '-', 

826 2: '--', 

827 3: '-.', 

828 4: ':', 

829 5: ':', 

830 } 

831 RK_strategies = [] 

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

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

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

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

836 else: 

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

838 

839 configurations[3] = { 

840 'custom_description': desc_poly, 

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

842 'num_procs': 1, 

843 'num_procs_sweeper': 3, 

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

845 } 

846 configurations[-1] = { 

847 'strategies': RK_strategies, 

848 'num_procs': 1, 

849 } 

850 configurations[2] = { 

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

852 'custom_description': desc, 

853 'num_procs': 4, 

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

855 } 

856 

857 elif mode == 'parallel_efficiency': 

858 """ 

859 Compare parallel runs of the step size adaptive SDC 

860 """ 

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

862 

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

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

865 else: 

866 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

867 generic_implicit_MPI as parallel_sweeper, 

868 ) 

869 

870 desc = {} 

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

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

873 

874 ls = { 

875 1: '-', 

876 2: '--', 

877 3: '-.', 

878 4: '--', 

879 5: ':', 

880 12: ':', 

881 } 

882 

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

884 

885 for num_procs in [4, 1]: 

886 plotting_params = ( 

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

888 ) 

889 configurations[num_procs] = { 

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

891 'custom_description': desc.copy(), 

892 'num_procs': num_procs, 

893 'plotting_params': plotting_params.copy(), 

894 } 

895 configurations[num_procs * 100 + 79] = { 

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

897 'strategies': [ 

898 AdaptivityPolynomialError( 

899 useMPI=True, newton_inexactness=newton_inexactness, linear_inexactness=True 

900 ) 

901 ], 

902 'num_procs_sweeper': 3, 

903 'num_procs': num_procs, 

904 'plotting_params': { 

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

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

907 }, 

908 } 

909 

910 configurations[num_procs * 200 + 79] = { 

911 'strategies': [ 

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

913 ], 

914 'num_procs': 1, 

915 } 

916 

917 elif mode == 'interpolate_between_restarts': 

918 """ 

919 Compare adaptivity with interpolation between restarts and without 

920 """ 

921 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

922 

923 i = 0 

924 for interpolate_between_restarts, handle, ls in zip( 

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

926 ): 

927 configurations[i] = { 

928 'strategies': [ 

929 AdaptivityPolynomialError(interpolate_between_restarts=interpolate_between_restarts, useMPI=True) 

930 ], 

931 'plotting_params': {'ls': ls}, 

932 'handle': handle, 

933 } 

934 i += 1 

935 elif mode == 'diagonal_SDC': 

936 """ 

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

938 """ 

939 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

940 

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

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

943 else: 

944 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

945 generic_implicit_MPI as parallel_sweeper, 

946 ) 

947 

948 for parallel in [False, True]: 

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

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

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

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

953 'strategies': [ 

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

955 ], 

956 'num_procs_sweeper': num_nodes if parallel else 1, 

957 'num_procs': 1, 

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

959 'plotting_params': { 

960 'ls': ls, 

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

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

963 }, 

964 } 

965 

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

967 """ 

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

969 """ 

970 from pySDC.projects.Resilience.strategies import ( 

971 AdaptivityStrategy, 

972 ERKStrategy, 

973 ESDIRKStrategy, 

974 AdaptivityPolynomialError, 

975 ) 

976 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

977 generic_implicit_MPI as parallel_sweeper, 

978 ) 

979 

980 Tends = { 

981 1000: 2000, 

982 100: 200, 

983 10: 20, 

984 0: 2, 

985 } 

986 mu = float(mode[14:]) 

987 Tend = Tends[mu] 

988 

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

990 

991 desc = {} 

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

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

994 desc['problem_params'] = problem_desc['problem_params'] 

995 

996 ls = { 

997 1: '-', 

998 2: '--', 

999 3: '-.', 

1000 4: ':', 

1001 5: ':', 

1002 'MIN-SR-S': '-', 

1003 'MIN-SR-NS': '--', 

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

1005 } 

1006 

1007 if mu < 100: 

1008 configurations[2] = { 

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

1010 'num_procs': 1, 

1011 'handle': mode, 

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

1013 'custom_description': problem_desc, 

1014 'Tend': Tend, 

1015 } 

1016 configurations[1] = { 

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

1018 'custom_description': desc, 

1019 'num_procs': 4, 

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

1021 'handle': mode, 

1022 'Tend': Tend, 

1023 } 

1024 configurations[4] = { 

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

1026 'num_procs': 1, 

1027 'handle': mode, 

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

1029 'custom_description': problem_desc, 

1030 'Tend': Tend, 

1031 } 

1032 for QI, i in zip( 

1033 [ 

1034 'MIN-SR-S', 

1035 # 'MIN-SR-FLEX', 

1036 ], 

1037 [9991, 12123127391, 1231723109247102731092], 

1038 ): 

1039 configurations[i] = { 

1040 'custom_description': { 

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

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

1043 'sweeper_class': parallel_sweeper, 

1044 }, 

1045 'strategies': [ 

1046 AdaptivityPolynomialError( 

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

1048 ) 

1049 ], 

1050 'num_procs_sweeper': 3, 

1051 'num_procs': 1, 

1052 'plotting_params': { 

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

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

1055 }, 

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

1057 'Tend': Tend, 

1058 } 

1059 

1060 elif mode == 'inexactness': 

1061 """ 

1062 Compare inexact SDC to exact SDC 

1063 """ 

1064 from pySDC.projects.Resilience.strategies import ( 

1065 AdaptivityPolynomialError, 

1066 ) 

1067 

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

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

1070 else: 

1071 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1072 generic_implicit_MPI as parallel_sweeper, 

1073 ) 

1074 

1075 strategies = [ 

1076 AdaptivityPolynomialError, 

1077 ] 

1078 

1079 inexactness = { 

1080 'newton_inexactness': True, 

1081 'linear_inexactness': True, 

1082 } 

1083 no_inexactness = { 

1084 'newton_inexactness': False, 

1085 'linear_inexactness': False, 

1086 'SDC_maxiter': 99, 

1087 'use_restol_rel': False, 

1088 } 

1089 

1090 configurations[1] = { 

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

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

1093 'num_procs_sweeper': 3, 

1094 'handle': 'exact', 

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

1096 } 

1097 configurations[0] = { 

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

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

1100 'handle': 'inexact', 

1101 'num_procs_sweeper': 3, 

1102 } 

1103 elif mode == 'compare_adaptivity': 

1104 """ 

1105 Compare various modes of adaptivity 

1106 """ 

1107 # TODO: configurations not final! 

1108 from pySDC.projects.Resilience.strategies import ( 

1109 # AdaptivityCollocationTypeStrategy, 

1110 # AdaptivityCollocationRefinementStrategy, 

1111 AdaptivityStrategy, 

1112 # AdaptivityExtrapolationWithinQStrategy, 

1113 ESDIRKStrategy, 

1114 ARKStrategy, 

1115 AdaptivityPolynomialError, 

1116 ) 

1117 

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

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

1120 else: 

1121 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1122 generic_implicit_MPI as parallel_sweeper, 

1123 ) 

1124 

1125 inexactness_params = { 

1126 # 'double_adaptivity': True, 

1127 'newton_inexactness': True, 

1128 'linear_inexactness': True, 

1129 } 

1130 

1131 strategies = [ 

1132 AdaptivityPolynomialError, 

1133 # AdaptivityCollocationTypeStrategy, 

1134 # AdaptivityExtrapolationWithinQStrategy, 

1135 ] 

1136 

1137 # restol = None 

1138 # for strategy in strategies: 

1139 # strategy.restol = restol 

1140 

1141 configurations[1] = { 

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

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

1144 'handle': 'parallel', 

1145 'num_procs_sweeper': 3, 

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

1147 } 

1148 configurations[2] = { 

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

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

1151 } 

1152 configurations[4] = { 

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

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

1155 } 

1156 

1157 desc_RK = {} 

1158 configurations[-1] = { 

1159 'strategies': [ 

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

1161 ], 

1162 'num_procs': 1, 

1163 'custom_description': desc_RK, 

1164 } 

1165 

1166 elif mode == 'preconditioners': 

1167 """ 

1168 Compare different preconditioners 

1169 """ 

1170 from pySDC.projects.Resilience.strategies import ( 

1171 AdaptivityStrategy, 

1172 IterateStrategy, 

1173 BaseStrategy, 

1174 ESDIRKStrategy, 

1175 ERKStrategy, 

1176 AdaptivityPolynomialError, 

1177 ) 

1178 

1179 inexacness = True 

1180 strategies = [ 

1181 AdaptivityPolynomialError( 

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

1183 ), 

1184 BaseStrategy(useMPI=True), 

1185 ] 

1186 

1187 desc = {} 

1188 desc['sweeper_params'] = { 

1189 'num_nodes': 3, 

1190 } 

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

1192 

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

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

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

1196 if i < len(precons): 

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

1198 handle = precons[i] 

1199 else: 

1200 handle = None 

1201 configurations[i] = { 

1202 'strategies': strategies, 

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

1204 'handle': handle, 

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

1206 } 

1207 elif mode == 'RK_comp_high_order': 

1208 """ 

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

1210 """ 

1211 from pySDC.projects.Resilience.strategies import ( 

1212 AdaptivityStrategy, 

1213 ERKStrategy, 

1214 ESDIRKStrategy, 

1215 ARKStrategy, 

1216 AdaptivityPolynomialError, 

1217 ) 

1218 

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

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

1221 else: 

1222 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import ( 

1223 generic_implicit_MPI as parallel_sweeper, 

1224 ) 

1225 

1226 desc = {} 

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

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

1229 

1230 desc_poly = {} 

1231 desc_poly['sweeper_class'] = parallel_sweeper 

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

1233 

1234 ls = { 

1235 1: '-', 

1236 2: '--', 

1237 3: '-.', 

1238 4: ':', 

1239 5: ':', 

1240 } 

1241 

1242 desc_RK = {} 

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

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

1245 

1246 configurations[3] = { 

1247 'custom_description': desc_poly, 

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

1249 'num_procs': 1, 

1250 'num_procs_sweeper': 4, 

1251 } 

1252 configurations[-1] = { 

1253 'strategies': [ 

1254 ERKStrategy(useMPI=True), 

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

1256 ], 

1257 'num_procs': 1, 

1258 'custom_description': desc_RK, 

1259 } 

1260 

1261 configurations[2] = { 

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

1263 'custom_description': desc, 

1264 'num_procs': 4, 

1265 } 

1266 elif mode == 'avoid_restarts': 

1267 """ 

1268 Test how well avoiding restarts works. 

1269 """ 

1270 from pySDC.projects.Resilience.strategies import ( 

1271 AdaptivityStrategy, 

1272 AdaptivityAvoidRestartsStrategy, 

1273 AdaptivityPolynomialStrategy, 

1274 ) 

1275 

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

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

1278 configurations[0] = { 

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

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

1281 'custom_description': desc, 

1282 'param_range': param_range, 

1283 } 

1284 configurations[1] = { 

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

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

1287 'custom_description': desc, 

1288 'param_range': param_range, 

1289 } 

1290 configurations[2] = { 

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

1292 'custom_description': desc, 

1293 'param_range': param_range, 

1294 } 

1295 else: 

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

1297 

1298 return configurations 

1299 

1300 

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

1302 """ 

1303 Get a figure to plot in. 

1304 

1305 Args: 

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

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

1308 

1309 Returns: 

1310 matplotlib.pyplot.Figure 

1311 """ 

1312 width = 1.0 

1313 ratio = 1.0 if y == 2 else 0.5 

1314 keyword_arguments = { 

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

1316 'layout': 'constrained', 

1317 **kwargs, 

1318 } 

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

1320 

1321 

1322def save_fig( 

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

1324): # pragma: no cover 

1325 """ 

1326 Save a figure with a legend on the bottom. 

1327 

1328 Args: 

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

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

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

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

1333 legend (bool): Put a legend or not 

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

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

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

1337 

1338 Returns: 

1339 None 

1340 """ 

1341 handles = [] 

1342 labels = [] 

1343 for ax in fig.get_axes(): 

1344 h, l = ax.get_legend_handles_labels() 

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

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

1347 if squares: 

1348 ax.set_box_aspect(1) 

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

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

1351 fig.legend( 

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

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

1354 loc='outside lower center', 

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

1356 frameon=False, 

1357 fancybox=True, 

1358 handlelength=2.2, 

1359 ) 

1360 

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

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

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

1364 

1365 

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

1367 """ 

1368 Make a plot comparing various strategies for all problems. 

1369 

1370 Args: 

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

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

1373 

1374 Returns: 

1375 None 

1376 """ 

1377 

1378 fig, axs = get_fig(2, 2) 

1379 

1380 shared_params = { 

1381 'work_key': 'k_SDC', 

1382 'precision_key': 'e_global', 

1383 'num_procs': 1, 

1384 'runs': 1, 

1385 'comm_world': MPI.COMM_WORLD, 

1386 'record': False, 

1387 'plotting': plotting, 

1388 **kwargs, 

1389 } 

1390 

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

1392 

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

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

1395 execute_configurations( 

1396 **shared_params, 

1397 problem=problems[i], 

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

1399 decorate=True, 

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

1401 mode=mode, 

1402 ) 

1403 

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

1405 ncols = { 

1406 'parallel_efficiency': 2, 

1407 'RK_comp': 2, 

1408 } 

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

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

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

1412 

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

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

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

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

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

1418 save_fig( 

1419 fig=fig, 

1420 name=mode, 

1421 work_key=shared_params['work_key'], 

1422 precision_key=shared_params['precision_key'], 

1423 legend=True, 

1424 base_path=base_path, 

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

1426 ) 

1427 

1428 

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

1430 """ 

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

1432 

1433 Args: 

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

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

1436 

1437 Returns: 

1438 None 

1439 """ 

1440 

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

1442 

1443 shared_params = { 

1444 'work_key': 'k_SDC', 

1445 'precision_key': 'e_global', 

1446 'num_procs': 1, 

1447 'runs': 1, 

1448 'comm_world': MPI.COMM_WORLD, 

1449 'record': False, 

1450 'plotting': plotting, 

1451 **kwargs, 

1452 } 

1453 

1454 problems = [run_vdp, run_Lorenz] 

1455 

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

1457 execute_configurations( 

1458 **shared_params, 

1459 problem=problems[i], 

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

1461 decorate=i == 0, 

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

1463 ) 

1464 

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

1466 save_fig( 

1467 fig=fig, 

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

1469 work_key=shared_params['work_key'], 

1470 precision_key=shared_params['precision_key'], 

1471 legend=True, 

1472 base_path=base_path, 

1473 ) 

1474 

1475 

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

1477 """ 

1478 Make a plot for a single problem 

1479 

1480 Args: 

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

1482 problem (function): A problem to run 

1483 """ 

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

1485 

1486 params = { 

1487 'work_key': 'k_SDC', 

1488 'precision_key': 'e_global', 

1489 'num_procs': 1, 

1490 'runs': 1, 

1491 'comm_world': MPI.COMM_WORLD, 

1492 'record': False, 

1493 'plotting': plotting, 

1494 **kwargs, 

1495 } 

1496 

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

1498 execute_configurations( 

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

1500 ) 

1501 

1502 if plotting: 

1503 save_fig( 

1504 fig=fig, 

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

1506 work_key=params['work_key'], 

1507 precision_key=params['precision_key'], 

1508 legend=False, 

1509 base_path=base_path, 

1510 ) 

1511 

1512 

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

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

1515 

1516 mus = [10, 100, 1000] 

1517 

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

1519 params = { 

1520 'runs': 1, 

1521 'problem': run_vdp, 

1522 'record': False, 

1523 'work_key': 't', 

1524 'precision_key': 'e_global', 

1525 'comm_world': MPI.COMM_WORLD, 

1526 **kwargs, 

1527 } 

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

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

1530 

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

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

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

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

1535 

1536 fig.suptitle('Van der Pol') 

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

1538 save_fig( 

1539 fig=fig, 

1540 name='vdp-stiffness', 

1541 work_key=params['work_key'], 

1542 precision_key=params['precision_key'], 

1543 legend=False, 

1544 base_path=base_path, 

1545 format=format, 

1546 ) 

1547 

1548 

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

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

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

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

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

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

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

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

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

1558 

1559 

1560def aggregate_parallel_efficiency_plot(): # pragma: no cover 

1561 """ 

1562 Make a "weak scaling" plot for diagonal SDC 

1563 """ 

1564 from pySDC.projects.Resilience.strategies import AdaptivityPolynomialError 

1565 

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

1567 

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

1569 num_procs = 1 

1570 num_procs_sweeper = 2 

1571 problem = run_quench 

1572 

1573 num_procs_sweeper_list = [2, 3, 4] 

1574 

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

1576 speedup = [] 

1577 for num_procs_sweeper in num_procs_sweeper_list: 

1578 s, e = plot_parallel_efficiency_diagonalSDC( 

1579 ax=_ax, 

1580 work_key='t', 

1581 precision_key='e_global_rel', 

1582 num_procs=num_procs, 

1583 num_procs_sweeper=num_procs_sweeper, 

1584 problem=problem, 

1585 strategy=AdaptivityPolynomialError(), 

1586 mode='diagonal_SDC', 

1587 handle=f'{num_procs_sweeper} nodes', 

1588 ) 

1589 speedup += [s] 

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

1591 

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

1593 ax.plot( 

1594 num_procs_sweeper_list, 

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

1596 label='parallel efficiency', 

1597 ) 

1598 

1599 fig.tight_layout() 

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

1601 

1602 

1603if __name__ == "__main__": 

1604 comm_world = MPI.COMM_WORLD 

1605 

1606 record = comm_world.size > 1 

1607 for mode in [ 

1608 # 'compare_strategies', 

1609 # 'RK_comp', 

1610 # 'parallel_efficiency', 

1611 ]: 

1612 params = { 

1613 'mode': mode, 

1614 'runs': 1, 

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

1616 } 

1617 params_single = { 

1618 **params, 

1619 'problem': run_RBC, 

1620 } 

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

1622 

1623 all_params = { 

1624 'record': True, 

1625 'runs': 5, 

1626 'work_key': 't', 

1627 'precision_key': 'e_global_rel', 

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

1629 } 

1630 

1631 for mode in [ 

1632 'RK_comp', 

1633 'parallel_efficiency', 

1634 'compare_strategies', 

1635 ]: 

1636 all_problems(**all_params, mode=mode) 

1637 comm_world.Barrier() 

1638 

1639 if comm_world.rank == 0: 

1640 plt.show()