Coverage for pySDC/projects/matrixPFASST/controller_matrix_nonMPI.py: 98%

1import numpy as np 

2 

3 

4from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

5from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit 

6 

7 

8class controller_matrix_nonMPI(controller_nonMPI): 

9 """ 

10 

11 PFASST controller, running serial matrix-based versions 

12 

13 """ 

14 

15 def __init__(self, num_procs, controller_params, description): 

16 """ 

17 Initialization routine for PFASST controller 

18 

19 Args: 

20 num_procs: number of parallel time steps (still serial, though), can be 1 

21 controller_params: parameter set for the controller and the steps 

22 description: all the parameters to set up the rest (levels, problems, transfer, ...) 

23 """ 

24 

25 assert description['sweeper_class'] is generic_implicit, ( 

26 'ERROR: matrix version will only work with generic_implicit sweeper, got %s' % description['sweeper_class'] 

27 ) 

28 

29 # call parent's initialization routine 

30 super(controller_matrix_nonMPI, self).__init__( 

31 num_procs=num_procs, controller_params=controller_params, description=description 

32 ) 

33 

34 self.nsteps = len(self.MS) 

35 self.nlevels = len(self.MS[0].levels) 

36 self.nnodes = self.MS[0].levels[0].sweep.coll.num_nodes 

37 self.nspace = self.MS[0].levels[0].prob.init[0] 

38 

39 self.dt = self.MS[0].levels[0].dt 

40 self.tol = self.MS[0].levels[0].params.restol 

41 self.maxiter = self.MS[0].params.maxiter 

42 

43 prob = self.MS[0].levels[0].prob 

44 

45 assert isinstance(self.nspace, int), 'ERROR: can only handle 1D data, got %s' % self.nspace 

46 assert [level.sweep.coll.right_is_node for step in self.MS for level in step.levels].count( 

47 True 

48 ) == self.nlevels * self.nsteps, 'ERROR: all collocation nodes have to be of Gauss-Radau type' 

49 assert self.nlevels <= 2, 'ERROR: cannot use matrix-PFASST with more than 2 levels' # TODO: fixme 

50 assert [level.dt for step in self.MS for level in step.levels].count( 

51 self.dt 

52 ) == self.nlevels * self.nsteps, 'ERROR: dt must be equal for all steps and all levels' 

53 # assert [level.sweep.coll.num_nodes for step in self.MS for level in step.levels].count(self.nnodes) == \ 

54 # self.nlevels * self.nsteps, 'ERROR: nnodes must be equal for all steps and all levels' 

55 assert [type(level.prob) for step in self.MS for level in step.levels].count( 

56 type(prob) 

57 ) == self.nlevels * self.nsteps, 'ERROR: all probem classes have to be the same' 

58 

59 assert self.params.predict_type is None, 'ERROR: no predictor for matrix controller yet' # TODO: fixme 

60 

61 assert hasattr(prob, 'A'), 'ERROR: need system matrix A for this (and linear problems!)' 

62 

63 A = prob.A.todense() 

64 Q = self.MS[0].levels[0].sweep.coll.Qmat[1:, 1:] 

65 Qd = self.MS[0].levels[0].sweep.QI[1:, 1:] 

66 

67 E = np.zeros((self.nsteps, self.nsteps)) 

68 np.fill_diagonal(E[1:, :], 1) 

69 

70 N = np.zeros((self.nnodes, self.nnodes)) 

71 N[:, -1] = 1 

72 

73 self.C = ( 

74 np.eye(self.nsteps * self.nnodes * self.nspace) 

75 - self.dt * np.kron(np.eye(self.nsteps), np.kron(Q, A)) 

76 - np.kron(E, np.kron(N, np.eye(self.nspace))) 

77 ) 

78 self.C = np.array(self.C) 

79 self.P = np.eye(self.nsteps * self.nnodes * self.nspace) - self.dt * np.kron( 

80 np.eye(self.nsteps), np.kron(Qd, A) 

81 ) 

