Coverage for pySDC/projects/Hamiltonian/simple_problems.py: 100%

1import os 

2from collections import defaultdict 

3 

4import dill 

5import numpy as np 

6 

7import pySDC.helpers.plot_helper as plt_helper 

8from pySDC.helpers.stats_helper import get_sorted, filter_stats 

9 

10from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

11from pySDC.implementations.problem_classes.HarmonicOscillator import harmonic_oscillator 

12from pySDC.implementations.problem_classes.HenonHeiles import henon_heiles 

13from pySDC.implementations.sweeper_classes.verlet import verlet 

14from pySDC.implementations.transfer_classes.TransferParticles_NoCoarse import particles_to_particles 

15from pySDC.projects.Hamiltonian.hamiltonian_output import hamiltonian_output 

16 

17 

18def setup_harmonic(): 

19 """ 

20 Helper routine for setting up everything for the harmonic oscillator 

21 

22 Returns: 

23 description (dict): description of the controller 

24 controller_params (dict): controller parameters 

25 """ 

26 

27 # initialize level parameters 

28 level_params = dict() 

29 level_params['restol'] = 1e-10 

30 level_params['dt'] = 0.5 

31 

32 # initialize sweeper parameters 

33 sweeper_params = dict() 

34 sweeper_params['quad_type'] = 'LOBATTO' 

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

36 sweeper_params['initial_guess'] = 'zero' 

37 

38 # initialize problem parameters for the Penning trap 

39 problem_params = dict() 

40 problem_params['k'] = 1.0 

41 problem_params['phase'] = 0.0 

42 problem_params['amp'] = 1.0 

43 

44 # initialize step parameters 

45 step_params = dict() 

46 step_params['maxiter'] = 50 

47 

48 # initialize controller parameters 

49 controller_params = dict() 

50 controller_params['hook_class'] = hamiltonian_output # specialized hook class for more statistics and output 

51 controller_params['logger_level'] = 30 

52 

53 # Fill description dictionary for easy hierarchy creation 

54 description = dict() 

55 description['problem_class'] = harmonic_oscillator 

56 description['problem_params'] = problem_params 

57 description['sweeper_class'] = verlet 

58 description['sweeper_params'] = sweeper_params 

59 description['level_params'] = level_params 

60 description['step_params'] = step_params 

61 description['space_transfer_class'] = particles_to_particles 

62 

63 return description, controller_params 

64 

65 

66def setup_henonheiles(): 

67 """ 

68 Helper routine for setting up everything for the Henon Heiles problem 

69 

70 Returns: 

71 description (dict): description of the controller 

72 controller_params (dict): controller parameters 

73 """ 

74 

75 # initialize level parameters 

76 level_params = dict() 

77 level_params['restol'] = 1e-10 

78 level_params['dt'] = 0.25 

79 

80 # initialize sweeper parameters 

81 sweeper_params = dict() 

82 sweeper_params['quad_type'] = 'LOBATTO' 

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

84 sweeper_params['initial_guess'] = 'zero' 

85 

86 # initialize problem parameters for the Penning trap 

87 problem_params = dict() 

88 

89 # initialize step parameters 

90 step_params = dict() 

91 step_params['maxiter'] = 50 

92 

93 # initialize controller parameters 

94 controller_params = dict() 

95 controller_params['hook_class'] = hamiltonian_output # specialized hook class for more statistics and output 

96 controller_params['logger_level'] = 30 

97 

98 # Fill description dictionary for easy hierarchy creation 

99 description = dict() 

100 description['problem_class'] = henon_heiles 

101 description['problem_params'] = problem_params 

102 description['sweeper_class'] = verlet 

103 description['sweeper_params'] = sweeper_params 

104 description['level_params'] = level_params 

105 description['step_params'] = step_params 

106 description['space_transfer_class'] = particles_to_particles 

107 

108 return description, controller_params 

109 

110 

111def run_simulation(prob=None): 

112 """ 

113 Routine to run the simulation of a second order problem 

114 

115 Args: 

116 prob (str): name of the problem 

117 

118 """ 

