Coverage for pySDC/implementations/convergence_controller_classes/adaptive_collocation.py: 99%

79 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-20 14:51 +0000

1import numpy as np 

2 

3from qmat.lagrange import LagrangeApproximation 

4from pySDC.core.convergence_controller import ConvergenceController, Status 

5from pySDC.core.collocation import CollBase 

6 

7 

8class AdaptiveCollocation(ConvergenceController): 

9 """ 

10 This convergence controller allows to change the underlying quadrature between iterations. 

11 

12 Supplying multiple quadrature rules will result in a change to a new quadrature whenever the previous one is 

13 converged until all methods given are converged, at which point the step is ended. 

14 Whenever the quadrature is changed, the solution is interpolated from the old nodes to the new nodes to ensure 

15 accelerated convergence compared to starting from the initial conditions. 

16 

17 Use this convergence controller by supplying parameters that the sweeper accepts as a list to the `params`. 

18 For instance, supplying 

19 

20 .. code-block:: python 

21 

22 params = { 

23 'num_nodes': [2, 3], 

24 } 

25 

26 will use collocation methods like you passed to the `sweeper_params` in the `description` object, but will change 

27 the number of nodes to 2 before the first iteration and to 3 as soon as the 2-node collocation problem is converged. 

28 This will override whatever you set for the number of nodes in the `sweeper_params`, but you still need to set 

29 something to allow instantiation of the levels before this convergence controller becomes active. 

30 Make sure all lists you supply here have the same length. 

31 

32 Feel free to set `logger_level = 15` in the controller parameters to get comprehensive text output on what exactly 

33 is happening. 

34 

35 This convergence controller has various applications. 

36 - You could try to obtain speedup by doing some inexactness. It is currently possible to set various residual 

37 tolerances, which will be passed to the levels, corresponding to the accuracy with which each collocation problem 

38 is solved. 

39 - You can compute multiple solutions to the same initial value problem with different order. This allows, for 

40 instance, to do adaptive time stepping. 

41 

42 When trying to obtain speedup with this, be ware that the interpolation is not for free. In particular, it is 

43 necessary to reevaluate the right hand side at all nodes afterwards. 

44 """ 

45 

46 def setup(self, controller, params, description, **kwargs): 

47 """ 

48 Record what variables we want to vary. 

49 

50 Args: 

51 controller (pySDC.Controller.controller): The controller 

52 params (dict): The params passed for this specific convergence controller 

53 description (dict): The description object used to instantiate the controller 

54 

55 Returns: 

56 (dict): The updated params dictionary 

57 """ 

58 

59 defaults = { 

60 'control_order': 300, 

61 'num_colls': 0, 

62 'sweeper_params': description['sweeper_params'], 

63 'vary_keys_sweeper': [], 

64 'vary_keys_level': [], 

65 } 

66 

67 # only these keys can be changed by this convergence controller 

68 self.allowed_sweeper_keys = ['quad_type', 'num_nodes', 'node_type', 'do_coll_update'] 

69 self.allowed_level_keys = ['restol'] 

70 

71 # add the keys to lists so we know what we need to change later 

72 for key in params.keys(): 

73 if type(params[key]) == list: 

74 if key in self.allowed_sweeper_keys: 

75 defaults['vary_keys_sweeper'] += [key] 

76 elif key in self.allowed_level_keys: 

77 defaults['vary_keys_level'] += [key] 

78 else: 

79 raise NotImplementedError(f'Don\'t know what to do with key {key} here!') 

80 

81 defaults['num_colls'] = max([defaults['num_colls'], len(params[key])]) 

82 

83 self.comm = description['sweeper_params'].get('comm', None) 

84 if self.comm: 

85 from mpi4py import MPI 

86 

87 self.MPI_SUM = MPI.SUM 

88 

89 return {**defaults, **super().setup(controller, params, description, **kwargs)} 

90 

91 def matmul(self, A, b): 

92 """ 

93 Matrix vector multiplication, possibly MPI parallel. 

94 The parallel implementation performs a reduce operation in every row of the matrix. While communicating the 

95 entire vector once could reduce the number of communications, this way we never need to store the entire vector 

96 on any specific rank. 

97 

98 Args: 

99 A (2d np.ndarray): Matrix 

100 b (list): Vector 

101 

102 Returns: 

103 List: Axb 

104 """ 

105 if self.comm: 

106 res = [A[i, 0] * b[0] if b[i] is not None else None for i in range(A.shape[0])] 

107 buf = b[0] * 0.0 

108 for i in range(1, A.shape[1]): 

109 self.comm.Reduce(A[i, self.comm.rank + 1] * b[self.comm.rank + 1], buf, op=self.MPI_SUM, root=i - 1) 

110 if i == self.comm.rank + 1: 

111 res[i] += buf 

112 return res 

113 else: 

114 return A @ b 

115 

116 def switch_sweeper(self, S): 

