Coverage for pySDC/implementations/sweeper_classes/boris_2nd_order.py: 100%

1import numpy as np 

2 

3from pySDC.core.errors import ParameterError 

4from pySDC.core.sweeper import Sweeper 

5 

6 

7class boris_2nd_order(Sweeper): 

8 """ 

9 SDC sweeper for charged particles in electromagnetic fields, with velocity-Verlet and a Boris push as base. 

10 

11 Second-order sweeper using velocity-Verlet with Boris scheme as base integrator 

12 

13 Attributes: 

14 S: node-to-node collocation matrix (first order) 

15 SQ: node-to-node collocation matrix (second order) 

16 ST: node-to-node trapezoidal matrix 

17 Sx: node-to-node Euler half-step for position update 

18 """ 

19 

20 def __init__(self, params, level): 

21 """ 

22 Initialization routine for the custom sweeper 

23 

24 Args: 

25 params: parameters for the sweeper 

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

27 """ 

28 

29 if "QI" not in params: 

30 params["QI"] = "IE" 

31 if "QE" not in params: 

32 params["QE"] = "EE" 

33 

34 super(boris_2nd_order, self).__init__(params, level) 

35 

36 # S- and SQ-matrices (derived from Q) and Sx- and ST-matrices for the integrator 

37 [ 

38 self.S, 

39 self.ST, 

40 self.SQ, 

41 self.Sx, 

42 self.QQ, 

43 self.QI, 

44 self.QT, 

45 self.Qx, 

46 self.Q, 

47 ] = self.__get_Qd() 

48 

49 # the Boris step adds back the trapezoid dt*d/2*(f[m]+f[m+1]) with d = diag(QI), so ST must be exactly that, 

50 # otherwise the sweep converges to a wrong fixed point 

51 M = self.coll.num_nodes 

52 d = np.diag(self.QI)[1:] 

53 T = np.zeros_like(self.ST) 

54 T[range(1, M + 1), range(1, M + 1)] = d / 2 

55 T[range(1, M + 1), range(M)] += d / 2 

56 # For a working pair ST == T in exact arithmetic. In floating point, QI and QE round the same node distances 

57 # independently (at most 1, so up to eps apart, eps / 2 after halving), and forming QT = (QI + QE) / 2 and its 

58 # row differences rounds twice more (eps / 2 each): |ST - T| <= 2 eps. The bar is twice that; a wrong pair is 

59 # off by O(0.1). 

60 if not np.allclose(self.ST[1:], T[1:], rtol=0, atol=4 * np.finfo(float).eps): 

61 raise ParameterError( 

62 f'boris_2nd_order converges to a wrong solution with QI={self.params.QI!r} and QE={self.params.QE!r}: ' 

63 'the velocity update needs QE to supply the trapezoid partner of the diagonal of QI. Use QI="IE" ' 

64 'with QE="EE" (the default), or QI="PIC" with QE="PIC".' 

65 ) 

66 

67 self.qQ = np.dot(self.coll.weights, self.coll.Qmat[1:, 1:]) 

68 

69 def __get_Qd(self): 

70 """ 

71 Get integration matrices for 2nd-order SDC 

72 

73 Returns: 

74 S: node-to-node collocation matrix (first order) 

75 SQ: node-to-node collocation matrix (second order) 

76 ST: node-to-node trapezoidal matrix 

77 Sx: node-to-node Euler half-step for position update 

78 """ 

79 

80 # set implicit and explicit Euler matrices (default, but can be changed) 

81 QI = self.get_Qdelta_implicit(qd_type=self.params.QI) 

82 QE = self.get_Qdelta_explicit(qd_type=self.params.QE) 

83 

84 # trapezoidal rule 

85 QT = 1 / 2 * (QI + QE) 

86 

87 # Qx as in the paper 

88 Qx = np.dot(QE, QT) + 1 / 2 * QE * QE 

89 

90 Sx = np.zeros(np.shape(self.coll.Qmat)) 

91 ST = np.zeros(np.shape(self.coll.Qmat)) 

92 S = np.zeros(np.shape(self.coll.Qmat)) 

93 

94 # fill-in node-to-node matrices 

95 Sx[0, :] = Qx[0, :] 

96 ST[0, :] = QT[0, :] 

97 S[0, :] = self.coll.Qmat[0, :] 

98 for m in range(self.coll.num_nodes): 

99 Sx[m + 1, :] = Qx[m + 1, :] - Qx[m, :] 

100 ST[m + 1, :] = QT[m + 1, :] - QT[m, :] 

101 S[m + 1, :] = self.coll.Qmat[m + 1, :] - self.coll.Qmat[m, :] 

102 # SQ via dot-product, could also be done via QQ 

