Coverage for pySDC/projects/Second_orderSDC/penningtrap_Simulation.py: 76%

1import numpy as np 

2 

3from pySDC.helpers.stats_helper import get_sorted 

4from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

5from pySDC.projects.Second_orderSDC.penningtrap_HookClass import particles_output 

6from pySDC.implementations.sweeper_classes.Runge_Kutta_Nystrom import RKN, Velocity_Verlet 

7from pySDC.projects.Second_orderSDC.plot_helper import PlotManager 

8 

9 

10class ComputeError(PlotManager): 

11 """ 

12 This class generates data for plots and computations for Second-order SDC 

13 """ 

14 

15 def __init__(self, controller_params, description, time_iter=3, K_iter=(1, 2, 3), Tend=2, axes=(1,), cwd=''): 

16 super().__init__( 

17 controller_params, description, time_iter=time_iter, K_iter=K_iter, Tend=Tend, axes=axes, cwd='' 

18 ) 

19 

20 def run_local_error(self): 

21 """ 

22 Controls everything to generate local convergence rate 

23 """ 

24 self.compute_local_error_data() 

25 self.plot_convergence() 

26 

27 def run_global_error(self): 

28 """ 

29 Computes global convergence order and finds approximate order 

30 """ 

31 self.error_type = 'global' 

32 self.compute_global_error_data() 

33 self.find_approximate_order(filename='data/dt_vs_global_errorSDC.csv') 

34 self.plot_convergence() 

35 

36 def run_work_precision(self, RK=True, VV=False, dt_cont=1): 

37 """ 

38 Implements work-precision of Second-order SDC 

39 Args: 

40 RK: True or False to include RKN method 

41 VV: True or False to include Velocity-Verlet Scheme 

42 dt_cont: moves RK and VV left to right 

43 """ 

44 self.RK = RK 

45 self.VV = VV 

46 self.compute_global_error_data(work_counter=True) 

47 self.compute_global_error_data(Picard=True, work_counter=True) 

48 if self.RK: 

49 self.compute_global_error_data(RK=RK, work_counter=True, dt_cont=dt_cont) 

50 if self.VV: 

51 self.compute_global_error_data(VV=VV, work_counter=True, dt_cont=dt_cont) 

52 self.plot_work_precision() 

53 

54 def compute_local_error_data(self): 

55 """ 

56 Computes local convergence rate and saves the data 

57 """ 

58 step_params = dict() 

59 dt_val = self.description['level_params']['dt'] 

60 

61 for order in self.K_iter: 

62 step_params['maxiter'] = order 

63 self.description['step_params'] = step_params 

64 

65 file_path = self.cwd + 'data/dt_vs_local_errorSDC.csv' 

66 mode = 'w' if order == self.K_iter[0] else 'a' 

67 

68 with open(file_path, mode) as file: 

69 if order == self.K_iter[0]: 

70 file.write("Time_steps | Order_pos | Abs_error_position | Order_vel | Abs_error_velocity\n") 

71 

72 for ii in range(0, self.time_iter): 

73 dt = dt_val / 2**ii 

74 

75 self.description['level_params']['dt'] = dt 

76 self.description['level_params'] = self.description['level_params'] 

77 

78 controller = controller_nonMPI( 

79 num_procs=1, controller_params=self.controller_params, description=self.description 

80 ) 

81 

82 t0, Tend = 0.0, dt 

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

84 uinit = P.u_init() 

85 

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

87 

88 uex = P.u_exact(Tend) 

89 coll_order = controller.MS[0].levels[0].sweep.coll.order 

90 order_pos = list(self.local_order_pos(order, coll_order)) 

91 order_vel = list(self.local_order_vel(order, coll_order)) 

92 error_pos = list(np.abs((uex - uend).pos).T[0]) 

93 error_vel = list(np.abs((uex - uend).vel).T[0]) 

94 

95 dt_omega = dt * self.description['problem_params']['omega_B'] 

96 file.write( 

97 f"{dt_omega}, {', '.join(map(str, order_pos))}, {', '.join(map(str, error_pos))}," 

98 f" {', '.join(map(str, order_vel))}, {', '.join(map(str, error_vel))}\n" 

99 ) 

100 

101 def compute_global_error_data(self, Picard=False, RK=False, VV=False, work_counter=False, dt_cont=1): 

102 """ 

103 Computes global convergence data and saves it into the data folder 

104 Args: 

105 Picard: bool, Picard iteration computation 

106 RK: bool, RKN method 

107 VV: bool, Velocity-Verlet scheme 

108 work_counter: bool, compute rhs for work precision 

109 dt_cont: moves the data left to right for RK and VV method 

110 """ 

111 K_iter = self.K_iter 

112 name, description = '', self.description 

113 

114 if Picard: 

115 name = 'Picard' 

116 description['sweeper_params']['QI'] = 'PIC' 

117 description['sweeper_params']['QE'] = 'PIC' 

118 elif RK: 

119 K_iter, name, description['sweeper_class'] = (1,), 'RKN', RKN 

120 elif VV: 

121 K_iter, name, description['sweeper_class'] = (1,), 'VV', Velocity_Verlet 

122 else: 

123 name = 'SDC' 

124 

125 self.controller_params['hook_class'] = particles_output 

126 step_params, dt_val = dict(), self.description['level_params']['dt'] 

127 values, error = ['position', 'velocity'], dict() 

128 

129 filename = f"data/{'rhs_eval_vs_global_error' if work_counter else 'dt_vs_global_error'}{name}.csv" 

130 

131 for order in K_iter: 

132 u_val, uex_val = dict(), dict() 

