Coverage for pySDC/implementations/sweeper_classes/Runge_Kutta_Nystrom.py: 96%

1import numpy as np 

2import logging 

3 

4from pySDC.core.sweeper import Sweeper, _Pars 

5from pySDC.core.errors import ParameterError 

6from pySDC.implementations.datatype_classes.particles import particles, fields, acceleration 

7from pySDC.implementations.sweeper_classes.Runge_Kutta import RungeKutta 

8 

9 

10class ButcherTableauNoCollUpdate(object): 

11 """Version of Butcher Tableau that does not need a collocation update because the weights are put in the last line of Q""" 

12 

13 def __init__(self, weights, nodes, matrix): 

14 """ 

15 Initialization routine to get a quadrature matrix out of a Butcher tableau 

16 

17 Args: 

18 weights (numpy.ndarray): Butcher tableau weights 

19 nodes (numpy.ndarray): Butcher tableau nodes 

20 matrix (numpy.ndarray): Butcher tableau entries 

21 """ 

22 # check if the arguments have the correct form 

23 if type(matrix) != np.ndarray: 

24 raise ParameterError('Runge-Kutta matrix needs to be supplied as a numpy array!') 

25 elif len(np.unique(matrix.shape)) != 1 or len(matrix.shape) != 2: 

26 raise ParameterError('Runge-Kutta matrix needs to be a square 2D numpy array!') 

27 

28 if type(weights) != np.ndarray: 

29 raise ParameterError('Weights need to be supplied as a numpy array!') 

30 elif len(weights.shape) != 1: 

31 raise ParameterError(f'Incompatible dimension of weights! Need 1, got {len(weights.shape)}') 

32 elif len(weights) != matrix.shape[0]: 

33 raise ParameterError(f'Incompatible number of weights! Need {matrix.shape[0]}, got {len(weights)}') 

34 

35 if type(nodes) != np.ndarray: 

36 raise ParameterError('Nodes need to be supplied as a numpy array!') 

37 elif len(nodes.shape) != 1: 

38 raise ParameterError(f'Incompatible dimension of nodes! Need 1, got {len(nodes.shape)}') 

39 elif len(nodes) != matrix.shape[0]: 

40 raise ParameterError(f'Incompatible number of nodes! Need {matrix.shape[0]}, got {len(nodes)}') 

41 

42 self.globally_stiffly_accurate = np.allclose(matrix[-1], weights) 

43 

44 self.tleft = 0.0 

45 self.tright = 1.0 

46 self.num_solution_stages = 0 if self.globally_stiffly_accurate else 1 

47 self.num_nodes = matrix.shape[0] + self.num_solution_stages 

48 self.weights = weights 

49 

50 if self.globally_stiffly_accurate: 

51 # For globally stiffly accurate methods, the last row of the Butcher tableau is the same as the weights. 

52 self.nodes = np.append([0], nodes) 

53 self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1]) 

54 self.Qmat[1:, 1:] = matrix 

55 else: 

56 self.nodes = np.append(np.append([0], nodes), [1]) 

57 self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1]) 

58 self.Qmat[1:-1, 1:-1] = matrix 

59 self.Qmat[-1, 1:-1] = weights # this is for computing the solution to the step from the previous stages 

60 

61 self.left_is_node = True 

62 self.right_is_node = self.nodes[-1] == self.tright 

63 

64 # compute distances between the nodes 

65 if self.num_nodes > 1: 

66 self.delta_m = self.nodes[1:] - self.nodes[:-1] 

67 else: 

68 self.delta_m = np.zeros(1) 

69 self.delta_m[0] = self.nodes[0] - self.tleft 

70 

71 # check if the RK scheme is implicit 

72 self.implicit = any(matrix[i, i] != 0 for i in range(self.num_nodes - self.num_solution_stages)) 

73 

74 

75class RungeKuttaNystrom(RungeKutta): 

76 """ 

77 Base class for Runge-Kutta-Nystrom methods for second-order problems on particle data, as a sweeper. 

78 Actually, the sweeper idea fits the Runge-Kutta idea when using only lower triangular rules, where solutions 

79 at the nodes are successively computed from earlier nodes. However, we only perform a single iteration of this. 

80 

81 We have two choices to realise a Runge-Kutta sweeper: We can choose Q = Q_Delta = <Butcher tableau>, but in this 

82 implementation, that would lead to a lot of wasted FLOPS from integrating with Q and then with Q_Delta and 

83 subtracting the two. For that reason, we built this new sweeper, which does not have a preconditioner. 

84 

85 This class only supports lower triangular Butcher tableaux such that the system can be solved with forward 

86 substitution. In this way, we don't get the maximum order that we could for the number of stages, but computing the 

87 stages is much cheaper. In particular, if the Butcher tableaux is strictly lower triangular, we get an explicit 

88 method, which does not require us to solve a system of equations to compute the stages. 

89 

90 Please be aware that all fundamental parameters of the Sweeper are ignored. These include 

91 

92 - num_nodes 

93 - collocation_class 

94 - initial_guess 

95 - QI 

96 

97 All of these variables are either determined by the RK rule, or are not part of an RK scheme. 

98 

99 Attribues: 

100 butcher_tableau (ButcherTableauNoCollUpdate): Butcher tableau for the Runge-Kutta scheme that you want 

101 """ 

102 

103 ButcherTableauClass = ButcherTableauNoCollUpdate 

104 weights_bar = None 

105 matrix_bar = None 

106 

107 def __init__(self, params, level): 

