Coverage for pySDC/projects/Resilience/dahlquist.py: 0%
147 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-29 09:02 +0000
1# script to run a simple advection problem
2from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d
3from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
4from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
5from pySDC.core.Hooks import hooks
6from pySDC.helpers.stats_helper import get_sorted
7import numpy as np
8import matplotlib.pyplot as plt
9from matplotlib.colors import Normalize
10from mpl_toolkits.axes_grid1 import make_axes_locatable
12from pySDC.implementations.hooks.log_solution import LogSolutionAfterIteration
13from pySDC.implementations.hooks.log_step_size import LogStepSize
14from pySDC.projects.Resilience.strategies import merge_descriptions
17class LogLambdas(hooks):
18 """
19 Store the lambda values at the beginning of the run
20 """
22 def pre_run(self, step, level_number):
23 super().pre_run(step, level_number)
24 L = step.levels[level_number]
25 self.add_to_stats(process=0, time=0, level=0, iter=0, sweep=0, type='lambdas', value=L.prob.lambdas)
28hooks = [LogLambdas, LogSolutionAfterIteration, LogStepSize]
31def run_dahlquist(
32 custom_description=None,
33 num_procs=1,
34 Tend=1.0,
35 hook_class=hooks,
36 fault_stuff=None,
37 custom_controller_params=None,
38 **kwargs,
39):
40 """
41 Run a Dahlquist problem with default parameters.
43 Args:
44 custom_description (dict): Overwrite presets
45 num_procs (int): Number of steps for MSSDC
46 Tend (float): Time to integrate to
47 hook_class (pySDC.Hook): A hook to store data
48 fault_stuff (dict): A dictionary with information on how to add faults
49 custom_controller_params (dict): Overwrite presets
51 Returns:
52 dict: The stats object
53 controller: The controller
54 Tend: The time that was supposed to be integrated to
55 """
57 # initialize level parameters
58 level_params = dict()
59 level_params['dt'] = 1.0
61 # initialize sweeper parameters
62 sweeper_params = dict()
63 sweeper_params['quad_type'] = 'RADAU-RIGHT'
64 sweeper_params['num_nodes'] = 3
65 sweeper_params['QI'] = 'IE'
66 sweeper_params['initial_guess'] = 'random'
68 # build lambdas
69 re = np.linspace(-30, 30, 400)
70 im = np.linspace(-50, 50, 400)
71 lambdas = np.array([[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]).reshape(
72 (len(re) * len(im))
73 )
75 problem_params = {
76 'lambdas': lambdas,
77 'u0': 1.0 + 0.0j,
78 }
80 # initialize step parameters
81 step_params = dict()
82 step_params['maxiter'] = 5
84 # initialize controller parameters
85 controller_params = dict()
86 controller_params['logger_level'] = 30
87 controller_params['hook_class'] = hook_class
88 controller_params['mssdc_jac'] = False
90 if custom_controller_params is not None:
91 controller_params = {**controller_params, **custom_controller_params}
93 # fill description dictionary for easy step instantiation
94 description = dict()
95 description['problem_class'] = testequation0d # pass problem class
96 description['problem_params'] = problem_params # pass problem parameters
97 description['sweeper_class'] = generic_implicit # pass sweeper
98 description['sweeper_params'] = sweeper_params # pass sweeper parameters
99 description['level_params'] = level_params # pass level parameters
100 description['step_params'] = step_params
102 if custom_description is not None:
103 description = merge_descriptions(description, custom_description)
105 # set time parameters
106 t0 = 0.0
108 # instantiate controller
109 controller = controller_nonMPI(num_procs=num_procs, controller_params=controller_params, description=description)
111 # insert faults
112 if fault_stuff is not None:
113 raise NotImplementedError('No fault stuff here...')
115 # get initial values on finest level
116 P = controller.MS[0].levels[0].prob
117 uinit = P.u_exact(t0)
119 # call main function to get things done...
120 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
121 return stats, controller, Tend
124def plot_stability(stats, ax=None, iter=None, colors=None, crosshair=True, fill=False, **kwargs):
125 """
126 Plot the domain of stability by checking if the solution grows.
128 Args:
129 stats (pySDC.stats): The stats object of the run
130 ax: Somewhere to plot
131 iter (list): Check the stability for different numbers of iterations
132 colors (list): Colors for the different iterations
133 crosshair (bool): Whether to highlight the axes
134 fill (bool): Fill the contours or not
136 Returns:
137 bool: If the method is A-stable or not
138 """
139 lambdas = get_sorted(stats, type='lambdas')[0][1]
140 u = get_sorted(stats, type='u', sortby='iter')
142 if ax is None:
143 fig, ax = plt.subplots(1, 1)
145 # decorate
146 if crosshair:
147 ax.axhline(0, color='black', alpha=1.0)
148 ax.axvline(0, color='black', alpha=1.0)
150 iter = [1] if iter is None else iter
151 colors = ['blue', 'red', 'violet', 'green'] if colors is None else colors
153 for i in iter:
154 # isolate the solutions from the iteration you want
155 U = np.reshape([me[1] for me in u if me[0] == i], (len(np.unique(lambdas.real)), len(np.unique(lambdas.imag))))
157 # get a grid for plotting
158 X, Y = np.meshgrid(np.unique(lambdas.real), np.unique(lambdas.imag))
159 if fill:
160 ax.contourf(X, Y, abs(U), levels=[-np.inf, 1 - np.finfo(float).eps], colors=colors[i - 1], alpha=0.5)
161 ax.contour(X, Y, abs(U), levels=[1], colors=colors[i - 1])
162 ax.plot([None], [None], color=colors[i - 1], label=f'k={i}')
164 # check if the method is A-stable
165 unstable = abs(U) > 1.0
166 Astable = not any(X[unstable] < 0)
168 ax.legend(frameon=False)
170 return Astable
173def plot_contraction(stats, fig=None, ax=None, iter=None, plot_increase=False, cbar=True, **kwargs):
174 """
175 Plot the contraction of the error.
177 Args:
178 stats (pySDC.stats): The stats object of the run
179 fig: Figure of the plot, needed for a colorbar
180 ax: Somewhere to plot
181 iter (list): Plot the contraction for different numbers of iterations
182 plot_increase (bool): Whether to also include increasing errors.
183 cbar (bool): Plot a color bar or not
185 Returns:
186 The plot
187 """
188 lambdas = get_sorted(stats, type='lambdas')[0][1]
189 real = np.unique(lambdas.real)
190 imag = np.unique(lambdas.imag)
192 u = get_sorted(stats, type='u', sortby='iter')
193 t = get_sorted(stats, type='u', sortby='time')[0][0]
194 u_exact = np.exp(lambdas * t)
196 kwargs['cmap'] = kwargs.get('cmap', 'seismic' if plot_increase else 'jet')
198 # decide which iterations to look at
199 iter = [0, 1] if iter is None else iter
200 assert len(iter) > 1, 'Need to compute the contraction factor across multiple iterations!'
202 # get solution for the specified iterations
203 us = [me[1] for me in u if me[0] in iter]
204 if 0 in iter: # ic's are not stored in stats, so we have to add them manually
205 us = np.append([np.ones_like(lambdas)], us, axis=0)
207 # get error for each iteration
208 e = abs(us - u_exact)
209 e[e == 0] = np.finfo(float).eps
211 # get contraction rates for each iteration
212 rho = e[1:, :] / e[:-1, :]
213 rho_avg = np.mean(rho, axis=0)
214 rho_log = np.log(np.reshape(rho_avg, (len(imag), len(real))))
216 # get spaceally averaged contraction factor
217 # rho_avg_space = np.mean(rho, axis=1)
218 # e_tot = np.sum(e, axis=1)
219 # rho_tot = e_tot[1:] / e_tot[:-1]
221 if ax is None:
222 fig, ax = plt.subplots(1, 1)
224 # get a grid for plotting
225 X, Y = np.meshgrid(real, imag)
226 if plot_increase:
227 ax.contour(X, Y, rho_log, levels=[0.0])
228 lim = max(np.abs([rho_log.min(), rho_log.max()]))
229 kwargs['vmin'] = kwargs.get('vmin', -lim)
230 kwargs['vmax'] = kwargs.get('vmax', lim)
231 cs = ax.contourf(X, Y, rho_log, **kwargs)
232 else:
233 cs = ax.contourf(X, Y, np.where(rho_log <= 0, rho_log, None), levels=500, **kwargs)
235 # decorate
236 ax.axhline(0, color='black')
237 ax.axvline(0, color='black')
239 # fix pdf plotting
240 ax.set_rasterized(True)
242 if cbar:
243 divider = make_axes_locatable(ax)
244 cbar_ax = divider.append_axes('right', 0.2, pad=0.1)
245 cb = fig.colorbar(cs, cbar_ax)
246 cb.set_label(r'$\rho$')
247 cbar_ax.set_rasterized(True)
248 return cs
251def plot_increment(stats, fig=None, ax=None, iter=None, cbar=True, **kwargs):
252 """
253 Plot the increment between iterations.
255 Args:
256 stats (pySDC.stats): The stats object of the run
257 fig: Figure of the plot, needed for a colorbar
258 ax: Somewhere to plot
259 iter (list): Plot the contraction for different numbers of iterations
260 cbar (bool): Plot a color bar or not
262 Returns:
263 None
264 """
265 lambdas = get_sorted(stats, type='lambdas')[0][1]
266 u = get_sorted(stats, type='u', sortby='iter')
268 kwargs['cmap'] = kwargs.get('cmap', 'jet')
270 # decide which iterations to look at
271 iter = [0, 1] if iter is None else iter
272 assert len(iter) > 1, 'Need to compute the increment across multiple iterations!'
274 # get solution for the specified iterations
275 u_iter = [me[1] for me in u if me[0] in iter]
276 if 0 in iter: # ics are not stored in stats, so we have to add them manually
277 u_iter = np.append(np.ones_like(lambdas), u_iter)
278 us = np.reshape(u_iter, (len(iter), len(lambdas)))
280 # get contraction rates for each iteration
281 rho = abs(us[1:, :] / us[:-1, :])
282 rho_avg = np.mean(rho, axis=0)
283 rho_log = np.log(np.reshape(rho_avg, (len(np.unique(lambdas.real)), len(np.unique(lambdas.imag)))))
285 if ax is None:
286 fig, ax = plt.subplots(1, 1)
288 # get a grid for plotting
289 X, Y = np.meshgrid(np.unique(lambdas.real), np.unique(lambdas.imag))
290 cs = ax.contourf(X, Y, rho_log, levels=500, **kwargs)
292 # outline the region where the increment is 0
293 ax.contour(X, Y, rho_log, levels=[-15], colors=['red'])
295 # decorate
296 ax.axhline(0, color='black')
297 ax.axvline(0, color='black')
299 # fix pdf plotting
300 ax.set_rasterized(True)
302 if cbar:
303 divider = make_axes_locatable(ax)
304 cbar_ax = divider.append_axes('right', 0.2, pad=0.1)
305 cb = fig.colorbar(cs, cbar_ax)
306 cb.set_label('increment')
307 cbar_ax.set_rasterized(True)
310def compare_contraction():
311 """
312 Make a plot comparing contraction factors between trapezoidal rule and implicit Euler.
313 """
314 fig, axs = plt.subplots(1, 2, figsize=(12, 5.5), gridspec_kw={'width_ratios': [0.8, 1]})
315 precons = ['TRAP', 'IE']
316 norm = Normalize(vmin=-7, vmax=0)
317 cbar = True
318 for i in range(len(precons)):
319 custom_description = {'sweeper_params': {'QI': precons[i]}}
320 stats, controller, Tend = run_dahlquist(custom_description=custom_description, res=(400, 400))
321 plot_contraction(stats, fig=fig, ax=axs[i], cbar=cbar, norm=norm, cmap='jet')
322 cbar = False
323 axs[i].set_title(precons[i])
324 fig.tight_layout()
327if __name__ == '__main__':
328 custom_description = None
329 stats, controller, Tend = run_dahlquist(custom_description=custom_description)
330 plot_stability(stats, iter=[1, 2, 3])
331 plot_contraction(stats, iter=[0, 4])
332 plt.show()