133 step_params['maxiter'], description['step_params'] = order, step_params 

134 

135 file_path = self.cwd + filename 

136 mode = 'w' if order == K_iter[0] else 'a' 

137 

138 with open(file_path, mode) as file: 

139 if order == K_iter[0]: 

140 file.write( 

141 "Time_steps/Work_counter | Order_pos | Abs_error_position | Order_vel | Abs_error_velocity\n" 

142 ) 

143 

144 cont = 2 if self.time_iter == 3 else 2 ** abs(3 - self.time_iter) 

145 cont = cont if not Picard else dt_cont 

146 

147 for ii in range(0, self.time_iter): 

148 dt = (dt_val * cont) / 2**ii 

149 

150 description['level_params']['dt'] = dt 

151 description['level_params'] = self.description['level_params'] 

152 

153 controller = controller_nonMPI( 

154 num_procs=1, controller_params=self.controller_params, description=description 

155 ) 

156 

157 t0, Tend = 0.0, self.Tend 

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

159 uinit = P.u_init() 

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

161 

162 func_eval = P.work_counters['Boris_solver'].niter + P.work_counters['rhs'].niter 

163 

164 for nn in values: 

165 u_val[nn] = get_sorted(stats, type=nn, sortby='time') 

166 uex_val[nn] = get_sorted(stats, type=nn + '_exact', sortby='time') 

167 error[nn] = self.relative_error(uex_val[nn], u_val[nn]) 

168 error[nn] = list(error[nn].T[0]) 

169 

170 if RK or VV: 

171 global_order = np.array([4, 4]) 

172 else: 

173 coll_order = controller.MS[0].levels[0].sweep.coll.order 

174 global_order = list(self.global_order(order, coll_order)) 

175 global_order = np.array([4, 4]) if RK or VV else list(self.global_order(order, coll_order)) 

176 

177 dt_omega = dt * self.description['problem_params']['omega_B'] 

178 save = func_eval if work_counter else dt_omega 

179 

180 file.write( 

181 f"{save}, {', '.join(map(str, global_order))}, {', '.join(map(str, error['position']))}," 

182 f" {', '.join(map(str, global_order))}, {', '.join(map(str, error['velocity']))}\n" 

183 ) 

184 

185 def local_order_pos(self, order_K, order_quad): 

186 if self.description['sweeper_params']['initial_guess'] == 'spread': 

187 if self.quad_type == 'GAUSS' or self.quad_type == 'RADAU-RIGHT': 

188 return np.array([np.min([order_K + 2 + 2, order_quad]), np.min([2 * order_K + 3, order_quad])]) 

189 elif self.quad_type == 'LOBATTO' or self.quad_type == 'RADAU-LEFT': 

190 return np.array([np.min([order_K + 2 + 2, order_quad]), np.min([2 * order_K + 3, order_quad])]) 

191 else: 

192 raise NotImplementedError('Order of convergence explicitly not implemented') 

193 else: 

194 if self.quad_type == 'GAUSS' or self.quad_type == 'RADAU-RIGHT': 

195 return np.array([np.min([order_K + 2, order_quad]), np.min([2 * order_K + 3, order_quad])]) 

196 elif self.quad_type == 'LOBATTO' or self.quad_type == 'RADAU-LEFT': 

197 return np.array([np.min([order_K + 2, order_quad]), np.min([2 * order_K + 3, order_quad])]) 

198 else: 

199 raise NotImplementedError('Order of convergence explicitly not implemented') 

200 

201 def local_order_vel(self, order_K, order_quad): 

202 if self.description['sweeper_params']['initial_guess'] == 'spread': 

203 if self.quad_type == 'GAUSS' or self.quad_type == 'RADAU-RIGHT': 

204 return np.array([np.min([order_K + 1 + 2, order_quad]), np.min([2 * order_K + 2, order_quad])]) 

205 elif self.quad_type == 'LOBATTO' or self.quad_type == 'RADAU-LEFT': 

206 return np.array([np.min([order_K + 1 + 2, order_quad]), np.min([2 * order_K + 2, order_quad])]) 

207 else: 

208 raise NotImplementedError('Order of convergence explicitly not implemented') 

209 else: 

210 if self.quad_type == 'GAUSS' or self.quad_type == 'RADAU-RIGHT': 

211 return np.array([np.min([order_K + 1, order_quad]), np.min([2 * order_K + 2, order_quad])]) 

212 elif self.quad_type == 'LOBATTO' or self.quad_type == 'RADAU-LEFT': 

213 return np.array([np.min([order_K + 1, order_quad]), np.min([2 * order_K + 2, order_quad])]) 

214 else: 

215 raise NotImplementedError('Order of convergence explicitly not implemented') 

216 

217 def global_order(self, order_K, order_quad): 

218 if self.quad_type == 'GAUSS' or self.quad_type == 'RADAU-RIGHT': 

219 return np.array([np.min([order_K, order_quad]), np.min([2 * order_K, order_quad])]) 

220 elif self.quad_type == 'LOBATTO' or self.quad_type == 'RADAU-LEFT': 

221 return np.array([np.min([order_K, order_quad]), np.min([2 * order_K, order_quad])]) + 2 

222 else: 

223 raise NotImplementedError('Order of convergence explicitly not implemented') 

224 

225 def relative_error(self, uex_data, u_data): 

226 u_ex = np.array([entry[1] for entry in uex_data]) 

227 u = np.array([entry[1] for entry in u_data]) 

228 return np.linalg.norm(np.abs((u_ex - u)), np.inf, 0) / np.linalg.norm(u_ex, np.inf, 0)