82 self.P = np.array(self.P) 

83 

84 if self.nlevels > 1: 

85 prob_c = self.MS[0].levels[1].prob 

86 self.nspace_c = prob_c.init[0] 

87 

88 Ac = prob_c.A.todense() 

89 Qdc = self.MS[0].levels[1].sweep.QI[1:, 1:] 

90 nnodesc = self.MS[0].levels[1].sweep.coll.num_nodes 

91 Nc = np.zeros((nnodesc, nnodesc)) 

92 Nc[:, -1] = 1 

93 

94 if hasattr(self.MS[0].base_transfer.space_transfer, 'Pspace'): 

95 TcfA = self.MS[0].base_transfer.space_transfer.Pspace.todense() 

96 else: 

97 TcfA = np.eye(self.nspace_c) 

98 if hasattr(self.MS[0].base_transfer.space_transfer, 'Rspace'): 

99 TfcA = self.MS[0].base_transfer.space_transfer.Rspace.todense() 

100 else: 

101 TfcA = np.eye(self.nspace) 

102 TcfQ = self.MS[0].base_transfer.Pcoll 

103 TfcQ = self.MS[0].base_transfer.Rcoll 

104 

105 self.Tcf = np.array(np.kron(np.eye(self.nsteps), np.kron(TcfQ, TcfA))) 

106 self.Tfc = np.array(np.kron(np.eye(self.nsteps), np.kron(TfcQ, TfcA))) 

107 

108 self.Pc = ( 

109 np.eye(self.nsteps * nnodesc * self.nspace_c) 

110 - self.dt * np.kron(np.eye(self.nsteps), np.kron(Qdc, Ac)) 

111 - np.kron(E, np.kron(Nc, np.eye(self.nspace_c))) 

112 ) 

113 self.Pc = np.array(self.Pc) 

114 

115 self.u = np.zeros(self.nsteps * self.nnodes * self.nspace) 

116 self.res = np.zeros(self.nsteps * self.nnodes * self.nspace) 

117 self.u0 = np.zeros(self.nsteps * self.nnodes * self.nspace) 

118 

119 def run(self, u0, t0, Tend): 

120 """ 

121 Main driver for running the serial matrix version of SDC, MSSDC, MLSDC and PFASST 

122 

123 Args: 

124 u0: initial values 

125 t0: starting time 

126 Tend: ending time 

127 

128 Returns: 

129 end values on the finest level 

130 stats object containing statistics for each step, each level and each iteration 

131 """ 

132 

133 # some initializations and reset of statistics 

134 uend = None 

135 num_procs = len(self.MS) 

136 for hook in self.hooks: 

137 hook.reset_stats() 

138 

139 assert ( 

140 (Tend - t0) / self.dt 

141 ).is_integer(), 'ERROR: dt, t0, Tend were not chosen correctly, do not divide interval to be computed equally' 

142 

143 assert int((Tend - t0) / self.dt) % num_procs == 0, 'ERROR: num_procs not chosen correctly' 

144 

145 # initial ordering of the steps: 0,1,...,Np-1 

146 slots = list(range(num_procs)) 

147 

148 # initialize time variables of each step 

149 time = [t0 + sum(self.dt for _ in range(p)) for p in slots] 

150 

151 # initialize block of steps with u0 

152 self.restart_block(slots, time, u0) 

153 

154 # call pre-run hook 

155 for S in self.MS: 

156 for hook in self.hooks: 

157 hook.pre_run(step=S, level_number=0) 

158 

159 nblocks = int((Tend - t0) / self.dt / num_procs) 

160 

161 for _ in range(nblocks): 

162 self.MS = self.pfasst(self.MS) 

163 

164 for p in slots: 

165 time[p] += num_procs * self.dt 

166 

167 # uend is uend of the last active step in the list 

168 uend = self.MS[-1].levels[0].uend 

169 self.restart_block(slots, time, uend) 

170 

171 # call post-run hook 

