Coverage for pySDC/projects/Hamiltonian/harmonic_oscillator.py: 0%

1import os 

2 

3import dill 

4import numpy as np 

5 

6import pySDC.helpers.plot_helper as plt_helper 

7from pySDC.helpers.stats_helper import get_sorted 

8 

9from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

10from pySDC.implementations.problem_classes.HarmonicOscillator import harmonic_oscillator 

11from pySDC.implementations.sweeper_classes.verlet import verlet 

12from pySDC.implementations.transfer_classes.TransferParticles_NoCoarse import particles_to_particles 

13from pySDC.projects.Hamiltonian.stop_at_error_hook import stop_at_error_hook 

14 

15 

16def run_simulation(): 

17 """ 

18 Routine to run the simulation of a second order problem 

19 

20 """ 

21 

22 # initialize level parameters 

23 level_params = dict() 

24 level_params['restol'] = 0.0 

25 level_params['dt'] = 1.0 

26 

27 # initialize sweeper parameters 

28 sweeper_params = dict() 

29 sweeper_params['quad_type'] = 'LOBATTO' 

30 sweeper_params['num_nodes'] = [5, 3] 

31 sweeper_params['initial_guess'] = 'zero' 

32 

33 # initialize problem parameters for the Penning trap 

34 problem_params = dict() 

35 problem_params['k'] = None # will be defined later 

36 problem_params['phase'] = 0.0 

37 problem_params['amp'] = 1.0 

38 

39 # initialize step parameters 

40 step_params = dict() 

41 step_params['maxiter'] = 50 

42 

43 # initialize controller parameters 

44 controller_params = dict() 

45 controller_params['hook_class'] = stop_at_error_hook 

46 controller_params['logger_level'] = 30 

47 

48 # Fill description dictionary for easy hierarchy creation 

49 description = dict() 

50 description['problem_class'] = harmonic_oscillator 

51 description['sweeper_class'] = verlet 

52 description['level_params'] = level_params 

53 description['step_params'] = step_params 

54 description['space_transfer_class'] = particles_to_particles 

55 

56 # set time parameters 

57 t0 = 0.0 

58 Tend = 4.0 

59 num_procs = 4 

60 

61 rlim_left = 0 

62 rlim_right = 16.0 

63 nstep = 34 

64 ks = np.linspace(rlim_left, rlim_right, nstep)[1:] 

65 

66 # qd_combinations = [('IE', 'EE'), ('IE', 'PIC'), 

67 # ('LU', 'EE'), ('LU', 'PIC'), 

68 # # ('MIN3', 'PIC'), ('MIN3', 'EE'), 

69 # ('PIC', 'EE'), ('PIC', 'PIC')] 

70 qd_combinations = [('IE', 'EE'), ('PIC', 'PIC')] 

71 

72 results = dict() 

73 results['ks'] = ks 

74 

75 for qd in qd_combinations: 

76 print('Working on combination (%s, %s)...' % qd) 

77 

78 niters = np.zeros(len(ks)) 

79 

80 for i, k in enumerate(ks): 

81 problem_params['k'] = k 

82 description['problem_params'] = problem_params 

83 

84 sweeper_params['QI'] = qd[0] 

85 sweeper_params['QE'] = qd[1] 

86 description['sweeper_params'] = sweeper_params 

87 

88 # instantiate the controller 

89 controller = controller_nonMPI( 

90 num_procs=num_procs, controller_params=controller_params, description=description 

91 ) 

92 

93 # get initial values on finest level 

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

95 uinit = P.u_exact(t=t0) 

96 

97 # call main function to get things done... 

98 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend) 

99 

100 uex = P.u_exact(Tend) 

101 

102 print('Error after run: %s' % abs(uex - uend)) 

103 

104 # filter statistics by type (number of iterations) 

105 iter_counts = get_sorted(stats, type='niter', sortby='time') 

106 

107 niters[i] = np.mean(np.array([item[1] for item in iter_counts])) 

108 

109 # print('Worked on k = %s, took %s iterations' % (k, results[i])) 

110 

111 results[qd] = niters 

112 

113 fname = 'data/harmonic_k.dat' 

114 f = open(fname, 'wb') 

115 dill.dump(results, f) 

116 f.close() 

117 

118 assert os.path.isfile(fname), 'Run did not create stats file' 

119 

120 

121def show_results(cwd=''): 

122 """ 

123 Helper function to plot the error of the Hamiltonian 

124 

125 Args: 

126 cwd (str): current working directory 

127 """ 

128 

129 plt_helper.mpl.style.use('classic') 

130 plt_helper.setup_mpl() 

131 plt_helper.newfig(textwidth=238.96, scale=0.89) 

132 

133 # read in the dill data 

134 f = open(cwd + 'data/harmonic_k.dat', 'rb') 

135 results = dill.load(f) 

136 f.close() 

137 

138 ks = results['ks'] 

139 

140 for qd in results: 

141 if qd != 'ks': 

142 plt_helper.plt.plot(ks, results[qd], label=qd) 

143 

144 plt_helper.plt.xlabel('k') 

145 plt_helper.plt.ylabel('Number of iterations') 

146 plt_helper.plt.legend( 

147 loc='upper left', 

148 ) 

149 plt_helper.plt.ylim([0, 15]) 

150 

151 fname = 'data/harmonic_qd_iterations' 

152 plt_helper.savefig(fname) 

153 

154 assert os.path.isfile(fname + '.pdf'), 'ERROR: plotting did not create PDF file' 

155 # assert os.path.isfile(fname + '.pgf'), 'ERROR: plotting did not create PGF file' 

156 assert os.path.isfile(fname + '.png'), 'ERROR: plotting did not create PNG file' 

157 

158 

159def main(): 

160 run_simulation() 

161 show_results() 

162 

163 

164if __name__ == "__main__": 

165 main()