103 SQ = np.dot(S, self.coll.Qmat) 

104 

105 # QQ-matrix via product of Q 

106 QQ = np.dot(self.coll.Qmat, self.coll.Qmat) 

107 

108 return [S, ST, SQ, Sx, QQ, QI, QT, Qx, self.coll.Qmat] 

109 

110 def update_nodes(self): 

111 """ 

112 Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes 

113 

114 Returns: 

115 None 

116 """ 

117 

118 # get current level and problem description 

119 L = self.level 

120 P = L.prob 

121 

122 # only if the level has been touched before 

123 assert L.status.unlocked 

124 

125 # get number of collocation nodes for easier access 

126 M = self.coll.num_nodes 

127 

128 # initialize integral terms with zeros, will add stuff later 

129 integral = [P.dtype_u(P.init, val=0.0) for l in range(M)] 

130 

131 # gather all terms which are known already (e.g. from the previous iteration) 

132 # this corresponds to SF(u^k) - SdF(u^k) + tau (note: have integrals in pos and vel!) 

133 for m in range(M): 

134 for j in range(M + 1): 

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

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

137 # add SQF(u^k) - SxF(u^k) for the position 

138 integral[m].pos += L.dt * (L.dt * (self.SQ[m + 1, j] - self.Sx[m + 1, j]) * f) 

139 # add SF(u^k) - STF(u^k) for the velocity 

140 integral[m].vel += L.dt * (self.S[m + 1, j] - self.ST[m + 1, j]) * f 

141 # add tau if associated 

142 if L.tau[m] is not None: 

143 integral[m] += L.tau[m] 

144 # tau is 0-to-node, need to change it to node-to-node here 

145 if m > 0: 

146 integral[m] -= L.tau[m - 1] 

147 

148 # do the sweep 

149 for m in range(0, M): 

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

151 tmp = P.dtype_u(integral[m]) 

152 for j in range(m + 1): 

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

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

155 # add SxF(u^{k+1}) 

156 tmp.pos += L.dt * (L.dt * self.Sx[m + 1, j] * f) 

157 # add pos at previous node + dt*v0 

158 tmp.pos += L.u[m].pos + L.dt * self.coll.delta_m[m] * L.u[0].vel 

159 # set new position, is explicit 

160 L.u[m + 1].pos = tmp.pos 

161 

162 # get E field with new positions and compute mean 

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

164 

165 ck = tmp.vel 

166 

167 # do the boris scheme 

168 L.u[m + 1].vel = P.boris_solver(ck, L.dt * np.diag(self.QI)[m + 1], L.f[m], L.f[m + 1], L.u[m]) 

169 

170 # indicate presence of new values at this level 

171 L.status.updated = True 

172 

173 return None 

174 

175 def integrate(self): 

176 """ 

177 Integrates the right-hand side 

178 

179 Returns: 

180 list of dtype_u: containing the integral as values 

181 """ 

182 

183 # get current level and problem description 

184 L = self.level 

185 P = L.prob 

186 

187 # create new instance of dtype_u, initialize values with 0 

188 p = [] 

189 

190 for m in range(1, self.coll.num_nodes + 1): 

191 p.append(P.dtype_u(P.init, val=0.0)) 

192 

193 # integrate RHS over all collocation nodes, RHS is here only f(x,v)! 

194 for j in range(1, self.coll.num_nodes + 1): 

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

196 p[-1].pos += L.dt * (L.dt * self.QQ[m, j] * f) + L.dt * self.coll.Qmat[m, j] * L.u[0].vel 

197 p[-1].vel += L.dt * self.coll.Qmat[m, j] * f 

198 

199 return p 

200 

201 def compute_end_point(self): 

202 """ 

203 Compute u at the right point of the interval 

204 

205 The value uend computed here is a full evaluation of the Picard formulation (always!) 

206 

207 Returns: 

208 None 

209 """ 

210 

211 # get current level and problem description 

212 L = self.level 

213 P = L.prob 

214 

215 # start with u0 and add integral over the full interval (using coll.weights) 

216 L.uend = P.dtype_u(L.u[0]) 

217 for m in range(self.coll.num_nodes): 

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

219 L.uend.pos += L.dt * (L.dt * self.qQ[m] * f) + L.dt * self.coll.weights[m] * L.u[0].vel 

220 L.uend.vel += L.dt * self.coll.weights[m] * f 

221 # add up tau correction of the full interval (last entry) 

222 if L.tau[-1] is not None: 

223 L.uend += L.tau[-1] 

224 

225 return None 

226 

227 def get_sweeper_mats(self): 

228 """ 

229 Returns the matrices Q, QQ, Qx, QT which define the sweeper. 

230 """ 

