Coverage for pySDC/projects/parallelSDC_reloaded/utils.py: 100%
164 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-29 09:02 +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, **kwargs):
214 dt = tEnd / nSteps
215 setupProblem(probName, params, dt, **kwargs)
217 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
219 prob = controller.MS[0].levels[0].prob
221 uInit = prob.u_exact(0)
222 uTmp = uInit.copy()
224 uSDC = np.zeros((nSteps + 1, uInit.size), dtype=uInit.dtype)
225 uSDC[0] = uInit
226 tVals = np.linspace(0, tEnd, nSteps + 1)
227 tBeg = time()
228 print(" -- computing numerical solution with pySDC ...")
229 warnings.filterwarnings("ignore")
230 for i in range(nSteps):
231 uTmp[:] = uSDC[i]
232 try:
233 uSDC[i + 1], _ = controller.run(u0=uTmp, t0=tVals[i], Tend=tVals[i + 1])
234 except Exception as e:
235 print(f" -- exception when running controller : {e}")
236 warnings.resetwarnings()
237 return None, (0, 0, 0), False
238 warnings.resetwarnings()
239 tComp = time() - tBeg
241 try:
242 nNewton = prob.work_counters["newton"].niter
243 except KeyError:
244 nNewton = 0
245 nRHS = prob.work_counters["rhs"].niter
246 print(f" done, newton:{nNewton}, rhs:{nRHS}, tComp:{tComp}")
247 try:
248 parallel = controller.MS[0].levels[0].sweep.parallelizable
249 except AttributeError: # pragma: no cover
250 parallel = False
252 return uSDC, (nNewton, nRHS, tComp), parallel
255def solutionExact(tEnd, nSteps, probName, **kwargs):
256 """Return the exact solution of the Van-der-Pol problem at tEnd"""
258 if probName == "VANDERPOL":
259 mu = kwargs.get('mu', 10)
260 key = f"{tEnd}_{nSteps}_{mu}"
261 cacheFile = '_solVanderpolExact.json'
262 elif probName == "LORENZ":
263 u0 = kwargs.get('u0', (1, 1, 1))
264 key = f"{tEnd}_{nSteps}_{u0}"
265 cacheFile = '_solLorenzExact.json'
266 elif probName == "CHEMREC":
267 key = f"{tEnd}_{nSteps}"
268 cacheFile = '_solChemicalReactionExact.json'
269 elif probName == "JACELL":
270 key = f"{tEnd}_{nSteps}"
271 cacheFile = '_solJacobiEllipticExact.json'
273 # Eventually load already computed solution from local cache
274 try:
275 with open(f"{PATH}/{cacheFile}", "r") as f:
276 cache = json.load(f)
277 if key in cache:
278 return np.array(cache[key])
279 except (FileNotFoundError, json.JSONDecodeError, UnboundLocalError):
280 cache = {}
282 # Compute solution
283 params = getParamsSDC()
284 dt = tEnd / nSteps
285 setupProblem(probName, params, dt, **kwargs)
287 controller = controller_nonMPI(num_procs=1, controller_params={'logger_level': 30}, description=params)
288 solver = controller.MS[0].levels[0].prob.u_exact
290 print(" -- computing analytical solution with P.u_exact ...")
291 tBeg = time()
292 tVals = np.linspace(0, tEnd, nSteps + 1)
293 uExact = [solver(0)]
294 for i in range(nSteps):
295 try:
296 uExact.append(solver(tVals[i + 1], uExact[-1], tVals[i]))
297 except TypeError:
298 uExact.append(solver(tVals[i + 1]))
299 uExact = np.array(uExact)
300 print(f" done in {time()-tBeg:1.2f}s")
302 try:
303 # Save solution in local cache
304 cache[key] = uExact.tolist()
305 with open(f"{PATH}/{cacheFile}", "w") as f:
306 json.dump(cache, f)
307 except UnboundLocalError:
308 pass
310 return uExact
313# Plotting functions
314def plotStabContour(reVals, imVals, stab, ax=None):
315 if ax is None:
316 ax = plt.gca()
317 ax.contour(reVals, imVals, stab, levels=[1.0], colors='black', linewidths=1)
318 ax.contourf(reVals, imVals, stab, levels=[1.0, np.inf], colors='gainsboro')
319 ax.hlines(0, min(reVals), max(reVals), linestyles='--', colors='black', linewidth=0.5)
320 ax.vlines(0, min(imVals), max(imVals), linestyles='--', colors='black', linewidth=0.5)
321 ax.set_aspect('equal', 'box')
322 ax.set_xlabel(r"$Re(z)$", labelpad=0, fontsize=10)
323 ax.set_ylabel(r"$Im(z)$", labelpad=0, fontsize=10)
324 return ax