Coverage for pySDC/projects/parallelSDC_reloaded/utils.py: 100%
168 statements
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-16 14:51 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Created on Sun Nov 12 18:50:39 2023
6Utility functions to investigate parallel SDC on non-linear problems
7"""
8import os
9import json
10import numpy as np
11from time import time
12import warnings
14from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
16from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol
17from pySDC.implementations.problem_classes.Lorenz import LorenzAttractor
18from pySDC.implementations.problem_classes.odeScalar import ProtheroRobinson
19from pySDC.implementations.problem_classes.odeSystem import (
20 ProtheroRobinsonAutonomous,
21 Kaps,
22 ChemicalReaction3Var,
23 JacobiElliptic,
24)
25from pySDC.implementations.problem_classes.AllenCahn_1D_FD import (
26 allencahn_front_fullyimplicit,
27 allencahn_periodic_fullyimplicit,
28)
29from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d
31from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
32import pySDC.implementations.sweeper_classes.Runge_Kutta as rk
34import matplotlib.pyplot as plt
36PATH = '/' + os.path.join(*__file__.split('/')[:-1])
38# General matplotlib settings
39plt.rc('font', size=12)
40plt.rcParams['lines.linewidth'] = 2
41plt.rcParams['axes.titlesize'] = 16
42plt.rcParams['axes.labelsize'] = 16
43plt.rcParams['xtick.labelsize'] = 15
44plt.rcParams['ytick.labelsize'] = 15
45plt.rcParams['xtick.major.pad'] = 3
46plt.rcParams['ytick.major.pad'] = 2
47plt.rcParams['axes.labelpad'] = 6
48plt.rcParams['markers.fillstyle'] = 'none'
49plt.rcParams['lines.markersize'] = 7.0
50plt.rcParams['lines.markeredgewidth'] = 1.5
51plt.rcParams['mathtext.fontset'] = 'cm'
52plt.rcParams['mathtext.rm'] = 'serif'
53plt.rcParams['figure.max_open_warning'] = 100
56def getParamsSDC(
57 quadType="RADAU-RIGHT", numNodes=4, qDeltaI="IE", nSweeps=3, nodeType="LEGENDRE", collUpdate=False, initType="copy"
58):
59 description = {
60 # Sweeper and its parameters
61 "sweeper_class": generic_implicit,
62 "sweeper_params": {
63 "quad_type": quadType,
64 "num_nodes": numNodes,
65 "node_type": nodeType,
66 "initial_guess": initType,
67 "do_coll_update": collUpdate,
68 "QI": qDeltaI,
69 'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
70 },
71 # Step parameters
72 "step_params": {
73 "maxiter": 1,
74 },
75 # Level parameters
76 "level_params": {
77 "restol": -1,
78 "nsweeps": nSweeps,
79 },
80 }
82 return description
85RK_SWEEPERS = {
86 "BE": rk.BackwardEuler,
87 "FE": rk.ForwardEuler,
88 "RK4": rk.RK4,
89 "DIRK43": rk.DIRK43,
90 "ESDIRK53": rk.ESDIRK53,
91 "ESDIRK43": rk.ESDIRK43,
92}
95def getParamsRK(method="RK4"):
96 description = {
97 # Sweeper and its parameters
98 "sweeper_class": RK_SWEEPERS[method],
99 "sweeper_params": {'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE')},
100 # Step parameters
101 "step_params": {
102 "maxiter": 1,
103 },
104 # Level parameters
105 "level_params": {
106 "restol": -1,
107 "nsweeps": 1,
108 },
109 }
111 return description
114def setupProblem(name, description, dt, **kwargs):
115 """Add problem settings to pySDC description parameters"""
117 # Common newton tolerance and max number of iterations
118 description["problem_params"] = {
119 'newton_tol': 1e-8,
120 'newton_maxiter': 300,
121 }
122 # Level parameters
123 description["level_params"]["dt"] = dt
125 if name == "VANDERPOL":
126 description["problem_class"] = vanderpol
127 description["problem_params"].update(
128 {
129 'newton_tol': 1e-12,
130 'mu': kwargs.get("mu", 10), # vanderpol parameter
131 'u0': np.array([2.0, 0]),
132 }
133 )
134 elif name == "LORENZ":
135 description["problem_class"] = LorenzAttractor
136 description["problem_params"].update(
137 {
138 'newton_tol': 1e-12,
139 'u0': kwargs.get("u0", (1, 1, 1)),
140 }
141 )
142 elif name == "PROTHERO-ROBINSON":
143 description["problem_class"] = ProtheroRobinson
144 description["problem_params"].update(
145 {
146 'newton_tol': 1e-12,
147 'epsilon': kwargs.get("epsilon", 1e-3),
148 }
149 )
150 elif name == "PROTHERO-ROBINSON-NL":
151 description["problem_class"] = ProtheroRobinson
152 description["problem_params"].update(
153 {
154 'newton_tol': 1e-12,
155 'epsilon': kwargs.get("epsilon", 1e-3),
156 'nonLinear': True,
157 }
158 )
159 elif name == "PROTHERO-ROBINSON-A":
160 description["problem_class"] = ProtheroRobinsonAutonomous
161 description["problem_params"].update(
162 {
163 'newton_tol': 1e-12,
164 'epsilon': kwargs.get("epsilon", 1e-3),
165 }
166 )
167 elif name == "PROTHERO-ROBINSON-A-NL":
168 description["problem_class"] = ProtheroRobinsonAutonomous
169 description["problem_params"].update(
170 {
171 'newton_tol': 1e-12,
172 'epsilon': kwargs.get("epsilon", 1e-3),
173 'nonLinear': True,
174 }
175 )
176 elif name == "KAPS":
177 description["problem_class"] = Kaps
178 description["problem_params"].update(
179 {
180 'epsilon': kwargs.get("epsilon", 1e-3),
181 }
182 )
183 elif name == "CHEMREC":
184 description["problem_class"] = ChemicalReaction3Var
185 elif name == "ALLEN-CAHN":
186 periodic = kwargs.get("periodic", False)
187 description["problem_class"] = allencahn_periodic_fullyimplicit if periodic else allencahn_front_fullyimplicit
188 description["problem_params"].update(
189 {
190 'newton_tol': 1e-8,
191 'nvars': kwargs.get("nvars", 128 if periodic else 127),
192 'eps': kwargs.get("epsilon", 0.04),
193 'stop_at_maxiter': True,
194 }
195 )
196 elif name == "JACELL":
197 description["problem_class"] = JacobiElliptic
198 elif name == "DAHLQUIST":
199 lambdas = kwargs.get("lambdas", None)
200 description["problem_class"] = testequation0d
201 description["problem_params"].update(
202 {
203 "lambdas": lambdas,
204 "u0": 1.0,
205 }
206 )
207 description["problem_params"].pop("newton_tol")
208 description["problem_params"].pop("newton_maxiter")
209 else:
210 raise NotImplementedError(f"problem {name} not implemented")
213def solutionSDC(tEnd, nSteps, params, probName, verbose=True, noExcept=False, **kwargs):
214 dt = tEnd / nSteps
215 setupProblem(probName, params, dt, **kwargs)
216 if noExcept:
217 params["problem_params"]["stop_at_nan"] = False
218 params["problem_params"]["stop_at_maxiter"] = False
220 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
222 prob = controller.MS[0].levels[0].prob
224 uInit = prob.u_exact(0)
225 uTmp = uInit.copy()
227 uSDC = np.zeros((nSteps + 1, uInit.size), dtype=uInit.dtype)
228 uSDC[0] = uInit
229 tVals = np.linspace(0, tEnd, nSteps + 1)
230 tBeg = time()
231 if verbose:
232 print(" -- computing numerical solution with pySDC ...")
233 warnings.filterwarnings("ignore")
234 for i in range(nSteps):
235 uTmp[:] = uSDC[i]
236 try:
237 uSDC[i + 1], _ = controller.run(u0=uTmp, t0=tVals[i], Tend=tVals[i + 1])
238 except Exception as e:
239 print(f" -- exception when running controller : {e}")
240 warnings.resetwarnings()
241 return None, (0, 0, 0), False
242 warnings.resetwarnings()
243 tComp = time() - tBeg
245 try:
246 nNewton = prob.work_counters["newton"].niter
247 except KeyError:
248 nNewton = 0
249 nRHS = prob.work_counters["rhs"].niter
250 if verbose:
251 print(f" done, newton:{nNewton}, rhs:{nRHS}, tComp:{tComp}")
252 try:
253 parallel = controller.MS[0].levels[0].sweep.parallelizable
254 except AttributeError: # pragma: no cover
255 parallel = False
257 return uSDC, (nNewton, nRHS, tComp), parallel
260def solutionExact(tEnd, nSteps, probName, **kwargs):
261 """Return the exact solution of the Van-der-Pol problem at tEnd"""
263 if probName == "VANDERPOL":
264 mu = kwargs.get('mu', 10)
265 key = f"{tEnd}_{nSteps}_{mu}"
266 cacheFile = '_solVanderpolExact.json'
267 elif probName == "LORENZ":
268 u0 = kwargs.get('u0', (1, 1, 1))
269 key = f"{tEnd}_{nSteps}_{u0}"
270 cacheFile = '_solLorenzExact.json'
271 elif probName == "CHEMREC":
272 key = f"{tEnd}_{nSteps}"
273 cacheFile = '_solChemicalReactionExact.json'
274 elif probName == "JACELL":
275 key = f"{tEnd}_{nSteps}"
276 cacheFile = '_solJacobiEllipticExact.json'
278 # Potentially load already computed solution from local cache
279 try:
280 with open(f"{PATH}/{cacheFile}", "r") as f:
281 cache = json.load(f)
282 if key in cache:
283 return np.array(cache[key])
284 except (FileNotFoundError, json.JSONDecodeError, UnboundLocalError):
285 cache = {}
287 # Compute solution
288 params = getParamsSDC()
289 dt = tEnd / nSteps
290 setupProblem(probName, params, dt, **kwargs)
292 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
293 solver = controller.MS[0].levels[0].prob.u_exact
295 print(" -- computing analytical solution with P.u_exact ...")
296 tBeg = time()
297 tVals = np.linspace(0, tEnd, nSteps + 1)
298 uExact = [solver(0)]
299 for i in range(nSteps):
300 try:
301 uExact.append(solver(tVals[i + 1], uExact[-1], tVals[i]))
302 except TypeError:
303 uExact.append(solver(tVals[i + 1]))
304 uExact = np.array(uExact)
305 print(f" done in {time()-tBeg:1.2f}s")
307 try:
308 # Save solution in local cache
309 cache[key] = uExact.tolist()
310 with open(f"{PATH}/{cacheFile}", "w") as f:
311 json.dump(cache, f)
312 except UnboundLocalError:
313 pass
315 return uExact
318# Plotting functions
319def plotStabContour(reVals, imVals, stab, ax=None):
320 if ax is None:
321 ax = plt.gca()
322 ax.contour(reVals, imVals, stab, levels=[1.0], colors='black', linewidths=1)
323 ax.contourf(reVals, imVals, stab, levels=[1.0, np.inf], colors='gainsboro')
324 ax.hlines(0, min(reVals), max(reVals), linestyles='--', colors='black', linewidth=0.5)
325 ax.vlines(0, min(imVals), max(imVals), linestyles='--', colors='black', linewidth=0.5)
326 ax.set_aspect('equal', 'box')
327 ax.set_xlabel(r"$Re(z)$", labelpad=0, fontsize=10)
328 ax.set_ylabel(r"$Im(z)$", labelpad=0, fontsize=10)
329 return ax