117 """ 

118 Update to the next sweeper in line. 

119 

120 Args: 

121 S (pySDC.Step.step): The current step 

122 

123 Returns: 

124 None 

125 """ 

126 

127 # generate dictionaries with the new parameters 

128 new_params_sweeper = { 

129 key: self.params.get(key)[self.status.active_coll] for key in self.params.vary_keys_sweeper 

130 } 

131 sweeper_params = self.params.sweeper_params.copy() 

132 update_params_sweeper = {**sweeper_params, **new_params_sweeper} 

133 

134 new_params_level = {key: self.params.get(key)[self.status.active_coll] for key in self.params.vary_keys_level} 

135 

136 # update sweeper for all levels 

137 for L in S.levels: 

138 P = L.prob 

139 

140 # store solution of current level which will be interpolated to new level 

141 u_old = [me.flatten() if me is not None else me for me in L.u] 

142 nodes_old = L.sweep.coll.nodes.copy() 

143 

144 # change sweeper 

145 L.sweep.__init__(update_params_sweeper) 

146 L.sweep.level = L 

147 

148 # reset level to tell it the new structure of the solution 

149 L.params.__dict__.update(new_params_level) 

150 L.reset_level(reset_status=False) 

151 

152 # interpolate solution of old collocation problem to new one 

153 nodes_new = L.sweep.coll.nodes.copy() 

154 interpolator = LagrangeApproximation(points=np.append(0, nodes_old)) 

155 

156 u_inter = self.matmul(interpolator.getInterpolationMatrix(np.append(0, nodes_new)), u_old) 

157 

158 # assign the interpolated values to the nodes in the level 

159 for i in range(0, len(u_inter)): 

160 if u_inter[i] is not None: 

161 me = P.dtype_u(P.init) 

162 me[:] = np.reshape(u_inter[i], P.init[0]) 

163 L.u[i] = me 

164 

165 # reevaluate rhs 

166 for i in range(L.sweep.coll.num_nodes + 1): 

167 if L.u[i] is not None: 

168 L.f[i] = L.prob.eval_f(L.u[i], L.time) 

169 

170 # log the new parameters 

171 self.log(f'Switching to collocation {self.status.active_coll + 1} of {self.params.num_colls}', S, level=20) 

172 msg = 'New quadrature:' 

173 for key in list(sweeper_params.keys()) + list(new_params_level.keys()): 

174 if key in self.params.vary_keys_sweeper: 

175 msg += f'\n--> {key}: {update_params_sweeper[key]}' 

176 elif key in self.params.vary_keys_level: 

177 msg += f'\n--> {key}: {new_params_level[key]}' 

178 else: 

179 msg += f'\n {key}: {update_params_sweeper[key]}' 

180 self.log(msg, S) 

181 

182 def setup_status_variables(self, controller, **kwargs): 

183 """ 

184 Add an index for which collocation method to use. 

185 

186 Args: 

187 controller (pySDC.Controller.controller): The controller 

188 

189 Returns: 

190 None 

191 """ 

192 self.status = Status(['active_coll']) 

193 

194 def reset_status_variables(self, controller, **kwargs): 

195 """ 

196 Reset the status variables between time steps. 

197 

198 Args: 

199 controller (pySDC.Controller.controller): The controller 

200 

201 Returns: 

202 None 

203 """ 

204 self.status.active_coll = 0 

205 

206 def post_iteration_processing(self, controller, S, **kwargs): 

207 """ 

208 Switch to the next collocation method if the current one is converged. 

209 

210 Args: 

211 controller (pySDC.Controller.controller): The controller 

212 S (pySDC.Step.step): The current step 

213 

214 Returns: 

215 None 

216 """ 

217 if (self.status.active_coll < self.params.num_colls - 1) and S.status.done: 

218 self.status.active_coll += 1 

219 S.status.done = False 

220 self.switch_sweeper(S) 

221 

222 def post_spread_processing(self, controller, S, **kwargs): 

223 """ 

224 Overwrite the sweeper parameters with the first collocation parameters set up here. 

225 

226 Args: 

227 controller (pySDC.Controller.controller): The controller 

228 S (pySDC.Step.step): The current step 

229 

230 Returns: 

231 None 

232 """ 

233 self.switch_sweeper(S) 

234 

235 def check_parameters(self, controller, params, description, **kwargs): 

236 """ 

237 Check if we allow the scheme to solve the collocation problems to convergence. 

238 

239 Args: 

240 controller (pySDC.Controller): The controller 

241 params (dict): The params passed for this specific convergence controller 

242 description (dict): The description object used to instantiate the controller 

243 

244 Returns: 

245 bool: Whether the parameters are compatible 

246 str: The error message 

247 """ 

248 if description["level_params"].get("restol", -1.0) <= 1e-16: 

249 return ( 

250 False, 

251 "Switching the collocation problems requires solving them to some tolerance that can be reached. Please set attainable `restol` in the level params", 

252 ) 

253 

254 return True, ""