Coverage for pySDC/implementations/convergence_controller_classes/adaptive_collocation.py: 99%
79 statements
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-16 14:51 +0000
1import numpy as np
3from qmat.lagrange import LagrangeApproximation
4from pySDC.core.convergence_controller 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
20 .. code-block:: python
22 params = {
23 'num_nodes': [2, 3],
24 }
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.
32 Feel free to set `logger_level = 15` in the controller parameters to get comprehensive text output on what exactly
33 is happening.
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.
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 """
46 def setup(self, controller, params, description, **kwargs):
47 """
48 Record what variables we want to vary.
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
55 Returns:
56 (dict): The updated params dictionary
57 """
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 }
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']
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!')
81 defaults['num_colls'] = max([defaults['num_colls'], len(params[key])])
83 self.comm = description['sweeper_params'].get('comm', None)
84 if self.comm:
85 from mpi4py import MPI
87 self.MPI_SUM = MPI.SUM
89 return {**defaults, **super().setup(controller, params, description, **kwargs)}
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.
98 Args:
99 A (2d np.ndarray): Matrix
100 b (list): Vector
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
116 def switch_sweeper(self, S):
117 """
118 Update to the next sweeper in line.
120 Args:
121 S (pySDC.Step.step): The current step
123 Returns:
124 None
125 """
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}
134 new_params_level = {key: self.params.get(key)[self.status.active_coll] for key in self.params.vary_keys_level}
136 # update sweeper for all levels
137 for L in S.levels:
138 P = L.prob
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()
144 # change sweeper
145 L.sweep.__init__(update_params_sweeper)
146 L.sweep.level = L
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)
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))
156 u_inter = self.matmul(interpolator.getInterpolationMatrix(np.append(0, nodes_new)), u_old)
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
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)
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)
182 def setup_status_variables(self, controller, **kwargs):
183 """
184 Add an index for which collocation method to use.
186 Args:
187 controller (pySDC.Controller.controller): The controller
189 Returns:
190 None
191 """
192 self.status = Status(['active_coll'])
194 def reset_status_variables(self, controller, **kwargs):
195 """
196 Reset the status variables between time steps.
198 Args:
199 controller (pySDC.Controller.controller): The controller
201 Returns:
202 None
203 """
204 self.status.active_coll = 0
206 def post_iteration_processing(self, controller, S, **kwargs):
207 """
208 Switch to the next collocation method if the current one is converged.
210 Args:
211 controller (pySDC.Controller.controller): The controller
212 S (pySDC.Step.step): The current step
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)
222 def post_spread_processing(self, controller, S, **kwargs):
223 """
224 Overwrite the sweeper parameters with the first collocation parameters set up here.
226 Args:
227 controller (pySDC.Controller.controller): The controller
228 S (pySDC.Step.step): The current step
230 Returns:
231 None
232 """
233 self.switch_sweeper(S)
235 def check_parameters(self, controller, params, description, **kwargs):
236 """
237 Check if we allow the scheme to solve the collocation problems to convergence.
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
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 )
254 return True, ""