172 for S in self.MS: 

173 for hook in self.hooks: 

174 hook.post_run(step=S, level_number=0) 

175 

176 return uend, self.return_stats() 

177 

178 def build_propagation_matrix(self, niter): 

179 """ 

180 Helper routine to create propagation matrix if requested 

181 

182 Args: 

183 niter: number of iterations 

184 

185 Returns: 

186 mat: propagation matrix 

187 """ 

188 

189 # build smoother iteration matrix and preconditioner using nsweeps 

190 Pinv = np.linalg.inv(self.P) 

191 precond_smoother = Pinv.copy() 

192 iter_mat_smoother = np.eye(self.nsteps * self.nnodes * self.nspace) - precond_smoother.dot(self.C) 

193 for k in range(1, self.MS[0].levels[0].params.nsweeps): 

194 precond_smoother += np.linalg.matrix_power(iter_mat_smoother, k).dot(Pinv) 

195 iter_mat_smoother = np.linalg.matrix_power(iter_mat_smoother, self.MS[0].levels[0].params.nsweeps) 

196 

197 # add coarse-grid correction (single sweep, though!) 

198 if self.nlevels > 1: 

199 precond_cgc = self.Tcf.dot(np.linalg.inv(self.Pc)).dot(self.Tfc) 

200 iter_mat_cgc = np.eye(self.nsteps * self.nnodes * self.nspace) - precond_cgc.dot(self.C) 

201 iter_mat = iter_mat_smoother.dot(iter_mat_cgc) 

202 precond = precond_smoother + precond_cgc - precond_smoother.dot(self.C).dot(precond_cgc) 

203 else: 

204 iter_mat = iter_mat_smoother 

205 precond = precond_smoother 

206 

207 # form span and reduce matrices and add to operator 

208 Tspread = np.kron(np.concatenate([[1] * (self.nsteps * self.nnodes)]), np.eye(self.nspace)).T 

209 Tnospread = np.kron( 

210 np.concatenate([[1], [0] * (self.nsteps - 1)]), np.kron(np.ones(self.nnodes), np.eye(self.nspace)) 

211 ).T 

212 Treduce = np.kron(np.concatenate([[0] * (self.nsteps * self.nnodes - 1), [1]]), np.eye(self.nspace)) 

213 

214 if self.MS[0].levels[0].sweep.params.initial_guess == 'spread': 

215 mat = np.linalg.matrix_power(iter_mat, niter).dot(Tspread) 

216 # mat = iter_mat_smoother.dot(Tspread) + precond_smoother.dot(Tnospread) 

217 else: 

218 mat = np.linalg.matrix_power(iter_mat, niter).dot(Tnospread) 

219 # mat = iter_mat_smoother.dot(Tnospread) + precond_smoother.dot(Tnospread) # No, the latter is not a typo! 

220 

221 # build propagation matrix 

222 # mat = np.linalg.matrix_power(iter_mat, niter - 1).dot(mat) 

223 for k in range(niter): 

224 mat += np.linalg.matrix_power(iter_mat, k).dot(precond).dot(Tnospread) 

225 mat = Treduce.dot(mat) 

226 

227 return mat 

228 

229 def restart_block(self, slots, time, u0): 

230 """ 

231 Helper routine to reset/restart block of steps 

232 

233 Args: 

234 slots: list of steps 

235 time: list of new times 

236 u0: initial value to distribute across the steps 

237 

238 """ 

239 

240 # loop over steps 

241 for p in slots: 

242 # store current slot number for diagnostics 

243 self.MS[p].status.slot = p 

244 

245 for p in slots: 

246 for lvl in self.MS[p].levels: 

247 lvl.status.time = time[p] 

248 P = lvl.prob 

249 for m in range(1, lvl.sweep.coll.num_nodes + 1): 

250 lvl.u[m] = P.dtype_u(init=P.init, val=0.0) 

251 lvl.f[m] = P.dtype_f(init=P.init, val=0.0) 

