Coverage for pySDC/tutorial/step_5/A_multistep_multilevel_hierarchy.py: 100%

1# --- 

2# jupyter: 

3# jupytext: 

4# formats: py:percent 

5# kernelspec: 

6# display_name: Python 3 

7# name: python3 

8# --- 

9 

10# %% [markdown] 

11# # Part A: Multistep multilevel hierarchy 

12# 

13# For SDC and MLSDC, one `step` was all we needed. PFASST works on several time steps at once, and a controller 

14# represents them simply as a list of steps, its attribute `MS`. The nice thing about going from MLSDC to PFASST is 

15# that only one number changes: `num_procs`, the number of time steps the controller works on at once. This 

16# controller only emulates working on them in parallel, see the box below; running them really in parallel takes the 

17# MPI controller as well, with one process per step. 

18 

19# %% 

20import matplotlib.pyplot as plt 

21 

22from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

23from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_forced 

24from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

25from pySDC.implementations.transfer_classes.TransferMesh import mesh_to_mesh 

26 

27# initialize level parameters 

28level_params = {'restol': 1e-10, 'dt': 0.5} 

29 

30# initialize sweeper parameters 

31sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': [3]} 

32 

33# initialize problem parameters 

34problem_params = { 

35 'nu': 0.1, # diffusion coefficient 

36 'freq': 4, # frequency for the test value 

37 'nvars': [31, 15, 7], # number of degrees of freedom for each level 

38 'bc': 'dirichlet-zero', # boundary conditions 

39} 

40 

41# initialize step parameters 

42step_params = {'maxiter': 20} 

43 

44# initialize space transfer parameters 

45space_transfer_params = {'rorder': 2, 'iorder': 6} 

46 

47# fill description dictionary for easy step instantiation 

48description = { 

49 'problem_class': heatNd_forced, # pass problem class 

50 'problem_params': problem_params, # pass problem parameters 

51 'sweeper_class': imex_1st_order, # pass sweeper (see part B) 

52 'sweeper_params': sweeper_params, # pass sweeper parameters 

53 'level_params': level_params, # pass level parameters 

54 'step_params': step_params, # pass step parameters 

55 'space_transfer_class': mesh_to_mesh, # pass spatial transfer class 

56 'space_transfer_params': space_transfer_params, # pass parameters for spatial transfer 

57} 

58 

59# %% [markdown] 

60# The description is the one of an MLSDC run with three levels in space. The controller gets 10 processes: 

61 

62# %% tags=["scroll-output"] 

63# instantiate controller 

64controller = controller_nonMPI(num_procs=10, controller_params={}, description=description) 

65 

66# %% 

67# check number of levels 

68for i in range(len(controller.MS)): 

69 print("Process %2i has %2i levels" % (i, len(controller.MS[i].levels))) 

70 

71# %% [markdown] 

72# Ten steps, each a full hierarchy of its own, with its own levels, problems and sweepers: 

73 

74# %% tags=["hide-input"] 

75fig, ax = plt.subplots(figsize=(8, 2.6)) 

76for i, S in enumerate(controller.MS): 

77 for L in S.levels: 

78 ax.add_patch( 

79 plt.Rectangle((i + 0.05, L.level_index + 0.05), 0.9, 0.9, color=plt.cm.Blues(0.8 - 0.25 * L.level_index)) 

80 ) 

81 ax.text(i + 0.5, L.level_index + 0.5, L.prob.nvars[0], ha='center', va='center', color='white', fontsize=9) 

82ax.set_xlim(0, len(controller.MS)) 

83ax.set_ylim(len(controller.MS[0].levels), 0) 

84ax.set_xticks([i + 0.5 for i in range(len(controller.MS))]) 

85ax.set_xticklabels([f'step {i}' for i in range(len(controller.MS))], fontsize=8) 

86ax.set_yticks([0.5, 1.5, 2.5]) 

87ax.set_yticklabels(['level 0 (fine)', 'level 1', 'level 2 (coarse)'], fontsize=8) 

88ax.set_title('the steps in controller.MS, with the number of unknowns on each level', fontsize=10) 

89ax.set_frame_on(False) 

90fig.tight_layout() 

91 

92# %% 

93print('Step 0 and step 1 share their problem:', controller.MS[0].levels[0].prob is controller.MS[1].levels[0].prob) 

94 

95# %% [markdown] 

96# :::{admonition} Important things to note 

97# - Controllers with the `_nonMPI` suffix only emulate parallelism, which saves the tedious installation of mpi4py 

98# and gives full access to all data at all times. The algorithm is the same, but the steps are computed 

99# serially: the controller moves all of them through the algorithm together, one stage at a time, and within 

100# each stage it handles one step after another. The `MPI` controllers run them really in parallel and should give 

101# the same results, see [Step 6](../step_6). 

102# - All steps of a controller are created from the same description, so they all have the same levels, as the 

103# check below confirms. 

104# ::: 

105# 

106# The check the tests run: 

107 

108# %% 

109assert all(len(S.levels) == 3 for S in controller.MS), "ERROR: not all steps have the same number of levels"