|
| 1 | +import numpy as np |
| 2 | +from pySDC.core.Problem import ptype |
| 3 | +from pySDC.implementations.datatype_classes.mesh import mesh |
| 4 | + |
| 5 | + |
| 6 | +class LorenzAttractor(ptype): |
| 7 | + """ |
| 8 | + Simple script to run a Lorenz attractor problem. |
| 9 | +
|
| 10 | + The Lorenz attractor is a system of three ordinary differential equations that exhibits some chaotic behaviour. |
| 11 | + It is well known for the "Butterfly Effect", because the solution looks like a butterfly (solve to Tend = 100 or |
| 12 | + so to see this with these initial conditions) and because of the chaotic nature. |
| 13 | +
|
| 14 | + Since the problem is non-linear, we need to use a Newton solver. |
| 15 | +
|
| 16 | + Problem and initial conditions do not originate from, but were taken from doi.org/10.2140/camcos.2015.10.1 |
| 17 | + """ |
| 18 | + |
| 19 | + def __init__(self, problem_params): |
| 20 | + """ |
| 21 | + Initialization routine |
| 22 | +
|
| 23 | + Args: |
| 24 | + problem_params (dict): custom parameters for the example |
| 25 | + """ |
| 26 | + |
| 27 | + # take care of essential parameters in defaults such that they always exist |
| 28 | + defaults = { |
| 29 | + 'sigma': 10.0, |
| 30 | + 'rho': 28.0, |
| 31 | + 'beta': 8.0 / 3.0, |
| 32 | + 'nvars': 3, |
| 33 | + 'newton_tol': 1e-9, |
| 34 | + 'newton_maxiter': 99, |
| 35 | + } |
| 36 | + |
| 37 | + for key in defaults.keys(): |
| 38 | + problem_params[key] = problem_params.get(key, defaults[key]) |
| 39 | + |
| 40 | + # invoke super init, passing number of dofs, dtype_u and dtype_f |
| 41 | + super().__init__( |
| 42 | + init=(problem_params['nvars'], None, np.dtype('float64')), |
| 43 | + dtype_u=mesh, |
| 44 | + dtype_f=mesh, |
| 45 | + params=problem_params, |
| 46 | + ) |
| 47 | + |
| 48 | + def eval_f(self, u, t): |
| 49 | + """ |
| 50 | + Routine to evaluate the RHS |
| 51 | +
|
| 52 | + Args: |
| 53 | + u (dtype_u): current values |
| 54 | + t (float): current time |
| 55 | +
|
| 56 | + Returns: |
| 57 | + dtype_f: the RHS |
| 58 | + """ |
| 59 | + # abbreviations |
| 60 | + sigma = self.params.sigma |
| 61 | + rho = self.params.rho |
| 62 | + beta = self.params.beta |
| 63 | + |
| 64 | + f = self.dtype_f(self.init) |
| 65 | + |
| 66 | + f[0] = sigma * (u[1] - u[0]) |
| 67 | + f[1] = rho * u[0] - u[1] - u[0] * u[2] |
| 68 | + f[2] = u[0] * u[1] - beta * u[2] |
| 69 | + return f |
| 70 | + |
| 71 | + def solve_system(self, rhs, dt, u0, t): |
| 72 | + """ |
| 73 | + Simple Newton solver for the nonlinear system. |
| 74 | + Notice that I did not go through the trouble of inverting the Jacobian beforehand. If you have some time on your |
| 75 | + hands feel free to do that! In the current implementation it is inverted using `numpy.linalg.solve`, which is a |
| 76 | + bit more expensive than simple matrix-vector multiplication. |
| 77 | +
|
| 78 | + Args: |
| 79 | + rhs (dtype_f): right-hand side for the linear system |
| 80 | + dt (float): abbrev. for the local stepsize (or any other factor required) |
| 81 | + u0 (dtype_u): initial guess for the iterative solver |
| 82 | + t (float): current time (e.g. for time-dependent BCs) |
| 83 | +
|
| 84 | + Returns: |
| 85 | + dtype_u: solution as mesh |
| 86 | + """ |
| 87 | + |
| 88 | + # abbreviations |
| 89 | + sigma = self.params.sigma |
| 90 | + rho = self.params.rho |
| 91 | + beta = self.params.beta |
| 92 | + |
| 93 | + # start Newton iterations |
| 94 | + u = self.dtype_u(u0) |
| 95 | + res = np.inf |
| 96 | + for _n in range(0, self.params.newton_maxiter): |
| 97 | + |
| 98 | + # assemble G such that G(u) = 0 at the solution to the step |
| 99 | + G = np.array( |
| 100 | + [ |
| 101 | + u[0] - dt * sigma * (u[1] - u[0]) - rhs[0], |
| 102 | + u[1] - dt * (rho * u[0] - u[1] - u[0] * u[2]) - rhs[1], |
| 103 | + u[2] - dt * (u[0] * u[1] - beta * u[2]) - rhs[2], |
| 104 | + ] |
| 105 | + ) |
| 106 | + |
| 107 | + # compute the residual and determine if we are done |
| 108 | + res = np.linalg.norm(G, np.inf) |
| 109 | + if res <= self.params.newton_tol or np.isnan(res): |
| 110 | + break |
| 111 | + |
| 112 | + # assemble Jacobian J of G |
| 113 | + J = np.array( |
| 114 | + [ |
| 115 | + [1.0 + dt * sigma, -dt * sigma, 0], |
| 116 | + [-dt * (rho - u[2]), 1 + dt, dt * u[0]], |
| 117 | + [-dt * u[1], -dt * u[0], 1.0 + dt * beta], |
| 118 | + ] |
| 119 | + ) |
| 120 | + |
| 121 | + # solve the linear system for the Newton correction J delta = G |
| 122 | + delta = np.linalg.solve(J, G) |
| 123 | + |
| 124 | + # update solution |
| 125 | + u = u - delta |
| 126 | + |
| 127 | + return u |
| 128 | + |
| 129 | + def u_exact(self, t, u_init=None, t_init=None): |
| 130 | + """ |
| 131 | + Routine to return initial conditions or to approximate exact solution using scipy. |
| 132 | +
|
| 133 | + Args: |
| 134 | + t (float): current time |
| 135 | + u_init (pySDC.implementations.problem_classes.Lorenz.dtype_u): initial conditions for getting the exact solution |
| 136 | + t_init (float): the starting time |
| 137 | +
|
| 138 | + Returns: |
| 139 | + dtype_u: exact solution |
| 140 | + """ |
| 141 | + me = self.dtype_u(self.init) |
| 142 | + |
| 143 | + if t > 0: |
| 144 | + from scipy.integrate import solve_ivp |
| 145 | + |
| 146 | + def rhs(t, u): |
| 147 | + """ |
| 148 | + Evaluate the right hand side, but switch the arguments for scipy. |
| 149 | +
|
| 150 | + Args: |
| 151 | + t (float): Time |
| 152 | + u (numpy.ndarray): Solution at time t |
| 153 | +
|
| 154 | + Returns: |
| 155 | + (numpy.ndarray): Right hand side evaluation |
| 156 | + """ |
| 157 | + return self.eval_f(u, t) |
| 158 | + |
| 159 | + tol = 100 * np.finfo(float).eps |
| 160 | + u_init = self.u_exact(t=0) if u_init is None else u_init |
| 161 | + t_init = 0 if t_init is None else t_init |
| 162 | + |
| 163 | + me[:] = solve_ivp(rhs, (t_init, t), u_init, rtol=tol, atol=tol).y[:, -1] |
| 164 | + else: |
| 165 | + me[:] = 1.0 |
| 166 | + return me |
0 commit comments