231 

232 Q = self.Q[1:, 1:] 

233 QQ = self.QQ[1:, 1:] 

234 Qx = self.Qx[1:, 1:] 

235 QT = self.QT[1:, 1:] 

236 

237 return Q, QQ, Qx, QT 

238 

239 def get_scalar_problems_sweeper_mats(self, lambdas=None): 

240 """ 

241 This function returns the corresponding matrices of an SDC sweep matrix formulation 

242 

243 Args: 

244 lambdas (numpy.narray): the first entry in lambdas is k-spring constant and the second is mu friction. 

245 """ 

246 

247 Q, QQ, Qx, QT = self.get_sweeper_mats() 

248 

249 if lambdas is None: 

250 pass 

251 # should use lambdas from attached problem and make sure it is scalar SDC 

252 raise NotImplementedError("At the moment, the values for the lambda have to be provided") 

253 else: 

254 k = lambdas[0] 

255 mu = lambdas[1] 

256 

257 nnodes = self.coll.num_nodes 

258 dt = self.level.dt 

259 

260 F = np.block( 

261 [ 

262 [-k * np.eye(nnodes), -mu * np.eye(nnodes)], 

263 [-k * np.eye(nnodes), -mu * np.eye(nnodes)], 

264 ] 

265 ) 

266 

267 C_coll = np.block([[np.eye(nnodes), dt * Q], [np.zeros([nnodes, nnodes]), np.eye(nnodes)]]) 

268 Q_coll = np.block( 

269 [ 

270 [dt**2 * QQ, np.zeros([nnodes, nnodes])], 

271 [np.zeros([nnodes, nnodes]), dt * Q], 

272 ] 

273 ) 

274 Q_vv = np.block( 

275 [ 

276 [dt**2 * Qx, np.zeros([nnodes, nnodes])], 

277 [np.zeros([nnodes, nnodes]), dt * QT], 

278 ] 

279 ) 

280 M_vv = np.eye(2 * nnodes) - np.dot(Q_vv, F) 

281 

282 return C_coll, Q_coll, Q_vv, M_vv, F 

283 

284 def get_scalar_problems_manysweep_mats(self, nsweeps, lambdas=None): 

285 """ 

286 For a scalar problem, K sweeps of SDC can be written in matrix form. 

287 

288 Args: 

289 nsweeps (int): number of sweeps 

290 lambdas (numpy.ndarray): the first entry in lambdas is k-spring constant and the second is mu friction. 

291 """ 

292 nnodes = self.coll.num_nodes 

293 

294 C_coll, Q_coll, Q_vv, M_vv, F = self.get_scalar_problems_sweeper_mats(lambdas=lambdas) 

295 

296 K_sdc = np.dot(np.linalg.inv(M_vv), Q_coll - Q_vv) @ F 

297 

298 Keig, Kvec = np.linalg.eig(K_sdc) 

299 

300 Kp_sdc = np.linalg.matrix_power(K_sdc, nsweeps) 

301 Kinv_sdc = np.linalg.inv(np.eye(2 * nnodes) - K_sdc) 

302 

303 Kdot_sdc = np.dot(np.eye(2 * nnodes) - Kp_sdc, Kinv_sdc) 

304 MC = np.dot(np.linalg.inv(M_vv), C_coll) 

305 

306 Mat_sweep = Kp_sdc + np.dot(Kdot_sdc, MC) 

307 

308 return Mat_sweep, np.max(np.abs(Keig)) 

309 

310 def get_scalar_problems_picardsweep_mats(self, nsweeps, lambdas=None): 

311 """ 

312 For a scalar problem, K sweeps of SDC can be written in matrix form. 

313 

314 Args: 

315 nsweeps (int): number of sweeps 

316 lambdas (numpy.ndarray): the first entry in lambdas is k-spring constant and the second is mu friction. 

317 """ 

318 nnodes = self.coll.num_nodes 

319 

320 C_coll, Q_coll, Q_vv, M_vv, F = self.get_scalar_problems_sweeper_mats(lambdas=lambdas) 

321 

322 K_sdc = np.dot(Q_coll, F) 

323 

324 Keig, Kvec = np.linalg.eig(K_sdc) 

325 

326 Kp_sdc = np.linalg.matrix_power(K_sdc, nsweeps) 

327 Kinv_sdc = np.linalg.inv(np.eye(2 * nnodes) - K_sdc) 

328 

329 Kdot_sdc = np.dot(np.eye(2 * nnodes) - Kp_sdc, Kinv_sdc) 

330 

331 Mat_sweep = Kp_sdc + np.dot(Kdot_sdc, C_coll) 

332 

333 return Mat_sweep, np.max(np.abs(Keig))