119 

120 # check what problem type we have and set up corresponding description and variables 

121 if prob == 'harmonic': 

122 description, controller_params = setup_harmonic() 

123 # set time parameters 

124 t0 = 0.0 

125 Tend = 50.0 

126 num_procs = 100 

127 maxmeaniter = 6.5 

128 elif prob == 'henonheiles': 

129 description, controller_params = setup_henonheiles() 

130 # set time parameters 

131 t0 = 0.0 

132 Tend = 25.0 

133 num_procs = 100 

134 maxmeaniter = 5.0 

135 else: 

136 raise NotImplementedError('Problem type not implemented, got %s' % prob) 

137 

138 # instantiate the controller 

139 controller = controller_nonMPI(num_procs=num_procs, controller_params=controller_params, description=description) 

140 

141 # get initial values on finest level 

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

143 uinit = P.u_exact(t=t0) 

144 

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

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

147 

148 # filter statistics by type (number of iterations) 

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

150 

151 # compute and print statistics 

152 # for item in iter_counts: 

153 # out = 'Number of iterations for time %4.2f: %2i' % item 

154 # print(out) 

155 

156 niters = np.array([item[1] for item in iter_counts]) 

157 out = ' Mean number of iterations: %4.2f' % np.mean(niters) 

158 print(out) 

159 out = ' Range of values for number of iterations: %2i ' % np.ptp(niters) 

160 print(out) 

161 out = ' Position of max/min number of iterations: %2i -- %2i' % (int(np.argmax(niters)), int(np.argmin(niters))) 

162 print(out) 

163 out = ' Std and var for number of iterations: %4.2f -- %4.2f' % (float(np.std(niters)), float(np.var(niters))) 

164 print(out) 

165 

166 assert np.mean(niters) <= maxmeaniter, 'Mean number of iterations is too high, got %s' % np.mean(niters) 

167 

168 fname = 'data/' + prob + '.dat' 

169 f = open(fname, 'wb') 

170 dill.dump(stats, f) 

171 f.close() 

172 

173 assert os.path.isfile(fname), 'Run for %s did not create stats file' % prob 

174 

175 

176def show_results(prob=None, cwd=''): 

177 """ 

178 Helper function to plot the error of the Hamiltonian 

179 

180 Args: 

181 prob (str): name of the problem 

182 cwd (str): current working directory 

183 """ 

184 

185 # read in the dill data 

186 f = open(cwd + 'data/' + prob + '.dat', 'rb') 

187 stats = dill.load(f) 

188 f.close() 

189 

190 # extract error in hamiltonian and prepare for plotting 

191 extract_stats = filter_stats(stats, type='err_hamiltonian') 

192 result = defaultdict(list) 

193 for k, v in extract_stats.items(): 

194 result[k.iter].append((k.time, v)) 

195 for k, _ in result.items(): 

196 result[k] = sorted(result[k], key=lambda x: x[0]) 

197 

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

199 plt_helper.setup_mpl() 

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

201 

202 # Rearrange data for easy plotting 

203 err_ham = 1 

204 for k, v in result.items(): 

205 time = [item[0] for item in v] 

206 ham = [item[1] for item in v] 

207 err_ham = ham[-1] 

208 plt_helper.plt.semilogy(time, ham, '-', lw=1, label='Iter ' + str(k)) 

209 

210 assert err_ham < 3.7e-08, 'Error in the Hamiltonian is too large for %s, got %s' % (prob, err_ham) 

211 

212 plt_helper.plt.xlabel('Time') 

213 plt_helper.plt.ylabel('Error in Hamiltonian') 

214 plt_helper.plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) 

215 

216 fname = 'data/' + prob + '_hamiltonian' 

217 plt_helper.savefig(fname) 

218 

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

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

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

222 

223 

224def main(): 

225 prob = 'harmonic' 

226 run_simulation(prob) 

227 show_results(prob) 

228 prob = 'henonheiles' 

229 run_simulation(prob) 

230 show_results(prob) 

231 

232 

233if __name__ == "__main__": 

234 main()