252 

253 self.u0 = np.kron(np.concatenate([[1], [0] * (self.nsteps - 1)]), np.kron(np.ones(self.nnodes), u0)) 

254 

255 if self.MS[0].levels[0].sweep.params.initial_guess == 'spread': 

256 self.u = np.kron(np.ones(self.nsteps * self.nnodes), u0) 

257 else: 

258 self.u = self.u0.copy() 

259 

260 self.res = np.zeros(self.nsteps * self.nnodes * self.nspace) 

261 

262 @staticmethod 

263 def update_data(MS, u, res, niter, level, stage): 

264 for S in MS: 

265 S.status.stage = stage 

266 S.status.iter = niter 

267 

268 L = S.levels[level] 

269 P = L.prob 

270 nnodes = L.sweep.coll.num_nodes 

271 nspace = P.init[0] 

272 

273 first = S.status.slot * nnodes * nspace 

274 last = (S.status.slot + 1) * nnodes * nspace 

275 

276 L.status.residual = np.linalg.norm(res[first:last], np.inf) 

277 

278 for m in range(1, nnodes + 1): 

279 mstart = first + (m - 1) * nspace 

280 mend = first + m * nspace 

281 L.u[m][:] = u[mstart:mend] 

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

283 

284 S.levels[level].sweep.compute_end_point() 

285 

286 return MS 

287 

288 def pfasst(self, MS): 

289 """ 

290 Main function including the stages of SDC, MLSDC and PFASST (the "controller") 

291 

292 Args: 

293 MS: all active steps 

294 

295 Returns: 

296 all active steps 

297 """ 

298 

299 niter = 0 

300 

301 self.res = self.u0 - self.C.dot(self.u) 

302 

303 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='PRE_STEP') 

304 for S in MS: 

305 for hook in self.hooks: 

306 hook.pre_step(step=S, level_number=0) 

307 

308 while np.linalg.norm(self.res, np.inf) > self.tol and niter < self.maxiter: 

309 niter += 1 

310 

311 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='PRE_ITERATION') 

312 for S in MS: 

313 for hook in self.hooks: 

314 hook.pre_iteration(step=S, level_number=0) 

315 

316 if self.nlevels > 1: 

317 for _ in range(MS[0].levels[1].params.nsweeps): 

318 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=1, stage='PRE_COARSE_SWEEP') 

319 for S in MS: 

320 for hook in self.hooks: 

321 hook.pre_sweep(step=S, level_number=1) 

322 

323 self.u += self.Tcf.dot(np.linalg.solve(self.Pc, self.Tfc.dot(self.res))) 

324 self.res = self.u0 - self.C.dot(self.u) 

325 

326 MS = self.update_data( 

327 MS=MS, u=self.u, res=self.res, niter=niter, level=1, stage='POST_COARSE_SWEEP' 

328 ) 

329 for S in MS: 

330 for hook in self.hooks: 

331 hook.post_sweep(step=S, level_number=1) 

332 

333 for _ in range(MS[0].levels[0].params.nsweeps): 

334 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='PRE_FINE_SWEEP') 

335 for S in MS: 

336 for hook in self.hooks: 

337 hook.pre_sweep(step=S, level_number=0) 

338 

339 self.u += np.linalg.solve(self.P, self.res) 

340 self.res = self.u0 - self.C.dot(self.u) 

341 

342 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='POST_FINE_SWEEP') 

343 for S in MS: 

344 for hook in self.hooks: 

345 hook.post_sweep(step=S, level_number=0) 

346 

347 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='POST_ITERATION') 

348 for S in MS: 

349 for hook in self.hooks: 

350 hook.post_iteration(step=S, level_number=0) 

351 

352 MS = self.update_data(MS=MS, u=self.u, res=self.res, niter=niter, level=0, stage='POST_STEP') 

353 for S in MS: 

354 for hook in self.hooks: 

355 hook.post_step(step=S, level_number=0) 

356 

357 return MS