Coverage for pySDC/implementations/convergence_controller_classes/adaptive_collocation.py: 99%
79 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-29 09:02 +0000
1import numpy as np
3from pySDC.core.Lagrange import LagrangeApproximation
4from pySDC.core.ConvergenceController import ConvergenceController, Status
5from pySDC.core.Collocation import CollBase
8class AdaptiveCollocation(ConvergenceController):
9 """
10 This convergence controller allows to change the underlying quadrature between iterations.
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.
17 Use this convergence controller by supplying parameters that the sweeper accepts as a list to the `params`.
18 For instance, supplying
19 ```
20 params = {
21 'num_nodes': [2, 3],
22 }
23 ```
24 will use collocation methods like you passed to the `sweeper_params` in the `description` object, but will change
25 the number of nodes to 2 before the first iteration and to 3 as soon as the 2-node collocation problem is converged.
26 This will override whatever you set for the number of nodes in the `sweeper_params`, but you still need to set
27 something to allow instantiation of the levels before this convergence controller becomes active.
28 Make sure all lists you supply here have the same length.
30 Feel free to set `logger_level = 15` in the controller parameters to get comprehensive text output on what exactly
31 is happening.
33 This convergence controller has various applications.
34 - You could try to obtain speedup by doing some inexactness. It is currently possible to set various residual
35 tolerances, which will be passed to the levels, corresponding to the accuracy with which each collocation problem
36 is solved.
37 - You can compute multiple solutions to the same initial value problem with different order. This allows, for
38 instance, to do adaptive time stepping.
40 When trying to obtain speedup with this, be ware that the interpolation is not for free. In particular, it is
41 necessary to reevaluate the right hand side at all nodes afterwards.
42 """
44 def setup(self, controller, params, description, **kwargs):
45 """
46 Record what variables we want to vary.
48 Args:
49 controller (pySDC.Controller.controller): The controller
50 params (dict): The params passed for this specific convergence controller
51 description (dict): The description object used to instantiate the controller
53 Returns:
54 (dict): The updated params dictionary
55 """
57 defaults = {
58 'control_order': 300,
59 'num_colls': 0,
60 'sweeper_params': description['sweeper_params'],
61 'vary_keys_sweeper': [],
62 'vary_keys_level': [],
63 }
65 # only these keys can be changed by this convergence controller
66 self.allowed_sweeper_keys = ['quad_type', 'num_nodes', 'node_type', 'do_coll_update']
67 self.allowed_level_keys = ['restol']
69 # add the keys to lists so we know what we need to change later
70 for key in params.keys():
71 if type(params[key]) == list:
72 if key in self.allowed_sweeper_keys:
73 defaults['vary_keys_sweeper'] += [key]
74 elif key in self.allowed_level_keys:
75 defaults['vary_keys_level'] += [key]
76 else:
77 raise NotImplementedError(f'Don\'t know what to do with key {key} here!')
79 defaults['num_colls'] = max([defaults['num_colls'], len(params[key])])
81 self.comm = description['sweeper_params'].get('comm', None)
82 if self.comm:
83 from mpi4py import MPI
85 self.MPI_SUM = MPI.SUM
87 return {**defaults, **super().setup(controller, params, description, **kwargs)}
89 def matmul(self, A, b):
90 """
91 Matrix vector multiplication, possibly MPI parallel.
92 The parallel implementation performs a reduce operation in every row of the matrix. While communicating the
93 entire vector once could reduce the number of communications, this way we never need to store the entire vector
94 on any specific rank.
96 Args:
97 A (2d np.ndarray): Matrix
98 b (list): Vector
100 Returns:
101 List: Axb
102 """
103 if self.comm:
104 res = [A[i, 0] * b[0] if b[i] is not None else None for i in range(A.shape[0])]
105 buf = b[0] * 0.0
106 for i in range(1, A.shape[1]):
107 self.comm.Reduce(A[i, self.comm.rank + 1] * b[self.comm.rank + 1], buf, op=self.MPI_SUM, root=i - 1)
108 if i == self.comm.rank + 1:
109 res[i] += buf
110 return res
111 else:
112 return A @ b
114 def switch_sweeper(self, S):
115 """
116 Update to the next sweeper in line.
118 Args:
119 S (pySDC.Step.step): The current step
121 Returns:
122 None
123 """
125 # generate dictionaries with the new parameters
126 new_params_sweeper = {
127 key: self.params.get(key)[self.status.active_coll] for key in self.params.vary_keys_sweeper
128 }
129 sweeper_params = self.params.sweeper_params.copy()
130 update_params_sweeper = {**sweeper_params, **new_params_sweeper}
132 new_params_level = {key: self.params.get(key)[self.status.active_coll] for key in self.params.vary_keys_level}
134 # update sweeper for all levels
135 for L in S.levels:
136 P = L.prob
138 # store solution of current level which will be interpolated to new level
139 u_old = [me.flatten() if me is not None else me for me in L.u]
140 nodes_old = L.sweep.coll.nodes.copy()
142 # change sweeper
143 L.sweep.__init__(update_params_sweeper)
144 L.sweep.level = L
146 # reset level to tell it the new structure of the solution
147 L.params.__dict__.update(new_params_level)
148 L.reset_level(reset_status=False)
150 # interpolate solution of old collocation problem to new one
151 nodes_new = L.sweep.coll.nodes.copy()
152 interpolator = LagrangeApproximation(points=np.append(0, nodes_old))
154 u_inter = self.matmul(interpolator.getInterpolationMatrix(np.append(0, nodes_new)), u_old)
156 # assign the interpolated values to the nodes in the level
157 for i in range(0, len(u_inter)):
158 if u_inter[i] is not None:
159 me = P.dtype_u(P.init)
160 me[:] = np.reshape(u_inter[i], P.init[0])
161 L.u[i] = me
163 # reevaluate rhs
164 for i in range(L.sweep.coll.num_nodes + 1):
165 if L.u[i] is not None:
166 L.f[i] = L.prob.eval_f(L.u[i], L.time)
168 # log the new parameters
169 self.log(f'Switching to collocation {self.status.active_coll + 1} of {self.params.num_colls}', S, level=20)
170 msg = 'New quadrature:'
171 for key in list(sweeper_params.keys()) + list(new_params_level.keys()):
172 if key in self.params.vary_keys_sweeper:
173 msg += f'\n--> {key}: {update_params_sweeper[key]}'
174 elif key in self.params.vary_keys_level:
175 msg += f'\n--> {key}: {new_params_level[key]}'
176 else:
177 msg += f'\n {key}: {update_params_sweeper[key]}'
178 self.log(msg, S)
180 def setup_status_variables(self, controller, **kwargs):
181 """
182 Add an index for which collocation method to use.
184 Args:
185 controller (pySDC.Controller.controller): The controller
187 Returns:
188 None
189 """
190 self.status = Status(['active_coll'])
192 def reset_status_variables(self, controller, **kwargs):
193 """
194 Reset the status variables between time steps.
196 Args:
197 controller (pySDC.Controller.controller): The controller
199 Returns:
200 None
201 """
202 self.status.active_coll = 0
204 def post_iteration_processing(self, controller, S, **kwargs):
205 """
206 Switch to the next collocation method if the current one is converged.
208 Args:
209 controller (pySDC.Controller.controller): The controller
210 S (pySDC.Step.step): The current step
212 Returns:
213 None
214 """
215 if (self.status.active_coll < self.params.num_colls - 1) and S.status.done:
216 self.status.active_coll += 1
217 S.status.done = False
218 self.switch_sweeper(S)
220 def post_spread_processing(self, controller, S, **kwargs):
221 """
222 Overwrite the sweeper parameters with the first collocation parameters set up here.
224 Args:
225 controller (pySDC.Controller.controller): The controller
226 S (pySDC.Step.step): The current step
228 Returns:
229 None
230 """
231 self.switch_sweeper(S)
233 def check_parameters(self, controller, params, description, **kwargs):
234 """
235 Check if we allow the scheme to solve the collocation problems to convergence.
237 Args:
238 controller (pySDC.Controller): The controller
239 params (dict): The params passed for this specific convergence controller
240 description (dict): The description object used to instantiate the controller
242 Returns:
243 bool: Whether the parameters are compatible
244 str: The error message
245 """
246 if description["level_params"].get("restol", -1.0) <= 1e-16:
247 return (
248 False,
249 "Switching the collocation problems requires solving them to some tolerance that can be reached. Please set attainable `restol` in the level params",
250 )
252 return True, ""