Coverage for pySDC/implementations/sweeper_classes/Runge_Kutta.py: 93%
281 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-09-09 14:59 +0000
« prev ^ index » next coverage.py v7.6.1, created at 2024-09-09 14:59 +0000
1import numpy as np
2import logging
3from qmat.qcoeff.butcher import RK_SCHEMES
5from pySDC.core.sweeper import Sweeper, _Pars
6from pySDC.core.errors import ParameterError
7from pySDC.core.level import Level
10class ButcherTableau(object):
11 def __init__(self, weights, nodes, matrix):
12 """
13 Initialization routine to get a quadrature matrix out of a Butcher tableau
15 Args:
16 weights (numpy.ndarray): Butcher tableau weights
17 nodes (numpy.ndarray): Butcher tableau nodes
18 matrix (numpy.ndarray): Butcher tableau entries
19 """
20 # check if the arguments have the correct form
21 if type(matrix) != np.ndarray:
22 raise ParameterError('Runge-Kutta matrix needs to be supplied as a numpy array!')
23 elif len(np.unique(matrix.shape)) != 1 or len(matrix.shape) != 2:
24 raise ParameterError('Runge-Kutta matrix needs to be a square 2D numpy array!')
26 if type(weights) != np.ndarray:
27 raise ParameterError('Weights need to be supplied as a numpy array!')
28 elif len(weights.shape) != 1:
29 raise ParameterError(f'Incompatible dimension of weights! Need 1, got {len(weights.shape)}')
30 elif len(weights) != matrix.shape[0]:
31 raise ParameterError(f'Incompatible number of weights! Need {matrix.shape[0]}, got {len(weights)}')
33 if type(nodes) != np.ndarray:
34 raise ParameterError('Nodes need to be supplied as a numpy array!')
35 elif len(nodes.shape) != 1:
36 raise ParameterError(f'Incompatible dimension of nodes! Need 1, got {len(nodes.shape)}')
37 elif len(nodes) != matrix.shape[0]:
38 raise ParameterError(f'Incompatible number of nodes! Need {matrix.shape[0]}, got {len(nodes)}')
40 # Set number of nodes, left and right interval boundaries
41 self.num_solution_stages = 1
42 self.num_nodes = matrix.shape[0] + self.num_solution_stages
43 self.tleft = 0.0
44 self.tright = 1.0
46 self.nodes = np.append(np.append([0], nodes), [1])
47 self.weights = weights
48 self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
49 self.Qmat[1:-1, 1:-1] = matrix
50 self.Qmat[-1, 1:-1] = weights # this is for computing the solution to the step from the previous stages
52 self.left_is_node = True
53 self.right_is_node = self.nodes[-1] == self.tright
55 # compute distances between the nodes
56 if self.num_nodes > 1:
57 self.delta_m = self.nodes[1:] - self.nodes[:-1]
58 else:
59 self.delta_m = np.zeros(1)
60 self.delta_m[0] = self.nodes[0] - self.tleft
62 # check if the RK scheme is implicit
63 self.implicit = any(matrix[i, i] != 0 for i in range(self.num_nodes - self.num_solution_stages))
66class ButcherTableauEmbedded(object):
67 def __init__(self, weights, nodes, matrix):
68 """
69 Initialization routine to get a quadrature matrix out of a Butcher tableau for embedded RK methods.
71 Be aware that the method that generates the final solution should be in the first row of the weights matrix.
73 Args:
74 weights (numpy.ndarray): Butcher tableau weights
75 nodes (numpy.ndarray): Butcher tableau nodes
76 matrix (numpy.ndarray): Butcher tableau entries
77 """
78 # check if the arguments have the correct form
79 if type(matrix) != np.ndarray:
80 raise ParameterError('Runge-Kutta matrix needs to be supplied as a numpy array!')
81 elif len(np.unique(matrix.shape)) != 1 or len(matrix.shape) != 2:
82 raise ParameterError('Runge-Kutta matrix needs to be a square 2D numpy array!')
84 if type(weights) != np.ndarray:
85 raise ParameterError('Weights need to be supplied as a numpy array!')
86 elif len(weights.shape) != 2:
87 raise ParameterError(f'Incompatible dimension of weights! Need 2, got {len(weights.shape)}')
88 elif len(weights[0]) != matrix.shape[0]:
89 raise ParameterError(f'Incompatible number of weights! Need {matrix.shape[0]}, got {len(weights[0])}')
91 if type(nodes) != np.ndarray:
92 raise ParameterError('Nodes need to be supplied as a numpy array!')
93 elif len(nodes.shape) != 1:
94 raise ParameterError(f'Incompatible dimension of nodes! Need 1, got {len(nodes.shape)}')
95 elif len(nodes) != matrix.shape[0]:
96 raise ParameterError(f'Incompatible number of nodes! Need {matrix.shape[0]}, got {len(nodes)}')
98 # Set number of nodes, left and right interval boundaries
99 self.num_solution_stages = 2
100 self.num_nodes = matrix.shape[0] + self.num_solution_stages
101 self.tleft = 0.0
102 self.tright = 1.0
104 self.nodes = np.append(np.append([0], nodes), [1, 1])
105 self.weights = weights
106 self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
107 self.Qmat[1:-2, 1:-2] = matrix
108 self.Qmat[-1, 1:-2] = weights[0] # this is for computing the higher order solution
109 self.Qmat[-2, 1:-2] = weights[1] # this is for computing the lower order solution
111 self.left_is_node = True
112 self.right_is_node = self.nodes[-1] == self.tright
114 # compute distances between the nodes
115 if self.num_nodes > 1:
116 self.delta_m = self.nodes[1:] - self.nodes[:-1]
117 else:
118 self.delta_m = np.zeros(1)
119 self.delta_m[0] = self.nodes[0] - self.tleft
121 # check if the RK scheme is implicit
122 self.implicit = any(matrix[i, i] != 0 for i in range(self.num_nodes - self.num_solution_stages))
125class RungeKutta(Sweeper):
126 nodes = None
127 weights = None
128 matrix = None
129 ButcherTableauClass = ButcherTableau
131 """
132 Runge-Kutta scheme that fits the interface of a sweeper.
133 Actually, the sweeper idea fits the Runge-Kutta idea when using only lower triangular rules, where solutions
134 at the nodes are successively computed from earlier nodes. However, we only perform a single iteration of this.
136 We have two choices to realise a Runge-Kutta sweeper: We can choose Q = Q_Delta = <Butcher tableau>, but in this
137 implementation, that would lead to a lot of wasted FLOPS from integrating with Q and then with Q_Delta and
138 subtracting the two. For that reason, we built this new sweeper, which does not have a preconditioner.
140 This class only supports lower triangular Butcher tableaux such that the system can be solved with forward
141 substitution. In this way, we don't get the maximum order that we could for the number of stages, but computing the
142 stages is much cheaper. In particular, if the Butcher tableaux is strictly lower triangular, we get an explicit
143 method, which does not require us to solve a system of equations to compute the stages.
145 Please be aware that all fundamental parameters of the Sweeper are ignored. These include
147 - num_nodes
148 - collocation_class
149 - initial_guess
150 - QI
152 All of these variables are either determined by the RK rule, or are not part of an RK scheme.
154 The entries of the Butcher tableau are stored as class attributes.
155 """
157 def __init__(self, params):
158 """
159 Initialization routine for the custom sweeper
161 Args:
162 params: parameters for the sweeper
163 """
164 # set up logger
165 self.logger = logging.getLogger('sweeper')
167 # check if some parameters are set which only apply to actual sweepers
168 for key in ['initial_guess', 'collocation_class', 'num_nodes']:
169 if key in params:
170 self.logger.warning(f'"{key}" will be ignored by Runge-Kutta sweeper')
172 # set parameters to their actual values
173 self.coll = self.get_Butcher_tableau()
174 params['initial_guess'] = 'zero'
175 params['collocation_class'] = type(self.ButcherTableauClass)
176 params['num_nodes'] = self.coll.num_nodes
178 # disable residual computation by default
179 params['skip_residual_computation'] = params.get(
180 'skip_residual_computation', ('IT_CHECK', 'IT_FINE', 'IT_COARSE', 'IT_UP', 'IT_DOWN')
181 )
183 # check if we can skip some usually unnecessary right hand side evaluations
184 params['eval_rhs_at_right_boundary'] = params.get('eval_rhs_at_right_boundary', False)
186 self.params = _Pars(params)
188 # This will be set as soon as the sweeper is instantiated at the level
189 self.__level = None
191 self.parallelizable = False
192 self.QI = self.coll.Qmat
194 @classmethod
195 def get_Q_matrix(cls):
196 return cls.get_Butcher_tableau().Qmat
198 @classmethod
199 def get_Butcher_tableau(cls):
200 return cls.ButcherTableauClass(cls.weights, cls.nodes, cls.matrix)
202 @classmethod
203 def get_update_order(cls):
204 """
205 Get the order of the lower order method for doing adaptivity. Only applies to embedded methods.
206 """
207 raise NotImplementedError(
208 f"There is not an update order for RK scheme \"{cls.__name__}\" implemented. Maybe it is not an embedded scheme?"
209 )
211 def get_full_f(self, f):
212 """
213 Get the full right hand side as a `mesh` from the right hand side
215 Args:
216 f (dtype_f): Right hand side at a single node
218 Returns:
219 mesh: Full right hand side as a mesh
220 """
221 if type(f).__name__ in ['mesh', 'cupy_mesh']:
222 return f
223 elif type(f).__name__ in ['imex_mesh', 'imex_cupy_mesh']:
224 return f.impl + f.expl
225 elif f is None:
226 prob = self.level.prob
227 return self.get_full_f(prob.dtype_f(prob.init, val=0))
228 else:
229 raise NotImplementedError(f'Type \"{type(f)}\" not implemented in Runge-Kutta sweeper')
231 def integrate(self):
232 """
233 Integrates the right-hand side
235 Returns:
236 list of dtype_u: containing the integral as values
237 """
239 # get current level and problem
240 lvl = self.level
241 prob = lvl.prob
243 me = []
245 # integrate RHS over all collocation nodes
246 for m in range(1, self.coll.num_nodes + 1):
247 # new instance of dtype_u, initialize values with 0
248 me.append(prob.dtype_u(prob.init, val=0.0))
249 for j in range(1, self.coll.num_nodes + 1):
250 me[-1] += lvl.dt * self.coll.Qmat[m, j] * self.get_full_f(lvl.f[j])
252 return me
254 def update_nodes(self):
255 """
256 Update the u- and f-values at the collocation nodes
258 Returns:
259 None
260 """
262 # get current level and problem
263 lvl = self.level
264 prob = lvl.prob
266 # only if the level has been touched before
267 assert lvl.status.unlocked
268 assert lvl.status.sweep <= 1, "RK schemes are direct solvers. Please perform only 1 iteration!"
270 # get number of collocation nodes for easier access
271 M = self.coll.num_nodes
273 for m in range(0, M):
274 # build rhs, consisting of the known values from above and new values from previous nodes (at k+1)
275 rhs = prob.dtype_u(lvl.u[0])
276 for j in range(1, m + 1):
277 rhs += lvl.dt * self.QI[m + 1, j] * self.get_full_f(lvl.f[j])
279 # implicit solve with prefactor stemming from the diagonal of Qd, use previous stage as initial guess
280 if self.coll.implicit:
281 lvl.u[m + 1][:] = prob.solve_system(
282 rhs, lvl.dt * self.QI[m + 1, m + 1], lvl.u[m], lvl.time + lvl.dt * self.coll.nodes[m + 1]
283 )
284 else:
285 lvl.u[m + 1][:] = rhs[:]
287 # update function values (we don't usually need to evaluate the RHS at the solution of the step)
288 if m < M - self.coll.num_solution_stages or self.params.eval_rhs_at_right_boundary:
289 lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
291 # indicate presence of new values at this level
292 lvl.status.updated = True
294 return None
296 def compute_end_point(self):
297 """
298 In this Runge-Kutta implementation, the solution to the step is always stored in the last node
299 """
300 self.level.uend = self.level.u[-1]
302 @property
303 def level(self):
304 """
305 Returns the current level
307 Returns:
308 pySDC.Level.level: Current level
309 """
310 return self.__level
312 @level.setter
313 def level(self, lvl):
314 """
315 Sets a reference to the current level (done in the initialization of the level)
317 Args:
318 lvl (pySDC.Level.level): Current level
319 """
320 assert isinstance(lvl, Level), f"You tried to set the sweeper's level with an instance of {type(lvl)}!"
321 if lvl.params.restol > 0:
322 lvl.params.restol = -1
323 self.logger.warning(
324 'Overwriting residual tolerance with -1 because RK methods are direct and hence may not compute a residual at all!'
325 )
327 self.__level = lvl
329 def predict(self):
330 """
331 Predictor to fill values at nodes before first sweep
332 """
334 # get current level and problem
335 lvl = self.level
336 prob = lvl.prob
338 for m in range(1, self.coll.num_nodes + 1):
339 lvl.u[m] = prob.dtype_u(init=prob.init, val=0.0)
341 # indicate that this level is now ready for sweeps
342 lvl.status.unlocked = True
343 lvl.status.updated = True
346class RungeKuttaIMEX(RungeKutta):
347 """
348 Implicit-explicit split Runge Kutta base class. Only supports methods that share the nodes and weights.
349 """
351 matrix_explicit = None
352 ButcherTableauClass_explicit = ButcherTableau
354 def __init__(self, params):
355 """
356 Initialization routine
358 Args:
359 params: parameters for the sweeper
360 """
361 super().__init__(params)
362 self.coll_explicit = self.get_Butcher_tableau_explicit()
363 self.QE = self.coll_explicit.Qmat
365 def predict(self):
366 """
367 Predictor to fill values at nodes before first sweep
368 """
370 # get current level and problem
371 lvl = self.level
372 prob = lvl.prob
374 for m in range(1, self.coll.num_nodes + 1):
375 lvl.u[m] = prob.dtype_u(init=prob.init, val=0.0)
376 lvl.f[m] = prob.dtype_f(init=prob.init, val=0.0)
378 # indicate that this level is now ready for sweeps
379 lvl.status.unlocked = True
380 lvl.status.updated = True
382 @classmethod
383 def get_Butcher_tableau_explicit(cls):
384 return cls.ButcherTableauClass_explicit(cls.weights, cls.nodes, cls.matrix_explicit)
386 def integrate(self):
387 """
388 Integrates the right-hand side
390 Returns:
391 list of dtype_u: containing the integral as values
392 """
394 # get current level and problem
395 lvl = self.level
396 prob = lvl.prob
398 me = []
400 # integrate RHS over all collocation nodes
401 for m in range(1, self.coll.num_nodes + 1):
402 # new instance of dtype_u, initialize values with 0
403 me.append(prob.dtype_u(prob.init, val=0.0))
404 for j in range(1, self.coll.num_nodes + 1):
405 me[-1] += lvl.dt * (
406 self.coll.Qmat[m, j] * lvl.f[j].impl + self.coll_explicit.Qmat[m, j] * lvl.f[j].expl
407 )
409 return me
411 def update_nodes(self):
412 """
413 Update the u- and f-values at the collocation nodes
415 Returns:
416 None
417 """
419 # get current level and problem
420 lvl = self.level
421 prob = lvl.prob
423 # only if the level has been touched before
424 assert lvl.status.unlocked
425 assert lvl.status.sweep <= 1, "RK schemes are direct solvers. Please perform only 1 iteration!"
427 # get number of collocation nodes for easier access
428 M = self.coll.num_nodes
430 for m in range(0, M):
431 # build rhs, consisting of the known values from above and new values from previous nodes (at k+1)
432 rhs = lvl.u[0]
433 for j in range(1, m + 1):
434 rhs += lvl.dt * (self.QI[m + 1, j] * lvl.f[j].impl + self.QE[m + 1, j] * lvl.f[j].expl)
436 # implicit solve with prefactor stemming from the diagonal of Qd, use previous stage as initial guess
437 lvl.u[m + 1][:] = prob.solve_system(
438 rhs, lvl.dt * self.QI[m + 1, m + 1], lvl.u[m], lvl.time + lvl.dt * self.coll.nodes[m + 1]
439 )
441 # update function values (we don't usually need to evaluate the RHS at the solution of the step)
442 if m < M - self.coll.num_solution_stages or self.params.eval_rhs_at_right_boundary:
443 lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
445 # indicate presence of new values at this level
446 lvl.status.updated = True
448 return None
451class ForwardEuler(RungeKutta):
452 """
453 Forward Euler. Still a classic.
455 Not very stable first order method.
456 """
458 generator = RK_SCHEMES["FE"]()
459 nodes, weights, matrix = generator.genCoeffs()
462class BackwardEuler(RungeKutta):
463 """
464 Backward Euler. A favorite among true connoisseurs of the heat equation.
466 A-stable first order method.
467 """
469 generator = RK_SCHEMES["BE"]()
470 nodes, weights, matrix = generator.genCoeffs()
473class CrankNicholson(RungeKutta):
474 """
475 Implicit Runge-Kutta method of second order, A-stable.
476 """
478 generator = RK_SCHEMES["CN"]()
479 nodes, weights, matrix = generator.genCoeffs()
482class ExplicitMidpointMethod(RungeKutta):
483 """
484 Explicit Runge-Kutta method of second order.
485 """
487 generator = RK_SCHEMES["RK2"]()
488 nodes, weights, matrix = generator.genCoeffs()
491class ImplicitMidpointMethod(RungeKutta):
492 """
493 Implicit Runge-Kutta method of second order.
494 """
496 generator = RK_SCHEMES["IMP"]()
497 nodes, weights, matrix = generator.genCoeffs()
500class RK4(RungeKutta):
501 """
502 Explicit Runge-Kutta of fourth order: Everybody's darling.
503 """
505 generator = RK_SCHEMES["RK4"]()
506 nodes, weights, matrix = generator.genCoeffs()
509class Heun_Euler(RungeKutta):
510 """
511 Second order explicit embedded Runge-Kutta method.
512 """
514 generator = RK_SCHEMES["HEUN"]()
515 nodes, weights, matrix = generator.genCoeffs()
517 @classmethod
518 def get_update_order(cls):
519 return 2
522class Cash_Karp(RungeKutta):
523 """
524 Fifth order explicit embedded Runge-Kutta. See [here](https://doi.org/10.1145/79505.79507).
525 """
527 generator = RK_SCHEMES["CashKarp"]()
528 nodes, weights, matrix = generator.genCoeffs(embedded=True)
529 ButcherTableauClass = ButcherTableauEmbedded
531 @classmethod
532 def get_update_order(cls):
533 return 5
536class DIRK43(RungeKutta):
537 """
538 Embedded A-stable diagonally implicit RK pair of order 3 and 4.
540 Taken from [here](https://doi.org/10.1007/BF01934920).
541 """
543 generator = RK_SCHEMES["EDIRK43"]()
544 nodes, weights, matrix = generator.genCoeffs(embedded=True)
545 ButcherTableauClass = ButcherTableauEmbedded
547 @classmethod
548 def get_update_order(cls):
549 return 4
552class DIRK43_2(RungeKutta):
553 """
554 L-stable Diagonally Implicit RK method with four stages of order 3.
555 Taken from [here](https://en.wikipedia.org/wiki/List_of_Runge%E2%80%93Kutta_methods).
556 """
558 generator = RK_SCHEMES["DIRK43"]()
559 nodes, weights, matrix = generator.genCoeffs()
562class EDIRK4(RungeKutta):
563 """
564 Stiffly accurate, fourth-order EDIRK with four stages. Taken from
565 [here](https://ntrs.nasa.gov/citations/20160005923), second one in eq. (216).
566 """
568 generator = RK_SCHEMES["EDIRK4"]()
569 nodes, weights, matrix = generator.genCoeffs()
572class ESDIRK53(RungeKutta):
573 """
574 A-stable embedded RK pair of orders 5 and 3, ESDIRK5(3)6L[2]SA.
575 Taken from [here](https://ntrs.nasa.gov/citations/20160005923)
576 """
578 generator = RK_SCHEMES["ESDIRK53"]()
579 nodes, weights, matrix = generator.genCoeffs(embedded=True)
580 ButcherTableauClass = ButcherTableauEmbedded
582 @classmethod
583 def get_update_order(cls):
584 return 4
587class ESDIRK43(RungeKutta):
588 """
589 A-stable embedded RK pair of orders 4 and 3, ESDIRK4(3)6L[2]SA.
590 Taken from [here](https://ntrs.nasa.gov/citations/20160005923)
591 """
593 generator = RK_SCHEMES["ESDIRK43"]()
594 nodes, weights, matrix = generator.genCoeffs(embedded=True)
595 ButcherTableauClass = ButcherTableauEmbedded
597 @classmethod
598 def get_update_order(cls):
599 return 4
602class ARK548L2SAERK(RungeKutta):
603 """
604 Explicit part of the ARK54 scheme.
605 """
607 generator = RK_SCHEMES["ARK548L2SAERK"]()
608 nodes, weights, matrix = generator.genCoeffs(embedded=True)
609 ButcherTableauClass = ButcherTableauEmbedded
611 @classmethod
612 def get_update_order(cls):
613 return 5
616class ARK548L2SAESDIRK(ARK548L2SAERK):
617 """
618 Implicit part of the ARK54 scheme. Be careful with the embedded scheme. It seems that both schemes are order 5 as opposed to 5 and 4 as claimed. This may cause issues when doing adaptive time-stepping.
619 """
621 generator_IMP = RK_SCHEMES["ARK548L2SAESDIRK"]()
622 matrix = generator_IMP.Q
625class ARK54(RungeKuttaIMEX):
626 """
627 Pair of pairs of ARK5(4)8L[2]SA-ERK and ARK5(4)8L[2]SA-ESDIRK from [here](https://doi.org/10.1016/S0168-9274(02)00138-1).
628 """
630 ButcherTableauClass = ButcherTableauEmbedded
631 ButcherTableauClass_explicit = ButcherTableauEmbedded
633 nodes = ARK548L2SAERK.nodes
634 weights = ARK548L2SAERK.weights
636 matrix = ARK548L2SAESDIRK.matrix
637 matrix_explicit = ARK548L2SAERK.matrix
639 @classmethod
640 def get_update_order(cls):
641 return 5
644class ARK548L2SAESDIRK2(RungeKutta):
645 """
646 Stiffly accurate singly diagonally L-stable implicit embedded Runge-Kutta pair of orders 5 and 4 with explicit first stage from [here](https://doi.org/10.1016/j.apnum.2018.10.007).
647 This method is part of the IMEX method ARK548L2SA.
648 """
650 generator = RK_SCHEMES["ARK548L2SAESDIRK2"]()
651 nodes, weights, matrix = generator.genCoeffs(embedded=True)
652 ButcherTableauClass = ButcherTableauEmbedded
654 @classmethod
655 def get_update_order(cls):
656 return 5
659class ARK548L2SAERK2(ARK548L2SAESDIRK2):
660 """
661 Explicit embedded pair of Runge-Kutta methods of orders 5 and 4 from [here](https://doi.org/10.1016/j.apnum.2018.10.007).
662 This method is part of the IMEX method ARK548L2SA.
663 """
665 generator_EXP = RK_SCHEMES["ARK548L2SAERK2"]()
666 matrix = generator_EXP.Q
669class ARK548L2SA(RungeKuttaIMEX):
670 """
671 IMEX Runge-Kutta method of order 5 based on the explicit method ARK548L2SAERK2 and the implicit method
672 ARK548L2SAESDIRK2 from [here](https://doi.org/10.1016/j.apnum.2018.10.007).
674 According to Kennedy and Carpenter (see reference), the two IMEX RK methods of order 5 are the only ones available
675 as of now. And we are not aware of higher order ones. This one is newer then the other one and apparently better.
676 """
678 ButcherTableauClass = ButcherTableauEmbedded
679 ButcherTableauClass_explicit = ButcherTableauEmbedded
681 nodes = ARK548L2SAERK2.nodes
682 weights = ARK548L2SAERK2.weights
684 matrix = ARK548L2SAESDIRK2.matrix
685 matrix_explicit = ARK548L2SAERK2.matrix
687 @classmethod
688 def get_update_order(cls):
689 return 5