108 """ 

109 Initialization routine for the custom sweeper 

110 

111 Args: 

112 params: parameters for the sweeper 

113 level (pySDC.Level.level): the level that uses this sweeper 

114 """ 

115 super().__init__(params, level) 

116 self.coll_bar = self.get_Butcher_tableau_bar() 

117 self.Qx = self.coll_bar.Qmat 

118 

119 @classmethod 

120 def get_Butcher_tableau_bar(cls): 

121 """ 

122 Build the Butcher tableau for the positions from the class attributes ``weights_bar``, ``nodes`` and 

123 ``matrix_bar``. 

124 

125 Returns: 

126 ButcherTableauNoCollUpdate: an instance of ``ButcherTableauClass`` 

127 """ 

128 return cls.ButcherTableauClass(cls.weights_bar, cls.nodes, cls.matrix_bar) 

129 

130 def get_full_f(self, f): 

131 """ 

132 Test the right hand side function is the correct type 

133 

134 Args: 

135 f (dtype_f): Right hand side at a single node 

136 

137 Returns: 

138 mesh: Full right hand side as a mesh 

139 """ 

140 

141 if type(f) in [particles, fields, acceleration]: 

142 return f 

143 else: 

144 raise NotImplementedError(f'Type \"{type(f)}\" not implemented') 

145 

146 def update_nodes(self): 

147 """ 

148 Update the u- and f-values at the collocation nodes 

149 

150 Returns: 

151 None 

152 """ 

153 

154 # get current level and problem description 

155 L = self.level 

156 P = L.prob 

157 

158 # only if the level has been touched before 

159 assert L.status.unlocked 

160 assert L.status.sweep <= 1, "RK schemes are direct solvers. Please perform only 1 iteration!" 

161 

162 # get number of collocation nodes for easier access 

163 M = self.coll.num_nodes 

164 

165 for m in range(0, M): 

166 # build rhs, consisting of the known values from above and new values from previous nodes (at k+1) 

167 rhs = P.dtype_u(L.u[0]) 

168 rhs.pos += L.dt * self.coll.nodes[m + 1] * L.u[0].vel 

169 

170 for j in range(1, m + 1): 

171 # build RHS from f-terms (containing the E field) and the B field 

172 

173 f = P.build_f(L.f[j], L.u[j], L.time + L.dt * self.coll.nodes[j]) 

174 rhs.pos += L.dt**2 * self.Qx[m + 1, j] * self.get_full_f(f) 

175 """ 

176 

177 Implicit part only works for Velocity-Verlet scheme 

178 Boris solver for the implicit part 

179 

180 """ 

181 

182 if self.coll.implicit: 

183 ck = rhs.vel * 0.0 

184 L.f[3] = P.eval_f(rhs, L.time + L.dt) 

185 rhs.vel = P.boris_solver(ck, L.dt, L.f[0], L.f[3], L.u[0]) 

186 

187 else: 

188 rhs.vel += L.dt * self.QI[m + 1, j] * self.get_full_f(f) 

189 

190 # implicit solve with prefactor stemming from the diagonal of Qd 

191 L.u[m + 1] = rhs 

192 # update function values 

193 if self.coll.implicit: 

194 # That is why it only works for the Velocity-Verlet scheme 

195 L.f[0] = P.eval_f(L.u[0], L.time) 

196 L.f[m + 1] = P.dtype_f(L.f[0]) 

197 else: 

198 if m != self.coll.num_nodes - 1: 

199 L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m + 1]) 

200 

201 # indicate presence of new values at this level 

202 

203 L.status.updated = True 

204 

205 return None 

206 

207 def compute_end_point(self): 

208 """ 

209 In this Runge-Kutta implementation, the solution to the step is always stored in the last node 

210 """ 

211 self.level.uend = self.level.u[-1] 

212 

213 

214class RKN(RungeKuttaNystrom): 

215 """ 

216 Classical fourth-order Runge-Kutta-Nystrom method, with the nodes and weights of RK4. 

217 

218 Runge-Kutta-Nystrom method 

219 https://link.springer.com/book/10.1007/978-3-540-78862-1 

220 page: 284 

221 Chapter: II.14 Numerical methods for Second order differential equations 

222 """ 

223 

224 nodes = np.array([0.0, 0.5, 0.5, 1]) 

225 weights = np.array([1.0, 2.0, 2.0, 1.0]) / 6.0 

226 matrix = np.zeros([4, 4]) 

227 matrix[1, 0] = 0.5 

228 matrix[2, 1] = 0.5 

229 matrix[3, 2] = 1.0 

230 

231 weights_bar = np.array([1.0, 1.0, 1.0, 0]) / 6.0 

232 matrix_bar = np.zeros([4, 4]) 

233 matrix_bar[1, 0] = 1 / 8 

234 matrix_bar[2, 0] = 1 / 8 

235 matrix_bar[3, 2] = 1 / 2 

236 

237 

238class Velocity_Verlet(RungeKuttaNystrom): 

239 """ 

240 Second-order velocity-Verlet scheme as a Runge-Kutta-Nystrom method, using a Boris solver for the velocity. 

241 

242 Velocity-Verlet scheme 

243 https://de.wikipedia.org/wiki/Verlet-Algorithmus 

244 """ 

245 

246 nodes = np.array([1.0, 1.0]) 

247 weights = np.array([1 / 2, 0]) 

248 matrix = np.zeros([2, 2]) 

249 matrix[1, 1] = 1 

250 weights_bar = np.array([1 / 2, 0]) 

251 matrix_bar = np.zeros([2, 2])